Add argument intrinsics (#725)

Author: christiangnrd

src/device/intrinsics/arguments.jl

@@ -17,11 +17,16 @@ const nodim_intr = [
 ]
 
 for (intr, offset) in nodim_intr
-    # XXX: these are also available as UInt16 (ushort)
+    intr_i16 = Symbol(intr, :_i16)
+    intr = Symbol(intr)
     @eval begin
-        export $(Symbol(intr))
+        # UInt32
+        export $(intr)
+        @device_function $(intr)() = ccall($"extern julia.air.$intr.i32", llvmcall, UInt32, ()) + UInt32($offset)
 
-        @device_function $(Symbol(intr))() = ccall($"extern julia.air.$intr.i32", llvmcall, UInt32, ()) + UInt32($offset)
+        # UInt16
+        export $(intr_i16)
+        @device_function $(intr_i16)() = ccall($"extern julia.air.$intr.i16", llvmcall, UInt16, ()) + UInt16($offset)
     end
 end
 
@@ -39,36 +44,47 @@ const dim_intr = [
 ]
 
 for (intr, offset) in dim_intr
-    # XXX: these are also available as UInt16 (ushort)
+    intr_i16 = Symbol(intr, :_i16)
+    intr = Symbol(intr)
     @eval begin
-        export $(Symbol(intr))
-
-        @device_function function $(Symbol(intr))()
+        # UInt32
+        export $(intr)
+        @device_function function $(intr)()
             vec = ccall($"extern julia.air.$intr.v3i32", llvmcall,
                         NTuple{3, VecElement{UInt32}}, ())
             (x = vec[1].value + UInt32($offset),
              y = vec[2].value + UInt32($offset),
              z = vec[3].value + UInt32($offset))
         end
+
+        # UInt16
+        export $(intr_i16)
+        @device_function function $(intr_i16)()
+            vec = ccall($"extern julia.air.$intr.v3i16", llvmcall,
+                        NTuple{3, VecElement{UInt16}}, ())
+            (x = vec[1].value + UInt16($offset),
+             y = vec[2].value + UInt16($offset),
+             z = vec[3].value + UInt16($offset))
+        end
     end
 
     # deprecated aliases
     @eval begin
-        export $(Symbol(intr * "_1d"))
-        export $(Symbol(intr * "_2d"))
-        export $(Symbol(intr * "_3d"))
+        export $(Symbol(intr, "_1d"))
+        export $(Symbol(intr, "_2d"))
+        export $(Symbol(intr, "_3d"))
 
-        function $(Symbol(intr * "_1d"))()
-            $(Symbol(intr))().x
+        function $(Symbol(intr, "_1d"))()
+            $(intr)().x
         end
 
-        function $(Symbol(intr * "_2d"))()
-            vec = $(Symbol(intr))()
+        function $(Symbol(intr, "_2d"))()
+            vec = $(intr)()
             (x = vec.x,
              y = vec.y)
         end
 
-        $(Symbol(intr * "_3d"))() = $(Symbol(intr))()
+        $(Symbol(intr, "_3d"))() = $(intr)()
     end
 end
 

test/device/intrinsics/arguments.jl

@@ -56,6 +56,29 @@
     @on_device threads_per_threadgroup_2d()
     @on_device threads_per_threadgroup_3d()
 
+    # UInt16
+    @on_device dispatch_quadgroups_per_threadgroup_i16()
+    @on_device dispatch_simdgroups_per_threadgroup_i16()
+    @on_device quadgroup_index_in_threadgroup_i16()
+    @on_device quadgroups_per_threadgroup_i16()
+    @on_device simdgroup_index_in_threadgroup_i16()
+    @on_device simdgroups_per_threadgroup_i16()
+    @on_device thread_index_in_quadgroup_i16()
+    @on_device thread_index_in_simdgroup_i16()
+    @on_device thread_index_in_threadgroup_i16()
+    @on_device thread_execution_width_i16()
+    @on_device threads_per_simdgroup_i16()
+
+    @on_device dispatch_threads_per_threadgroup_i16()
+    @on_device grid_origin_i16()
+    @on_device grid_size_i16()
+    @on_device thread_position_in_grid_i16()
+    @on_device thread_position_in_threadgroup_i16()
+    @on_device threadgroup_position_in_grid_i16()
+    @on_device threadgroups_per_grid_i16()
+    @on_device threads_per_grid_i16()
+    @on_device threads_per_threadgroup_i16()
+
     global const CPU_ONLY_ERR = "This function is not intended for use on the CPU"
 
     @test_throws CPU_ONLY_ERR dispatch_quadgroups_per_threadgroup()
@@ -114,4 +137,27 @@
     @test_throws CPU_ONLY_ERR threads_per_threadgroup_1d()
     @test_throws CPU_ONLY_ERR threads_per_threadgroup_2d()
     @test_throws CPU_ONLY_ERR threads_per_threadgroup_3d()
+
+    # UInt16
+    @test_throws CPU_ONLY_ERR dispatch_quadgroups_per_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR dispatch_simdgroups_per_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR quadgroup_index_in_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR quadgroups_per_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR simdgroup_index_in_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR simdgroups_per_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR thread_index_in_quadgroup()
+    @test_throws CPU_ONLY_ERR thread_index_in_simdgroup_i16()
+    @test_throws CPU_ONLY_ERR thread_index_in_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR thread_execution_width_i16()
+    @test_throws CPU_ONLY_ERR threads_per_simdgroup_i16()
+
+    @test_throws CPU_ONLY_ERR dispatch_threads_per_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR grid_origin_i16()
+    @test_throws CPU_ONLY_ERR grid_size_i16()
+    @test_throws CPU_ONLY_ERR thread_position_in_grid_i16()
+    @test_throws CPU_ONLY_ERR thread_position_in_threadgroup_i16()
+    @test_throws CPU_ONLY_ERR threadgroup_position_in_grid_i16()
+    @test_throws CPU_ONLY_ERR threadgroups_per_grid_i16()
+    @test_throws CPU_ONLY_ERR threads_per_grid_i16()
+    @test_throws CPU_ONLY_ERR threads_per_threadgroup_i16()
 end