Fix size error when folding vector shuffle. (#1721)

When folding a vector shuffle that feeds another vector shuffle causes
the size of the first operand to change, when other indices have to be
adjusted reletive to the new size.

3 files changed, 55 insertions, 1 deletions

diff --git a/source/opt/folding_rules.cpp b/source/opt/folding_rules.cpp
index 3d3cf31d..787af41c 100644
--- a/source/opt/folding_rules.cpp
+++ b/source/opt/folding_rules.cpp
@@ -2090,6 +2090,22 @@ FoldingRule VectorShuffleFeedingShuffle() {
     }
 
     if (feeder_is_op0) {
+      // If the size of the first vector operand changed then the indices
+      // referring to the second operand need to be adjusted.
+      Instruction* new_feeder_inst = def_use_mgr->GetDef(new_feeder_id);
+      analysis::Type* new_feeder_type =
+          type_mgr->GetType(new_feeder_inst->type_id());
+      uint32_t new_op0_size = new_feeder_type->AsVector()->element_count();
+      int32_t adjustment = op0_length - new_op0_size;
+
+      if (adjustment != 0) {
+        for (uint32_t i = 2; i < new_operands.size(); i++) {
+          if (inst->GetSingleWordInOperand(i) >= op0_length) {
+            new_operands[i].words[0] -= adjustment;
+          }
+        }
+      }
+
       new_operands[0].words[0] = new_feeder_id;
       new_operands[1] = inst->GetInOperand(1);
     } else {
diff --git a/source/val/validate_id.cpp b/source/val/validate_id.cpp
index e1d76ac7..1ceab202 100644
--- a/source/val/validate_id.cpp
+++ b/source/val/validate_id.cpp
@@ -1781,7 +1781,7 @@ bool idUsage::isValid<SpvOpVectorShuffle>(const spv_instruction_t* inst,
   for (size_t i = firstLiteralIndex; i < inst->words.size(); ++i) {
     auto literal = inst->words[i];
     if (literal != 0xFFFFFFFF && literal >= N) {
-      DIAG(module_.FindDef(inst->words[i]))
+      DIAG(module_.FindDef(inst->words[0]))
           << "Component literal value " << literal << " is greater than "
           << N - 1 << ".";
       return false;
diff --git a/test/opt/fold_test.cpp b/test/opt/fold_test.cpp
index ff4a1b59..3be8cde9 100644
--- a/test/opt/fold_test.cpp
+++ b/test/opt/fold_test.cpp
@@ -5816,6 +5816,44 @@ INSTANTIATE_TEST_CASE_P(VectorShuffleMatchingTest, MatchingInstructionWithNoResu
             "%9 = OpVectorShuffle %v4double %7 %8 2 0 7 5\n" +
             "OpReturn\n" +
             "OpFunctionEnd",
+        9, true),
+    // Test case 9: Replace first operand with a smaller vector.
+    InstructionFoldingCase<bool>(
+        Header() +
+            "; CHECK: OpVectorShuffle\n" +
+            "; CHECK: OpVectorShuffle {{%\\w+}} %5 %7 0 0 5 3\n" +
+            "; CHECK: OpReturn\n" +
+            "%main = OpFunction %void None %void_func\n" +
+            "%main_lab = OpLabel\n" +
+            "%2 = OpVariable %_ptr_v2double Function\n" +
+            "%3 = OpVariable %_ptr_v4double Function\n" +
+            "%4 = OpVariable %_ptr_v4double Function\n" +
+            "%5 = OpLoad %v2double %2\n" +
+            "%6 = OpLoad %v4double %3\n" +
+            "%7 = OpLoad %v4double %4\n" +
+            "%8 = OpVectorShuffle %v4double %5 %5 0 1 2 3\n" +
+            "%9 = OpVectorShuffle %v4double %8 %7 2 0 7 5\n" +
+            "OpReturn\n" +
+            "OpFunctionEnd",
+        9, true),
+    // Test case 10: Replace first operand with a larger vector.
+    InstructionFoldingCase<bool>(
+        Header() +
+            "; CHECK: OpVectorShuffle\n" +
+            "; CHECK: OpVectorShuffle {{%\\w+}} %5 %7 3 0 7 5\n" +
+            "; CHECK: OpReturn\n" +
+            "%main = OpFunction %void None %void_func\n" +
+            "%main_lab = OpLabel\n" +
+            "%2 = OpVariable %_ptr_v4double Function\n" +
+            "%3 = OpVariable %_ptr_v4double Function\n" +
+            "%4 = OpVariable %_ptr_v4double Function\n" +
+            "%5 = OpLoad %v4double %2\n" +
+            "%6 = OpLoad %v4double %3\n" +
+            "%7 = OpLoad %v4double %4\n" +
+            "%8 = OpVectorShuffle %v2double %5 %5 0 3\n" +
+            "%9 = OpVectorShuffle %v4double %8 %7 1 0 5 3\n" +
+            "OpReturn\n" +
+            "OpFunctionEnd",
         9, true)
 ));
 #endif