SPIR-V: declare native integer widths in the datalayout

src/interface.jl

@@ -321,6 +321,8 @@ uses_julia_runtime(@nospecialize(job::CompilerJob)) = false
 #   for `InstCombinePass`, retaining only the simplification subset of the peephole
 #   transforms (useful e.g. for downstream rewriters like Enzyme that get confused by
 #   InstCombine's more aggressive rewrites).
+# - `switch_to_lookup::Bool = true`: when `false`, SimplifyCFG does not turn switches into
+#   lookup tables or bitmaps, which can introduce integer types of arbitrary width.
 #
 # Returning a `NamedTuple` keeps this single extension point lightweight: downstream
 # users add new keys without GPUCompiler having to grow an interface method per option.

src/optim.jl

@@ -111,6 +111,13 @@ const AggressiveSimplifyCFGOptions =
        #sink_common_insts=true
     )
 
+# Build a SimplifyCFG pass with `options`, letting `optimization_options(job).switch_to_lookup`
+# disable the conversion of switches to lookup tables.
+function simplifycfg_pass(@nospecialize(job::CompilerJob), options)
+    switch_to_lookup = get(optimization_options(job), :switch_to_lookup, true)
+    SimplifyCFGPass(; options..., switch_to_lookup)
+end
+
 function buildEarlySimplificationPipeline(mpm, @nospecialize(job::CompilerJob), opt_level)
     if should_verify()
         add!(mpm, FunctionPassManager()) do fpm
