spirv-opt: Fix stacked CompositeExtract constant folds (#4932)

This was spotted in the Validation Layers where OpSpecConstantOp %x CompositeExtract %y 0 was being folded to a constant, but anything that was using it wasn't recognizing it as a constant, the simple fix was to add a const_mgr->MapInst(new_const_inst); so the next instruction knew it was a const

1 files changed, 204 insertions, 1 deletions

diff --git a/test/opt/fold_spec_const_op_composite_test.cpp b/test/opt/fold_spec_const_op_composite_test.cpp
index 7eddf7e99..c98a44c3d 100644
--- a/test/opt/fold_spec_const_op_composite_test.cpp
+++ b/test/opt/fold_spec_const_op_composite_test.cpp
@@ -105,6 +105,209 @@ TEST_F(FoldSpecConstantOpAndCompositePassBasicTest,
       builder.GetCode(), builder.GetCode(), /* skip_nop = */ true);
 }
 
+// Test where OpSpecConstantOp depends on another OpSpecConstantOp with
+// CompositeExtract
+TEST_F(FoldSpecConstantOpAndCompositePassBasicTest, StackedCompositeExtract) {
+  AssemblyBuilder builder;
+  builder.AppendTypesConstantsGlobals({
+      // clang-format off
+    "%uint = OpTypeInt 32 0",
+    "%v3uint = OpTypeVector %uint 3",
+    "%uint_2 = OpConstant %uint 2",
+    "%uint_3 = OpConstant %uint 3",
+    // Folding target:
+    "%composite_0 = OpSpecConstantComposite %v3uint %uint_2 %uint_3 %uint_2",
+    "%op_0 = OpSpecConstantOp %uint CompositeExtract %composite_0 0",
+    "%op_1 = OpSpecConstantOp %uint CompositeExtract %composite_0 1",
+    "%op_2 = OpSpecConstantOp %uint IMul %op_0 %op_1",
+    "%composite_1 = OpSpecConstantComposite %v3uint %op_0 %op_1 %op_2",
+    "%op_3 = OpSpecConstantOp %uint CompositeExtract %composite_1 0",
+    "%op_4 = OpSpecConstantOp %uint IMul %op_2 %op_3",
+      // clang-format on
+  });
+
+  std::vector<const char*> expected = {
+      // clang-format off
+        "OpCapability Shader",
+        "OpCapability Float64",
+    "%1 = OpExtInstImport \"GLSL.std.450\"",
+        "OpMemoryModel Logical GLSL450",
+        "OpEntryPoint Vertex %main \"main\"",
+        "OpName %void \"void\"",
+        "OpName %main_func_type \"main_func_type\"",
+        "OpName %main \"main\"",
+        "OpName %main_func_entry_block \"main_func_entry_block\"",
+        "OpName %uint \"uint\"",
+        "OpName %v3uint \"v3uint\"",
+        "OpName %uint_2 \"uint_2\"",
+        "OpName %uint_3 \"uint_3\"",
+        "OpName %composite_0 \"composite_0\"",
+        "OpName %op_0 \"op_0\"",
+        "OpName %op_1 \"op_1\"",
+        "OpName %op_2 \"op_2\"",
+        "OpName %composite_1 \"composite_1\"",
+        "OpName %op_3 \"op_3\"",
+        "OpName %op_4 \"op_4\"",
+    "%void = OpTypeVoid",
+"%main_func_type = OpTypeFunction %void",
+    "%uint = OpTypeInt 32 0",
+  "%v3uint = OpTypeVector %uint 3",
+  "%uint_2 = OpConstant %uint 2",
+  "%uint_3 = OpConstant %uint 3",
+"%composite_0 = OpConstantComposite %v3uint %uint_2 %uint_3 %uint_2",
+    "%op_0 = OpConstant %uint 2",
+    "%op_1 = OpConstant %uint 3",
+    "%op_2 = OpConstant %uint 6",
+"%composite_1 = OpConstantComposite %v3uint %op_0 %op_1 %op_2",
+"%op_3 = OpConstant %uint 2",
+ "%op_4 = OpConstant %uint 12",
+    "%main = OpFunction %void None %main_func_type",
+"%main_func_entry_block = OpLabel",
+            "OpReturn",
+            "OpFunctionEnd",
+      // clang-format on
+  };
+  SinglePassRunAndCheck<FoldSpecConstantOpAndCompositePass>(
+      builder.GetCode(), JoinAllInsts(expected), /* skip_nop = */ true);
+}
+
+// Test where OpSpecConstantOp depends on another OpSpecConstantOp with
+// VectorShuffle
+TEST_F(FoldSpecConstantOpAndCompositePassBasicTest, StackedVectorShuffle) {
+  AssemblyBuilder builder;
+  builder.AppendTypesConstantsGlobals({
+      // clang-format off
+    "%uint = OpTypeInt 32 0",
+    "%v3uint = OpTypeVector %uint 3",
+    "%uint_1 = OpConstant %uint 1",
+    "%uint_2 = OpConstant %uint 2",
+    "%uint_3 = OpConstant %uint 3",
+    "%uint_4 = OpConstant %uint 4",
+    "%uint_5 = OpConstant %uint 5",
+    "%uint_6 = OpConstant %uint 6",
+    // Folding target:
+    "%composite_0 = OpSpecConstantComposite %v3uint %uint_1 %uint_2 %uint_3",
+    "%composite_1 = OpSpecConstantComposite %v3uint %uint_4 %uint_5 %uint_6",
+    "%vecshuffle = OpSpecConstantOp %v3uint VectorShuffle %composite_0 %composite_1 0 5 3",
+    "%op = OpSpecConstantOp %uint CompositeExtract %vecshuffle 1",
+      // clang-format on
+  });
+
+  std::vector<const char*> expected = {
+      // clang-format off
+        "OpCapability Shader",
+        "OpCapability Float64",
+        "%1 = OpExtInstImport \"GLSL.std.450\"",
+        "OpMemoryModel Logical GLSL450",
+        "OpEntryPoint Vertex %main \"main\"",
+        "OpName %void \"void\"",
+        "OpName %main_func_type \"main_func_type\"",
+        "OpName %main \"main\"",
+        "OpName %main_func_entry_block \"main_func_entry_block\"",
+        "OpName %uint \"uint\"",
+        "OpName %v3uint \"v3uint\"",
+        "OpName %uint_1 \"uint_1\"",
+        "OpName %uint_2 \"uint_2\"",
+        "OpName %uint_3 \"uint_3\"",
+        "OpName %uint_4 \"uint_4\"",
+        "OpName %uint_5 \"uint_5\"",
+        "OpName %uint_6 \"uint_6\"",
+        "OpName %composite_0 \"composite_0\"",
+        "OpName %composite_1 \"composite_1\"",
+        "OpName %vecshuffle \"vecshuffle\"",
+        "OpName %op \"op\"",
+    "%void = OpTypeVoid",
+"%main_func_type = OpTypeFunction %void",
+    "%uint = OpTypeInt 32 0",
+  "%v3uint = OpTypeVector %uint 3",
+  "%uint_1 = OpConstant %uint 1",
+  "%uint_2 = OpConstant %uint 2",
+  "%uint_3 = OpConstant %uint 3",
+  "%uint_4 = OpConstant %uint 4",
+  "%uint_5 = OpConstant %uint 5",
+  "%uint_6 = OpConstant %uint 6",
+"%composite_0 = OpConstantComposite %v3uint %uint_1 %uint_2 %uint_3",
+"%composite_1 = OpConstantComposite %v3uint %uint_4 %uint_5 %uint_6",
+"%vecshuffle = OpConstantComposite %v3uint %uint_1 %uint_6 %uint_4",
+      "%op = OpConstant %uint 6",
+    "%main = OpFunction %void None %main_func_type",
+"%main_func_entry_block = OpLabel",
+        "OpReturn",
+        "OpFunctionEnd",
+      // clang-format on
+  };
+  SinglePassRunAndCheck<FoldSpecConstantOpAndCompositePass>(
+      builder.GetCode(), JoinAllInsts(expected), /* skip_nop = */ true);
+}
+
+// Test CompositeExtract with matrix
+TEST_F(FoldSpecConstantOpAndCompositePassBasicTest, CompositeExtractMaxtrix) {
+  AssemblyBuilder builder;
+  builder.AppendTypesConstantsGlobals({
+      // clang-format off
+    "%uint = OpTypeInt 32 0",
+    "%v3uint = OpTypeVector %uint 3",
+    "%mat3x3 = OpTypeMatrix %v3uint 3",
+    "%uint_1 = OpConstant %uint 1",
+    "%uint_2 = OpConstant %uint 2",
+    "%uint_3 = OpConstant %uint 3",
+    // Folding target:
+    "%a = OpSpecConstantComposite %v3uint %uint_1 %uint_1 %uint_1",
+    "%b = OpSpecConstantComposite %v3uint %uint_1 %uint_1 %uint_3",
+    "%c = OpSpecConstantComposite %v3uint %uint_1 %uint_2 %uint_1",
+    "%op = OpSpecConstantComposite %mat3x3 %a %b %c",
+    "%x = OpSpecConstantOp %uint CompositeExtract %op 2 1",
+    "%y = OpSpecConstantOp %uint CompositeExtract %op 1 2",
+      // clang-format on
+  });
+
+  std::vector<const char*> expected = {
+      // clang-format off
+        "OpCapability Shader",
+        "OpCapability Float64",
+   "%1 = OpExtInstImport \"GLSL.std.450\"",
+        "OpMemoryModel Logical GLSL450",
+        "OpEntryPoint Vertex %main \"main\"",
+        "OpName %void \"void\"",
+        "OpName %main_func_type \"main_func_type\"",
+        "OpName %main \"main\"",
+        "OpName %main_func_entry_block \"main_func_entry_block\"",
+        "OpName %uint \"uint\"",
+        "OpName %v3uint \"v3uint\"",
+        "OpName %mat3x3 \"mat3x3\"",
+        "OpName %uint_1 \"uint_1\"",
+        "OpName %uint_2 \"uint_2\"",
+        "OpName %uint_3 \"uint_3\"",
+        "OpName %a \"a\"",
+        "OpName %b \"b\"",
+        "OpName %c \"c\"",
+        "OpName %op \"op\"",
+        "OpName %x \"x\"",
+        "OpName %y \"y\"",
+    "%void = OpTypeVoid",
+"%main_func_type = OpTypeFunction %void",
+    "%uint = OpTypeInt 32 0",
+  "%v3uint = OpTypeVector %uint 3",
+  "%mat3x3 = OpTypeMatrix %v3uint 3",
+  "%uint_1 = OpConstant %uint 1",
+  "%uint_2 = OpConstant %uint 2",
+  "%uint_3 = OpConstant %uint 3",
+       "%a = OpConstantComposite %v3uint %uint_1 %uint_1 %uint_1",
+       "%b = OpConstantComposite %v3uint %uint_1 %uint_1 %uint_3",
+       "%c = OpConstantComposite %v3uint %uint_1 %uint_2 %uint_1",
+      "%op = OpConstantComposite %mat3x3 %a %b %c",
+       "%x = OpConstant %uint 2",
+       "%y = OpConstant %uint 3",
+    "%main = OpFunction %void None %main_func_type",
+"%main_func_entry_block = OpLabel",
+        "OpReturn",
+        "OpFunctionEnd",
+      // clang-format on
+  };
+  SinglePassRunAndCheck<FoldSpecConstantOpAndCompositePass>(
+      builder.GetCode(), JoinAllInsts(expected), /* skip_nop = */ true);
+}
+
 // All types and some common constants that are potentially required in
 // FoldSpecConstantOpAndCompositeTest.
 std::vector<std::string> CommonTypesAndConstants() {
@@ -199,7 +402,7 @@ std::string StripOpNameInstructions(const std::string& str) {
 struct FoldSpecConstantOpAndCompositePassTestCase {
   // Original constants with unfolded spec constants.
   std::vector<std::string> original;
-  // Expected cosntants after folding.
+  // Expected constant after folding.
   std::vector<std::string> expected;
 };