Allow early returns (#92)

Author: AntonOresten

src/compiler/codegen/control_flow.jl

@@ -453,6 +453,59 @@ function emit_terminator!(ctx::CGCtx, ::ConditionOp)
     # ConditionOp is handled specially by emit_while_op!, not emitted as a terminator
 end
 
+#=============================================================================
+ Early Return Hoisting
+
+ tileiras rejects ReturnNode (cuda_tile.return) inside IfOp (cuda_tile.if)
+ regions. This pre-pass rewrites the structured IR so that ReturnNode only
+ appears at the top level, replacing nested returns with YieldOp.
+=============================================================================#
+
+"""
+    hoist_returns!(block::Block)
+
+Rewrite `ReturnNode` terminators inside `IfOp` regions into `YieldOp`,
+hoisting the return to the parent block. Operates recursively so that
+nested early returns (multiple successive `if ... return end` patterns)
+are handled automatically.
+
+Only handles the case where BOTH branches of an IfOp terminate with
+ReturnNode (REGION_TERMINATION with 3 children). The 2-child case
+(early return inside a loop) is not handled.
+"""
+function hoist_returns!(block::Block)
+    # First, recurse into all nested control flow ops
+    for (_, entry) in block.body
+        stmt = entry.stmt
+        if stmt isa IfOp
+            hoist_returns!(stmt.then_region)
+            hoist_returns!(stmt.else_region)
+        elseif stmt isa ForOp
+            hoist_returns!(stmt.body)
+        elseif stmt isa WhileOp
+            hoist_returns!(stmt.before)
+            hoist_returns!(stmt.after)
+        elseif stmt isa LoopOp
+            hoist_returns!(stmt.body)
+        end
+    end
+
+    # Now check: does this block contain an IfOp where both branches return?
+    # If so, replace branch ReturnNodes with YieldOp and set block terminator.
+    for (_, entry) in block.body
+        entry.stmt isa IfOp || continue
+        op = entry.stmt::IfOp
+        op.then_region.terminator isa ReturnNode || continue
+        op.else_region.terminator isa ReturnNode || continue
+
+        # Both branches return — hoist to parent block.
+        # Replace branch terminators with YieldOp (void — no values to yield).
+        op.then_region.terminator = YieldOp()
+        op.else_region.terminator = YieldOp()
+        block.terminator = ReturnNode(nothing)
+    end
+end
+
 """
     emit_getfield!(ctx, args) -> Union{CGVal, Nothing}
 

src/compiler/codegen/kernel.jl

@@ -157,6 +157,9 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
     ctx.token_type = token_type
     ctx.token = encode_MakeTokenOp!(cb, token_type)
 
+    # Hoist early returns out of IfOp regions (tileiras rejects ReturnOp inside IfOp)
+    hoist_returns!(ctx.sci.entry)
+
     # Emit the structured IR (uses original Julia SSA indices everywhere)
     emit_block!(ctx, ctx.sci.entry)
 

test/execution/basic.jl

@@ -1154,3 +1154,66 @@ end
     end
 end
 
+@testset "early return — taken" begin
+    function early_return_skip(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, flag::Int32)
+        pid = ct.bid(1)
+        tile = ct.load(a, pid, (16,))
+        if flag == Int32(0)
+            return nothing
+        end
+        ct.store(b, pid, tile .* 2.0f0)
+        return nothing
+    end
+
+    a = CUDA.rand(Float32, 64)
+    b = CUDA.zeros(Float32, 64)
+    ct.launch(early_return_skip, 4, a, b, Int32(0))
+    @test all(Array(b) .== 0.0f0)
+end
+
+@testset "early return — not taken" begin
+    function early_return_store(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, flag::Int32)
+        pid = ct.bid(1)
+        tile = ct.load(a, pid, (16,))
+        if flag == Int32(0)
+            return nothing
+        end
+        ct.store(b, pid, tile .* 2.0f0)
+        return nothing
+    end
+
+    a = CUDA.rand(Float32, 64)
+    b = CUDA.zeros(Float32, 64)
+    ct.launch(early_return_store, 4, a, b, Int32(1))
+    @test Array(b) ≈ Array(a) .* 2.0f0
+end
+
+@testset "multiple early returns" begin
+    function multi_early_return(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1},
+                                flag1::Int32, flag2::Int32)
+        pid = ct.bid(1)
+        tile = ct.load(a, pid, (16,))
+        if flag1 == Int32(0)
+            return nothing
+        end
+        if flag2 == Int32(0)
+            return nothing
+        end
+        ct.store(b, pid, tile .* 2.0f0)
+        return nothing
+    end
+
+    a = CUDA.rand(Float32, 64)
+
+    b1 = CUDA.zeros(Float32, 64)
+    ct.launch(multi_early_return, 4, a, b1, Int32(1), Int32(1))
+    @test Array(b1) ≈ Array(a) .* 2.0f0
+
+    b2 = CUDA.zeros(Float32, 64)
+    ct.launch(multi_early_return, 4, a, b2, Int32(0), Int32(1))
+    @test all(Array(b2) .== 0.0f0)
+
+    b3 = CUDA.zeros(Float32, 64)
+    ct.launch(multi_early_return, 4, a, b3, Int32(1), Int32(0))
+    @test all(Array(b3) .== 0.0f0)
+end