Port work item intrinsics.

Author: maleadt

lib/intrinsics/src/work_item.jl

@@ -1,28 +1,60 @@
 # Work-Item Functions
 
-export get_work_dim,
-       get_global_size, get_global_id,
-       get_local_size, get_enqueued_local_size, get_local_id,
-       get_num_groups, get_group_id,
-       get_global_offset,
-       get_global_linear_id, get_local_linear_id
-
 # NOTE: these functions now unsafely truncate to Int to avoid top bit checks.
 #       we should probably use range metadata instead.
 
-@device_function get_work_dim() = @builtin_ccall("get_work_dim", UInt32, ()) % Int
-
-@device_function get_global_size(dimindx::Integer=1) = @builtin_ccall("get_global_size", UInt, (UInt32,), dimindx-1) % Int
-@device_function get_global_id(dimindx::Integer=1) = @builtin_ccall("get_global_id", UInt, (UInt32,), dimindx-1) % Int + 1
-
-@device_function get_local_size(dimindx::Integer=1) = @builtin_ccall("get_local_size", UInt, (UInt32,), dimindx-1) % Int
-@device_function get_enqueued_local_size(dimindx::Integer=1) = @builtin_ccall("get_enqueued_local_size", UInt, (UInt32,), dimindx-1) % Int
-@device_function get_local_id(dimindx::Integer=1) = @builtin_ccall("get_local_id", UInt, (UInt32,), dimindx-1) % Int + 1
-
-@device_function get_num_groups(dimindx::Integer=1) = @builtin_ccall("get_num_groups", UInt, (UInt32,), dimindx-1) % Int
-@device_function get_group_id(dimindx::Integer=1) = @builtin_ccall("get_group_id", UInt, (UInt32,), dimindx-1) % Int + 1
-
-@device_function get_global_offset(dimindx::Integer=1) = @builtin_ccall("get_global_offset", UInt, (UInt32,), dimindx-1) % Int + 1
-
-@device_function get_global_linear_id() = @builtin_ccall("get_global_linear_id", UInt, ()) % Int + 1
-@device_function get_local_linear_id() = @builtin_ccall("get_local_linear_id", UInt, ()) % Int + 1
+# 1D values
+for (julia_name, (spirv_name, offset)) in [
+        # indices
+        :get_global_linear_id           => (:BuiltInGlobalLinearId, 1),
+        :get_local_linear_id            => (:BuiltInLocalInvocationIndex, 1),
+        :get_sub_group_id               => (:BuiltInSubgroupId, 1),
+        :get_sub_group_local_id         => (:BuiltInSubgroupLocalInvocationId, 1),
+        # sizes
+        :get_work_dim                   => (:BuiltInWorkDim, 0),
+        :get_sub_group_size             => (:BuiltInSubgroupSize, 0),
+        :get_max_sub_group_size         => (:BuiltInSubgroupMaxSize, 0),
+        :get_num_sub_groups             => (:BuiltInNumSubgroups, 0),
+        :get_enqueued_num_sub_groups    => (:BuiltInNumEnqueuedSubgroups, 0)]
+    gvar_name = Symbol("@__spirv_$(spirv_name)")
+    @eval begin
+        export $julia_name
+        @device_function $julia_name() =
+            Base.llvmcall(
+                $("""$gvar_name = external addrspace(1) global i32
+                     define i32 @entry() #0 {
+                         %val = load i32, i32 addrspace(1)* $gvar_name
+                         ret i32 %val
+                     }
+                     attributes #0 = { alwaysinline }
+                """, "entry"), UInt32, Tuple{}) % Int + $offset
+    end
+end
+
+# 3D values
+for (julia_name, (spirv_name, offset)) in [
+        # indices
+        :get_global_id              => (:BuiltInGlobalInvocationId, 1),
+        :get_global_offset          => (:BuiltInGlobalOffset, 1),
+        :get_local_id               => (:BuiltInLocalInvocationId, 1),
+        :get_group_id               => (:BuiltInWorkgroupId, 1),
+        # sizes
+        :get_global_size            => (:BuiltInGlobalSize, 0),
+        :get_local_size             => (:BuiltInWorkgroupSize, 0),
+        :get_enqueued_local_size    => (:BuiltInEnqueuedWorkgroupSize, 0),
+        :get_num_groups             => (:BuiltInNumWorkgroups, 0)]
+    gvar_name = Symbol("@__spirv_$(spirv_name)")
+    @eval begin
+        export $julia_name
+        @device_function $julia_name(dimindx::Integer=1) =
+            Base.llvmcall(
+                $("""$gvar_name = external addrspace(1) global <3 x i32>
+                     define i32 @entry(i32 %idx) #0 {
+                         %val = load <3 x i32>, <3 x i32> addrspace(1)* $gvar_name
+                         %element = extractelement <3 x i32> %val, i32 %idx
+                         ret i32 %element
+                     }
+                     attributes #0 = { alwaysinline }
+                """, "entry"), UInt32, Tuple{UInt32}, UInt32(dimindx - 1)) % Int + $offset
+    end
+end