Rewriting the lower-buffer-element-type pass to avoid unnecessary packing/unpacking. (#8526)

Part of the effort to improve the performance of generated SPIRV code.

The existing lower-buffer-element-type pass works by loading the entire
buffer element content from memory, and translate it to logical type
stored in a local variable at the earliest reference of a buffer handle.
This means that is can generate inefficient code that reads more than
necessary.

Consider this example:
```
struct BigStruct { bool values[1024]; }
ConstantBuffer<BigStruct> cb;

void test(BigStruct v)
{
      if (v.values[0]) { printf("ok"); }
}

[numthreads(1,1,1)]
void computeMain()
{
    test(cb);
}
```

In IR, the `computeMain` function before lower-buffer-element-type pass
is something like following:
```
func test:
   %v = param : BigStruct
   %barr = fieldExtract(%v, "values")
   %element = elementExtract(%barr, 0)
    ... // uses %element 

func computeMain:
  %v = load(cb)
  call %test %v
```

The existing lower-buffer-element-type pass will rewrite the bool array
in `BigStruct` into `int` array so it is legal in SPIRV. However, it
does so by inserting the translation on the first `load` of the constant
buffer:

```
struct BigStruct_std430 {
    int values[1024];
}
var cb : ConstantBuffer<BigStruct_std430>;
func computeMain:
   %tmpVar : var<BigStruct>
    call %unpackStorage(%tmpVar, cb)
   %v : BigStruct = load %tmpVar
   call %test %v
```

This means that the entire array will be loaded and translated to int,
before calling `test`, which only uses one element. It turns out that
the downstream compiler isn't always able to optimize out this
inefficient translation/copy.

This PR completely rewrites the way buffer-element-type lowering is
handled to avoid producing this inefficient code. It works in two parts:
first we turn on the `transformParamsToConstRef` pass for SPIRV target
as well, so we will translate the `test` function to take the `v`
parameter as `constref`. The second part is a redesigned
buffer-element-type pass that defers the storage-type to logical-type
translation until a value is actually used by a `load` instruction.

In this example, after `transformParamsToConstRef`, the IR is:

```
func test:
   %v = param : ConstRef<BigStruct>
   %barr = fieldAddr(%v, "values")
   %elementPtr = elementAddr(%barr, 0)
   %element = load(%elementPtr)
    ... // uses %element 

func computeMain:
  call %test %cb
```

The new `buffer-element-type-lowering` pass will take this IR, and
insert translation at latest possible time across the entire call graph,
and translate the IR into:

```
func test:
   %v = param : ConstRef<BigStruct_std430>
   %barr = fieldAddr(%v, "values")
   %elementPtr : ptr<int> = elementAddr(%barr, 0)
   %element_int = load(%elementPtr)
    %element = cast(%element_int) : %bool
    ... // uses %element 

func computeMain:
  call %test %cb
```

In this new IR, there is no longer a load and conversion of the entire
array.

See new comment in `slang-ir-lower-buffer-element-type.cpp` for more
details of how the pass works.

This PR also address many other issues surfaced by turning on
`transformParamsToConstRef` pass on SPIRV backend.

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>

11 files changed, 115 insertions, 40 deletions

diff --git a/tests/compute/byte-address-buffer-array.slang b/tests/compute/byte-address-buffer-array.slang
index 90cdb2261..58862d1ac 100644
--- a/tests/compute/byte-address-buffer-array.slang
+++ b/tests/compute/byte-address-buffer-array.slang
@@ -1,7 +1,7 @@
 // byte-address-buffer-array.slang
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -d3d12 -profile cs_6_0 -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -cuda -profile cs_6_0 -shaderobj -output-using-type
-//DISABLED_TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -vk -shaderobj -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -vk -shaderobj -output-using-type
 
 //TEST:SIMPLE(filecheck=CHECK2):-target hlsl -entry computeMain -stage compute
 //TEST:SIMPLE(filecheck=CHECK3):-target spirv -entry computeMain -stage compute
@@ -10,7 +10,7 @@
 // Confirm compilation of `(RW)ByteAddressBuffer` with aligned load / stores to wider data types.
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=buffer
-[vk::binding(2, 3)] RWByteAddressBuffer buffer;
+RWByteAddressBuffer buffer;
 struct Block {
     float4 val[2];
 };
diff --git a/tests/optimization/arrray-storage-lowering.slang b/tests/optimization/arrray-storage-lowering.slang
new file mode 100644
index 000000000..42bb8f127
--- /dev/null
+++ b/tests/optimization/arrray-storage-lowering.slang
@@ -0,0 +1,42 @@
+// TEST:SIMPLE(filecheck=SPV): -target spirv
+
+// TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-vk -compute -shaderobj -output-using-type -emit-spirv-directly
+
+struct DoubleNested
+{
+    int4x3 matrix;
+    int getMatVal(int i, int j) { return matrix[i][j]; }
+}
+
+struct Nested
+{
+    bool values[4];
+    DoubleNested doubleNested;
+    int getVal(int id) { return (int)values[0] + doubleNested.getMatVal(0, 1); }
+}
+
+struct Params
+{
+    Nested nested;
+
+    int getVal(int id) { return nested.getVal(id) + nested.getVal(id + 1); }
+}
+
+// TEST_INPUT: set outputBuffer = out ubuffer(data=[0], stride=4)
+RWStructuredBuffer<int4> outputBuffer;
+
+// TEST_INPUT:set gParams = cbuffer(data=[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1])
+ConstantBuffer<Params> gParams;
+
+// TEST_INPUT: set gDoubleNested = ubuffer(data=[1 2 3 4 5 6 7 8 9 10 11 12])
+uniform DoubleNested *gDoubleNested;
+
+// CHECK: 9
+
+[numthreads(1,1,1)]
+void computeMain(int id: SV_DispatchThreadID)
+{
+    outputBuffer[0].xyz = gParams.getVal(id) + gDoubleNested.getMatVal(1, 1);
+}
+
+// SPV-NOT: OpCompositeConstruct
diff --git a/tests/optimization/get-array-element.slang b/tests/optimization/get-array-element.slang
new file mode 100644
index 000000000..16a71aee2
--- /dev/null
+++ b/tests/optimization/get-array-element.slang
@@ -0,0 +1,17 @@
+//TEST:SIMPLE(filecheck=CHECK):-target spirv
+
+int test(int arr[32]) {
+    int sum = 0;
+    for (int i =0; i < 32; i++) sum += arr[i];
+    return sum;
+}
+
+uniform int gArr[32];
+uniform int* result;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    // CHECK-NOT: OpCompositeConstruct
+    *result = test(gArr);
+}
\ No newline at end of file
diff --git a/tests/pipeline/rasterization/get-attribute-at-vertex.slang b/tests/pipeline/rasterization/get-attribute-at-vertex.slang
index 342796b90..f73ebe86f 100644
--- a/tests/pipeline/rasterization/get-attribute-at-vertex.slang
+++ b/tests/pipeline/rasterization/get-attribute-at-vertex.slang
@@ -9,8 +9,11 @@
 // CHECK-SPIRV: OpExtension "SPV_KHR_fragment_shader_barycentric"
 
 // CHECK-SPIRV: OpEntryPoint Fragment %main "main"
