SPIR-V: switch on an i32 instead of an i1

src/spirv.jl

@@ -318,6 +318,11 @@ end
     # (after `prepare_execution!`, whose InstCombine would narrow the widened operations)
     widen_bfloat16!(job, mod)
 
+    # the LLVM back-end translates a switch on an `i1` to an OpSwitch on a boolean, which SPIR-V
+    # doesn't allow. InstCombine narrows a switch's condition to the bits it can take, so this
+    # too runs after `prepare_execution!`.
+    widen_bool_switches!(mod)
+
     # The SPIRV Tools don't handle Julia's debug info, rejecting DW_LANG_Julia...
     strip_debuginfo!(mod)
 
@@ -908,6 +913,38 @@ end
 
 ## LLVM passes
 
+# make switches on an `i1` switch on it zero-extended to an `i32`, with the same successors, so
+# that the `phi`s in those keep their incoming blocks
+function widen_bool_switches!(mod::LLVM.Module)
+    changed = false
+    @tracepoint "widen bool switches" begin
+
+    T_int32 = LLVM.Int32Type()
+    for f in mod.functions, bb in f.blocks
+        sw = bb.terminator
+        sw isa LLVM.SwitchInst || continue
+        ops = sw.operands
+        cond, default = ops[1], ops[2]
+        cond.value_type == LLVM.Int1Type() || continue
+        @dispose builder=IRBuilder() begin
+            position!(builder, LLVM.before(sw))
+            builder.debug_location = sw.debug_location
+            cond32 = zext!(builder, cond, T_int32)
+            ncases = (length(ops) - 2) ÷ 2
+            new = switch!(builder, cond32, default, ncases)
+            for i in 1:ncases
+                val = convert(Bool, ops[2 + 2i - 1])
+                LLVM.API.LLVMAddCase(new, ConstantInt(T_int32, val), ops[2 + 2i])
+            end
+        end
+        erase!(sw)
+        changed = true
+    end
+
+    end
+    return changed
+end
+
 # flatten `insertvalue`s with multiple indices into single-index ones, extracting and
 # re-inserting the intermediate aggregates: `insertvalue %agg, %val, 1, 0` becomes
 #   %sub = extractvalue %agg, 1

test/spirv.jl

@@ -718,6 +718,34 @@ end
 
 end
 
+@testset "bool switch widening" begin
+    # `widen_bool_switches!` makes a switch on an `i1`, which InstCombine can leave and the LLVM
+    # back-end translates to an invalid OpSwitch on a boolean, switch on an `i32`
+    Context() do ctx
+        mod = parse(LLVM.Module, """
+            define i32 @entry(i1 %c, i32 %x) {
+            entry:
+              switch i1 %c, label %other [
+                i1 false, label %no
+                i1 true, label %join
+              ]
+            no:
+              br label %join
+            other:
+              ret i32 0
+            join:
+              %r = phi i32 [ 1, %no ], [ %x, %entry ]
+              ret i32 %r
+            }""")
+        @test GPUCompiler.widen_bool_switches!(mod)
+        verify(mod)
+        ir = string(mod)
+        @test !occursin("switch i1", ir)
+        @test occursin(r"zext i1 %c to i32\n\s+switch i32 %\d+, label %other \[\n\s+i32 0, label %no\n\s+i32 1, label %join", ir)
+        @test !GPUCompiler.widen_bool_switches!(mod)
+    end
+end
+
 @testset "min and max lowering" begin
     # `lower_minimum_maximum!` lowers every family of floating-point minimum and maximum to
     # `llvm.minnum`/`llvm.maxnum`, which the back-ends translate to OpenCL's `fmin`/`fmax`.