From b5cdd8322bd51603c217dfb7662306628b144c78 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 10 Jun 2024 11:45:36 -0700
Subject: Support all integer typed indices in StructuredBuffer Load/Store/[].
 (#4311)

* Support all integer typed indices in StructuredBuffer Load/Store/[].

* Fix tests.

---------

Co-authored-by: Jay Kwak <82421531+jkwak-work@users.noreply.github.com>
---
 source/slang/hlsl.meta.slang                     | 14 ++++++++------
 tests/language-feature/overload-resolution.slang |  8 ++++----
 tests/spirv/i64-structured-buffer.slang          | 16 ++++++++++++++++
 tests/spirv/subgroup-size-2.slang                |  2 +-
 4 files changed, 29 insertions(+), 11 deletions(-)
 create mode 100644 tests/spirv/i64-structured-buffer.slang

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index f9831bef2..e215fd93b 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -3655,13 +3655,14 @@ struct StructuredBuffer
     __intrinsic_op($(kIROp_StructuredBufferLoad))
     [__readNone]
     [require(cpp_cuda_glsl_hlsl_spirv, structuredbuffer)]
-    T Load(int location);
+    T Load<TIndex : __BuiltinIntegerType>(TIndex location);
 
     __intrinsic_op($(kIROp_StructuredBufferLoadStatus))
     [require(cpp_cuda_glsl_hlsl_spirv, structuredbuffer)]
-    T Load(int location, out uint status);
+    T Load<TIndex : __BuiltinIntegerType>(TIndex location, out uint status);
 
-    __subscript(uint index) -> T
+    __generic<TIndex : __BuiltinIntegerType>
+    __subscript(TIndex index) -> T
     {
         [__readNone]
         __intrinsic_op($(kIROp_StructuredBufferLoad))
@@ -4940,13 +4941,14 @@ struct $(item.name)
 
     [__NoSideEffect]
     __intrinsic_op($(kIROp_RWStructuredBufferLoad))
-    T Load(int location);
+    T Load<TIndex : __BuiltinIntegerType>(TIndex location);
 
     [__NoSideEffect]
     __intrinsic_op($(kIROp_RWStructuredBufferLoadStatus))
-    T Load(int location, out uint status);
+    T Load<TIndex : __BuiltinIntegerType>(TIndex location, out uint status);
 
-    __subscript(uint index) -> T
+    __generic<TIndex : __BuiltinIntegerType>
+    __subscript(TIndex index) -> T
     {
         [__NoSideEffect]
         __intrinsic_op($(kIROp_RWStructuredBufferGetElementPtr))
diff --git a/tests/language-feature/overload-resolution.slang b/tests/language-feature/overload-resolution.slang
index 9c135137a..c95269477 100644
--- a/tests/language-feature/overload-resolution.slang
+++ b/tests/language-feature/overload-resolution.slang
@@ -25,10 +25,10 @@ T myGenF<T : __BuiltinIntegerType>(inout T a, T b)
         return a - b;
     }
 }
-// CHECK: result{{.*}}[0{{U?}}] = 1
-// CHECK: result{{.*}}[1{{U?}}] = 4
-// CHECK: result{{.*}}[2{{U?}}] = 1
-// CHECK: result{{.*}}[3{{U?}}] = 4
+// CHECK: result{{.*}}[int(0)] = 1
+// CHECK: result{{.*}}[int(1)] = 4
+// CHECK: result{{.*}}[int(2)] = 1
+// CHECK: result{{.*}}[int(3)] = 4
 [numthreads(1,1,1)]
 void main()
 {
diff --git a/tests/spirv/i64-structured-buffer.slang b/tests/spirv/i64-structured-buffer.slang
new file mode 100644
index 000000000..10cd220f6
--- /dev/null
+++ b/tests/spirv/i64-structured-buffer.slang
@@ -0,0 +1,16 @@
+//TEST:SIMPLE(filecheck=CHECK): -stage compute -target spirv -emit-spirv-directly -entry main
+
+RWStructuredBuffer<float> output;
+
+// Check that 64bit integer index can be used in structured buffers without conversion to int.
+
+// CHECK: %[[INDEX:[A-Za-z0-9_]+]] = OpLoad %long %{{.*}}
+// CHECK: OpAccessChain %_ptr_StorageBuffer_float %output %int_0 %[[INDEX]]
+
+uniform int64_t index;
+
+[numthreads(1,1,1)]
+void main()
+{
+    output[index] = 1;
+}
diff --git a/tests/spirv/subgroup-size-2.slang b/tests/spirv/subgroup-size-2.slang
index 68fee6fe6..bd5ae5eec 100644
--- a/tests/spirv/subgroup-size-2.slang
+++ b/tests/spirv/subgroup-size-2.slang
@@ -19,7 +19,7 @@ void compute1()
     const int x = f().x;
     outputBuffer[0] = x;
 
-    // CHECK-DAG: %[[PTR:[A-Za-z0-9_]+]] = OpAccessChain %_ptr_StorageBuffer_int %outputBuffer %int_0 %uint_1
+    // CHECK-DAG: %[[PTR:[A-Za-z0-9_]+]] = OpAccessChain %_ptr_StorageBuffer_int %outputBuffer %int_0 %int_1
     // CHECK: OpStore %[[PTR]] %int_2
     outputBuffer[1] = WorkgroupSize().y;
 }
-- 
cgit v1.2.3