@@ -133,7 +140,7 @@ function buildEarlySimplificationPipeline(mpm, @nospecialize(job::CompilerJob),
         # DCE must come before simplifycfg: codegen can generate unused
         # statements that would otherwise alter how simplifycfg optimizes the CFG.
         add!(fpm, DCEPass())
-        add!(fpm, SimplifyCFGPass(; BasicSimplifyCFGOptions...))
+        add!(fpm, simplifycfg_pass(job, BasicSimplifyCFGOptions))
         if opt_level >= 1
             add!(fpm, SROAPass())
             add!(fpm, EarlyCSEPass())
@@ -216,7 +223,7 @@ function buildLoopOptimizerPipeline(fpm, @nospecialize(job::CompilerJob), opt_le
     if opt_level >= 2
         add!(fpm, IRCEPass())
     end
-    add!(fpm, SimplifyCFGPass(; BasicSimplifyCFGOptions...))
+    add!(fpm, simplifycfg_pass(job, BasicSimplifyCFGOptions))
     add!(fpm, instcombine_pass(job))
     add!(fpm, LoopPassManager()) do lpm
         if opt_level >= 2
@@ -268,16 +275,16 @@ function buildScalarOptimizerPipeline(fpm, @nospecialize(job::CompilerJob), opt_
         if LLVM.version() >= v"17"
             add!(fpm, PeepholeCallbacks(; opt_level))
         end
-        add!(fpm, SimplifyCFGPass(; AggressiveSimplifyCFGOptions...))
+        add!(fpm, simplifycfg_pass(job, AggressiveSimplifyCFGOptions))
         add!(fpm, AllocOptPass())
         add!(fpm, LoopPassManager(; use_memory_ssa=true)) do lpm
             add!(lpm, LICMPass())
             add!(lpm, JuliaLICMPass())
         end
-        add!(fpm, SimplifyCFGPass(; AggressiveSimplifyCFGOptions...))
+        add!(fpm, simplifycfg_pass(job, AggressiveSimplifyCFGOptions))
         add!(fpm, instcombine_pass(job))
     elseif opt_level >= 1
-        add!(fpm, SimplifyCFGPass(; AggressiveSimplifyCFGOptions...))
+        add!(fpm, simplifycfg_pass(job, AggressiveSimplifyCFGOptions))
     end
     if LLVM.version() >= v"17"
         add!(fpm, ScalarOptimizerLateCallbacks(; opt_level))
@@ -294,7 +301,7 @@ function buildVectorPipeline(fpm, @nospecialize(job::CompilerJob), opt_level)
     add!(fpm, InjectTLIMappings())
     add!(fpm, LoopVectorizePass())
     add!(fpm, LoopLoadEliminationPass())
-    add!(fpm, SimplifyCFGPass(; AggressiveSimplifyCFGOptions...))
+    add!(fpm, simplifycfg_pass(job, AggressiveSimplifyCFGOptions))
     add!(fpm, LoopPassManager(; use_memory_ssa=true)) do lpm
         add!(lpm, LICMPass())
     end
@@ -394,7 +401,7 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op
     if opt_level >= 1
         add!(mpm, FunctionPassManager()) do fpm
             add!(fpm, instcombine_pass(job))
-            add!(fpm, SimplifyCFGPass(; AggressiveSimplifyCFGOptions...))
+            add!(fpm, simplifycfg_pass(job, AggressiveSimplifyCFGOptions))
             # raise the alignment of memory accesses to what can be proven, e.g., from the
             # alignment of a shared memory global. back-ends only look at the alignment of
             # each access, e.g., to vectorize them. Julia's own pipeline does the same since

src/spirv.jl

@@ -140,8 +140,8 @@ function spirv_extensions(target::SPIRVCompilerTarget)
 end
 
 llvm_datalayout(::SPIRVCompilerTarget) = Int===Int64 ?
-    "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-G1" :
-    "e-p:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-G1"
+    "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1" :
+    "e-p:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
 
 
 ## job
@@ -151,6 +151,11 @@ llvm_datalayout(::SPIRVCompilerTarget) = Int===Int64 ?
 isintrinsic(@nospecialize(job::CompilerJob{SPIRVCompilerTarget}), fn::String) =
     occursin(r"^(_Z\d+)?__spirv_", fn)
 
+# lookup bitmaps are integers of the table's bit size (e.g. i24), which the Khronos
+# translator does not support
+optimization_options(@nospecialize(job::CompilerJob{SPIRVCompilerTarget})) =
+    (; switch_to_lookup=false)
+
 function finish_module!(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module,
                         entry::LLVM.Function)
     # update calling convention

test/spirv.jl

@@ -623,6 +623,68 @@ end
     end
 end
 
+@testset "native integer widths" begin
+    # without native integer widths in the datalayout, InstCombine narrows this switch to
+    # an i2, which the Khronos translator does not support
+    mod = @eval module $(gensym())
+        function kernel(out, x)
+            v = unsafe_load(x) & 3
+            if v == 1
+                unsafe_store!(out, 10)
+            elseif v == 2
+                unsafe_store!(out, 20)
+            elseif v == 3
+                unsafe_store!(out, 30)
+            end
+            return
+        end
+    end
+
+    @test @filecheck begin
+        @check_label "define spir_kernel void @_Z6kernel"
+        @check_not "{{[^0-9A-Za-z_]i2[^0-9]}}"
+        @check "switch i64"
+        SPIRV.code_llvm(mod.kernel, Tuple{Ptr{Int64}, Ptr{Int64}}; backend, kernel=true)
+    end
+end
+
+@testset "no switch lookup tables" begin
+    # SimplifyCFG would turn these into an odd-width bitmap and a lookup table once the
+    # datalayout declares native integer widths
+    mod = @eval module $(gensym())
+        function bitmap(out, x)
+            v = unsafe_load(x)
+            r = v == 0 ? Int8(5) : v == 1 ? Int8(17) : v == 2 ? Int8(9) :
+                v == 3 ? Int8(100) : v == 4 ? Int8(42) : Int8(0)
+            unsafe_store!(out, r)
+            return
+        end
+        function table(out, x)
+            v = unsafe_load(x)
+            r = v == 0 ? Int32(1000) : v == 1 ? Int32(77) : v == 2 ? Int32(123456) :
+                v == 3 ? Int32(9) : v == 4 ? Int32(31337) : Int32(0)
+            unsafe_store!(out, r)
+            return
+        end
+    end
+
+    @test @filecheck begin
+        @check_label "define spir_kernel void @_Z6bitmap"
+        @check_not "switch."
+        @check "ret void"
+        SPIRV.code_llvm(mod.bitmap, Tuple{Ptr{Int8}, Ptr{Int64}}; backend, kernel=true)
+    end
+
+    @test @filecheck begin
+        @check_not "switch.table"
+        @check_label "define spir_kernel void @_Z5table"
+        @check_not "switch."
+        @check "ret void"
+        SPIRV.code_llvm(mod.table, Tuple{Ptr{Int32}, Ptr{Int64}}; backend, kernel=true,
+                        dump_module=true)
+    end
+end
+
 end
 
 @testset "min and max lowering" begin