-// CHECK-SPIRV: OpDecorate %{{.*}} BuiltIn BaryCoordKHR
-// CHECK-SPIRV: OpDecorate %{{.*}} BuiltIn BaryCoordNoPerspKHR
+
+// CHECK-SPIRV-DAG: OpDecorate %{{.*}} BuiltIn BaryCoordKHR
+
+// CHECK-SPIRV-DAG: OpDecorate %{{.*}} BuiltIn BaryCoordNoPerspKHR
+
 // CHECK-SPIRV: %{{.*}} = OpAccessChain %_ptr_Input_{{.*}} %{{.*}} %uint_0
 // CHECK-SPIRV: %{{.*}} = OpAccessChain %_ptr_Input_{{.*}} %{{.*}} %uint_1
 // CHECK-SPIRV: %{{.*}} = OpAccessChain %_ptr_Input_{{.*}} %{{.*}} %uint_2
diff --git a/tests/spirv/aligned-load-store.slang b/tests/spirv/aligned-load-store.slang
index c2f50b66c..8131e1cc7 100644
--- a/tests/spirv/aligned-load-store.slang
+++ b/tests/spirv/aligned-load-store.slang
@@ -4,8 +4,6 @@
 // CHECK: OpStore {{.*}} Aligned 16
 
 // CHECK: OpLoad {{.*}} Aligned 16
-// CHECK: OpLoad {{.*}} Aligned 16
-// CHECK: OpStore {{.*}} Aligned 16
 // CHECK: OpStore {{.*}} Aligned 16
 
 uniform float4* data;
