Use stacked method tables

Author: vchuravy

Project.toml

@@ -34,3 +34,6 @@ SPIRV_LLVM_Backend_jll = "20"
 SPIRV_Tools_jll = "2025.1"
 StaticArrays = "1"
 julia = "1.10"
+
+[sources]
+GPUCompiler = {url="https://github.com/JuliaGPU/GPUCompiler.jl", rev="vc/mtv"}

lib/intrinsics/src/SPIRVIntrinsics.jl

@@ -32,6 +32,7 @@ macro import_all()
         # bring all the names of this module in scope
         name in (:SPIRVIntrinsics, :eval, :include) && continue
         startswith(string(name), "#") && continue
+        name in (:method_table, :@device_function, :@device_override) && continue
         push!(code.args, :(using .SPIRVIntrinsics: $name))
     end
 

src/OpenCL.jl

@@ -15,10 +15,40 @@ include("../lib/cl/CL.jl")
 @reexport using .cl
 export cl
 
+## device overrides
+
+# local method table for device functions
+Base.Experimental.@MethodTable(method_table)
+
+macro device_override(ex)
+    esc(quote
+        Base.Experimental.@overlay($method_table, $ex)
+    end)
+end
+
+macro device_function(ex)
+    ex = macroexpand(__module__, ex)
+    def = ExprTools.splitdef(ex)
+
+    # generate a function that errors
+    def[:body] = quote
+        error("This function is not intended for use on the CPU")
+    end
+
+    esc(quote
+        $(ExprTools.combinedef(def))
+        @device_override $ex
+    end)
+end
+
+
 # device functionality
 import SPIRVIntrinsics
 SPIRVIntrinsics.@import_all
 SPIRVIntrinsics.@reexport_public
+
+const spirv_method_table = SPIRVIntrinsics.method_table
+
 include("device/runtime.jl")
 include("device/array.jl")
 include("device/quirks.jl")

src/compiler/compilation.jl

@@ -6,7 +6,7 @@ const OpenCLCompilerJob = CompilerJob{SPIRVCompilerTarget,OpenCLCompilerParams}
 
 GPUCompiler.runtime_module(::CompilerJob{<:Any,OpenCLCompilerParams}) = OpenCL
 
-GPUCompiler.method_table(::OpenCLCompilerJob) = method_table
+GPUCompiler.method_table_view(job::OpenCLCompilerJob) = GPUCompiler.StackedMethodTable(job.world, method_table, spirv_method_table)
 
 # filter out OpenCL built-ins
 # TODO: eagerly lower these using the translator API