Fix shuffle feeding shuffle with undef literal (#4883)

When folding a vector shuffle with an undef literal, it is possible that the
literal is adjusted so that it will then be interpreted as an index into
the input operands.  This is fixed by special casing that case, and not
adjusting those operands.

Fixes #4859

2 files changed, 23 insertions, 1 deletions

diff --git a/source/opt/folding_rules.cpp b/source/opt/folding_rules.cpp
index 2d778b94a..3f10bd009 100644
--- a/source/opt/folding_rules.cpp
+++ b/source/opt/folding_rules.cpp
@@ -2748,7 +2748,8 @@ FoldingRule VectorShuffleFeedingShuffle() {
 
       if (adjustment != 0) {
         for (uint32_t i = 2; i < new_operands.size(); i++) {
-          if (inst->GetSingleWordInOperand(i) >= op0_length) {
+          uint32_t operand = inst->GetSingleWordInOperand(i);
+          if (operand >= op0_length && operand != undef_literal) {
             new_operands[i].words[0] -= adjustment;
           }
         }
diff --git a/test/opt/fold_test.cpp b/test/opt/fold_test.cpp
index f6c94ca05..a034e959a 100644
--- a/test/opt/fold_test.cpp
+++ b/test/opt/fold_test.cpp
@@ -7963,6 +7963,27 @@ INSTANTIATE_TEST_SUITE_P(VectorShuffleMatchingTest, MatchingInstructionWithNoRes
             "%9 = OpVectorShuffle %v4double %7 %8 2 0 1 4294967295\n" +
             "OpReturn\n" +
             "OpFunctionEnd",
+        9, true),
+    // Test case 14: Shuffle with undef literal and change size of first input vector.
+    InstructionFoldingCase<bool>(
+        Header() +
+            "; CHECK: [[double:%\\w+]] = OpTypeFloat 64\n" +
+            "; CHECK: [[v4double:%\\w+]] = OpTypeVector [[double]] 2\n" +
+            "; CHECK: OpVectorShuffle\n" +
+            "; CHECK: OpVectorShuffle {{%\\w+}} %5 %7 0 1 4 4294967295\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 1\n" +
+            "%9 = OpVectorShuffle %v4double %8 %7 0 1 2 4294967295\n" +
+            "OpReturn\n" +
+            "OpFunctionEnd",
         9, true)
 ));