diff --git a/tests/spirv/buffer-pointer-matrix-layout.slang b/tests/spirv/buffer-pointer-matrix-layout.slang
index cbb8f2857..4d80419e4 100644
--- a/tests/spirv/buffer-pointer-matrix-layout.slang
+++ b/tests/spirv/buffer-pointer-matrix-layout.slang
@@ -1,33 +1,50 @@
-//TEST:SIMPLE(filecheck=CHECK): -target spirv -emit-spirv-directly -stage compute -entry main -matrix-layout-column-major
+//TEST:COMPARE_COMPUTE(filecheck-buffer=CHECK): -vk -output-using-type -xslang -matrix-layout-column-major -emit-spirv-directly
 
-// CHECK: OpLoad {{.*}} Aligned 4
+// TEST_INPUT: set ptr = ubuffer(data=[1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0 11.0 12.0],stride=4)
+uniform float3x4 *ptr;
 
-struct Push
-{
-    float3x4* ptr;
-};
+// TEST_INPUT: set outputBuffer = out ubuffer(data=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0],stride=4)
+RWStructuredBuffer<float> outputBuffer;
 
-[[vk::push_constant]] Push push;
 [shader("compute")]
 [numthreads(1, 1, 1)]
-void main(uint3 dtid : SV_DispatchThreadID)
-{    
+void computeMain(uint3 dtid: SV_DispatchThreadID)
+{
     // This matrix is in memry column major. Slang respects this here and load it properly!
-    float3x4 correctly_read_matrix = *push.ptr;
-    printf("(%f,%f,%f,%f)\n(%f,%f,%f,%f)\n",
-        correctly_read_matrix[0][0], correctly_read_matrix[0][1], correctly_read_matrix[0][2], correctly_read_matrix[0][3],
-        correctly_read_matrix[1][0], correctly_read_matrix[1][1], correctly_read_matrix[1][2], correctly_read_matrix[1][3]
-    );
-    printf("(%f,%f,%f,%f)\n\n",
-        correctly_read_matrix[2][0], correctly_read_matrix[2][1], correctly_read_matrix[2][2], correctly_read_matrix[2][3]
-    );
-    // With this syntax however, Slang ignores the column major setting and loads it as it it was row major!
-    float3x4 broken_matrix = push.ptr[0];
-    printf("(%f,%f,%f,%f)\n(%f,%f,%f,%f)\n",
-        broken_matrix[0][0], broken_matrix[0][1], broken_matrix[0][2], broken_matrix[0][3],
-        broken_matrix[1][0], broken_matrix[1][1], broken_matrix[1][2], broken_matrix[1][3]
-    );
-    printf("(%f,%f,%f,%f)\n\n",
-        broken_matrix[2][0], broken_matrix[2][1], broken_matrix[2][2], broken_matrix[2][3]
-    );
+    float3x4 correctly_read_matrix = *ptr;
+    outputBuffer[0] = correctly_read_matrix[0][0];
+    outputBuffer[1] = correctly_read_matrix[0][1];
+    outputBuffer[2] = correctly_read_matrix[0][2];
+    outputBuffer[3] = correctly_read_matrix[0][3];
+    outputBuffer[4] = correctly_read_matrix[1][0];
+    outputBuffer[5] = correctly_read_matrix[1][1];
+    outputBuffer[6] = correctly_read_matrix[1][2];
+    outputBuffer[7] = correctly_read_matrix[1][3];
+    // CHECK: 1.0
+    // CHECK: 4.0
+    // CHECK: 7.0
+    // CHECK: 10.0
+    // CHECK: 2.0
+    // CHECK: 5.0
+    // CHECK: 8.0
+    // CHECK: 11.0
+
+    // With this syntax however, Slang was ignoring the column major setting and loads it as it it was row major!
+    float3x4 broken_matrix = ptr[0];
+    outputBuffer[8] = broken_matrix[0][0];
+    outputBuffer[9] = broken_matrix[0][1];
+    outputBuffer[10] = broken_matrix[0][2];
+    outputBuffer[11] = broken_matrix[0][3];
+    outputBuffer[12] = broken_matrix[1][0];
+    outputBuffer[13] = broken_matrix[1][1];
+    outputBuffer[14] = broken_matrix[1][2];
+    outputBuffer[15] = broken_matrix[1][3];
+    // CHECK: 1.0
+    // CHECK: 4.0
+    // CHECK: 7.0
+    // CHECK: 10.0
+    // CHECK: 2.0
+    // CHECK: 5.0
+    // CHECK: 8.0
+    // CHECK: 11.0
 }
\ No newline at end of file
diff --git a/tests/spirv/geometry-shader-sub-func.slang b/tests/spirv/geometry-shader-sub-func.slang
index 6c6944f31..20634ea67 100644
--- a/tests/spirv/geometry-shader-sub-func.slang
+++ b/tests/spirv/geometry-shader-sub-func.slang
@@ -36,7 +36,7 @@ void main(
     {
         CoarseVertex coarseVertex = coarseVertices[ii];
         RasterVertex rasterVertex;
-        rasterVertex.position   = coarseVertex.position;
+        rasterVertex.position = coarseVertex.position;
         rasterVertex.color      = coarseVertex.color;
         rasterVertex.id         = coarseVertex.id + primitiveID;
         appendVertex(outputStream, rasterVertex);
diff --git a/tests/spirv/large-struct.slang b/tests/spirv/large-struct.slang
index 2d79c0aaf..e4cbd6d1c 100644
--- a/tests/spirv/large-struct.slang
+++ b/tests/spirv/large-struct.slang
@@ -1,5 +1,5 @@
 //TEST:SIMPLE(filecheck=CHECK): -target spirv -emit-spirv-directly -profile glsl_460
-//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-vk -compute -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-vk -compute -output-using-type -xslang -g0
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-d3d12 -compute -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-cpu -compute -output-using-type
 
diff --git a/tests/spirv/pointer-2.slang b/tests/spirv/pointer-2.slang
index 1f2b2d0ea..b93ca32b4 100644
--- a/tests/spirv/pointer-2.slang
+++ b/tests/spirv/pointer-2.slang
@@ -1,6 +1,6 @@
-//TEST:SIMPLE(filecheck=CHECK_SPV): -entry vertexMain -stage vertex -emit-spirv-directly -target spirv
-//TEST:SIMPLE(filecheck=CHECK_SPV_VIA_GLSL): -entry vertexMain -stage vertex -emit-spirv-via-glsl -target spirv
-//TEST:SIMPLE(filecheck=CHECK_GLSL): -entry vertexMain -stage vertex -target glsl
+// TEST:SIMPLE(filecheck=CHECK_GLSL): -entry vertexMain -stage vertex -target glsl
+// TEST:SIMPLE(filecheck=CHECK_SPV): -entry vertexMain -stage vertex -emit-spirv-directly -target spirv
+// TEST:SIMPLE(filecheck=CHECK_SPV_VIA_GLSL): -entry vertexMain -stage vertex -emit-spirv-via-glsl -target spirv
 
 struct Inner1
 {
diff --git a/tests/spirv/spec-constant-operations.slang b/tests/spirv/spec-constant-operations.slang
index 86d16ef34..7c7b9bc60 100644
--- a/tests/spirv/spec-constant-operations.slang
+++ b/tests/spirv/spec-constant-operations.slang
@@ -11,8 +11,6 @@ RWStructuredBuffer<float> outputBuffer;
 // CHECK-DAG: OpSpecConstant %float 1
 // CHECK-DAG: OpSpecConstant %ulong 256
 // CHECK-DAG: OpSpecConstant %float 100
-// CHECK-DAG: OpSpecConstantOp %half FConvert
-// CHECK-DAG: OpSpecConstantOp %int UConvert
 
 // CHECK-NOT: OpSpecConstantOp {{.*}} FAdd
 // CHECK-NOT: OpSpecConstantOp {{.*}} FSub
diff --git a/tests/spirv/spirv-debug-break.slang b/tests/spirv/spirv-debug-break.slang
index e57024037..67b3e975b 100644
--- a/tests/spirv/spirv-debug-break.slang
+++ b/tests/spirv/spirv-debug-break.slang
@@ -1,5 +1,5 @@
 // spirv-instruction.slang
-//TEST(compute, vulkan):SIMPLE(filecheck=CHECK):-target glsl -entry computeMain -stage compute
+// TEST:SIMPLE(filecheck=CHECK):-target glsl -entry computeMain -stage compute
 
 [[vk::spirv_instruction(1, "NonSemantic.DebugBreak")]]
 void _spvDebugBreak(int v);