[DAGCombiner] Fix PR33368 (vector extend/truncate optimization)

[wolfgangp]

Summary

The DAGCombiner, when attempting to turn vector shuffles into vector in_reg_extends and truncates, eliminates extend/truncate sequences when the source and destination element sizes are the same. This fails when the scales of the extensions and truncations involved don't match.

This fixes PR33368.

[RKSimon]

If possible I'd like to see a regular llc test as well as the mir test.

[wolfgangp]

In D34069#777451, @RKSimon wrote:

If possible I'd like to see a regular llc test as well as the mir test.

Actually, I could simply remove the YAML markups and turn this into a regular llc test, if you prefer. The substance of it would not change.

[wolfgangp]

Replaced .mir test with a .ll test, although the test is substantially the same.

[wolfgangp]

Ping...

The test is now a .ll test, do you think this is sufficient? It's relatively tricky to isolate the condition under which this occurs. The original test case was a randomized test.

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15085 ↗

(On Diff #102213)

Add a check that (ExtDstSizeInBits % ExtSrcSizeInBits) == 0? Either an assert or early-out

[wolfgangp]

Addressed review comment: added check to ensure that the extension size is a multiple of the source's size.

[wolfgangp]

Ping...

[RKSimon]

LGTM

[zvi]

Hi @wolfgangp,

I'm afraid that D34076 adds a combine which replaces the BUILD_VECTOR with a TRUNCATE, so if D34076 will be applied, the test-case in this commit will no longer be covering the bug this commit fixes.
Do you have a suggestion for an alternative test case?

[wolfgangp]

In D34069#794414, @zvi wrote:

Hi @wolfgangp,

I'm afraid that D34076 adds a combine which replaces the BUILD_VECTOR with a TRUNCATE, so if D34076 will be applied, the test-case in this commit will no longer be covering the bug this commit fixes.
Do you have a suggestion for an alternative test case?

Unfortunately not. It may not be possible to create a test case that is sufficiently insulated from future changes to DAGCombine. The only thing that sounds reasonable would be a unit test for the specific routine that had the bug. That would probably be immune against other optimizations, but I'm not even sure that the infrastructure for DAGCombine unit tests exists. In any case, I would be happy with updating the test case to reflect the code that is generated with the new optimization in place, since it will still demonstrate that the right code is generated.

[zvi]

Err, the patch i am referring to is actually D34077 and not what i mentioned above. I will add you to the review to discuss what can be done about this test.

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (Opcode != ISD::ANY_EXTEND_VECTOR_INREG &&
       Opcode != ISD::ZERO_EXTEND_VECTOR_INREG)
     return SDValue();
 
   SDValue N00 = N0.getOperand(0);
   ArrayRef<int> Mask = SVN->getMask();
   unsigned NumElts = VT.getVectorNumElements();
   unsigned EltSizeInBits = VT.getScalarSizeInBits();
   unsigned ExtSrcSizeInBits = N00.getScalarValueSizeInBits();
+  unsigned ExtDstSizeInBits = N0.getScalarValueSizeInBits();
+
+  if (ExtDstSizeInBits % ExtSrcSizeInBits != 0)
+    return SDValue();
+  unsigned ExtScale = ExtDstSizeInBits / ExtSrcSizeInBits;
 
   // (v4i32 truncate_vector_inreg(v2i64)) == shuffle<0,2-1,-1>
   // (v8i16 truncate_vector_inreg(v4i32)) == shuffle<0,2,4,6,-1,-1,-1,-1>
   // (v8i16 truncate_vector_inreg(v2i64)) == shuffle<0,4,-1,-1,-1,-1,-1,-1>
   auto isTruncate = [&Mask, &NumElts](unsigned Scale) {
     for (unsigned i = 0; i != NumElts; ++i) {
       if (Mask[i] < 0)
         continue;
       if ((i * Scale) < NumElts && Mask[i] == (int)(i * Scale))
         continue;
       return false;
     }
     return true;
   };
 
   // At the moment we just handle the case where we've truncated back to the
   // same size as before the extension.
   // TODO: handle more extension/truncation cases as cases arise.
   if (EltSizeInBits != ExtSrcSizeInBits)
     return SDValue();
 
-  // Attempt to match a 'truncate_vector_inreg' shuffle, we just search for
-  // power-of-2 truncations as they are the most likely.
-  for (unsigned Scale = 2; Scale < NumElts; Scale *= 2)
-    if (isTruncate(Scale))
-      return DAG.getBitcast(VT, N00);
+  // We can remove *extend_vector_inreg only if the truncation happens at
+  // the same scale as the extension.
+  if (isTruncate(ExtScale))
 
   return SDValue();
 }
 
 // Combine shuffles of splat-shuffles of the form:
 // shuffle (shuffle V, undef, splat-mask), undef, M
 // If splat-mask contains undef elements, we need to be careful about
 // introducing undef's in the folded mask which are not the result of composing

llvm/trunk/test/CodeGen/X86/vector-truncate-combine.ll

+; RUN: llc -mtriple=x86_64-- -O2 -start-after=stack-protector -stop-before=loops %s -o - | FileCheck %s
+
+; This test verifies the fix for PR33368.
+;
+; The expected outcome of the operation is to store bytes 0 and 2 of the incoming
+; parameter into c2 (a 2 x i8 vector). DAGCombine converts shuffles into a
+; sequence of extend and subsequent truncate operations. The bug was that an extension
+; by 4 followed by a truncation by 8 was completely eliminated.
+
+; The test checks for the correct sequence of operations that results from the
+; preservation of the extend/truncate operations mentioned above (2 extend and
+; 3 truncate instructions).
+;
+; NOTE: This operation could be collapsed in to a single truncate. Once that is done
+; this test will have to be adjusted.
+
+; CHECK:      PUNPCKLBWrr
+; CHECK:      PUNPCKLWDrr
+; CHECK:      PACKUSWBrr
+; CHECK:      PACKUSWBrr
+; CHECK:      PACKUSWBrr
+
+define void @test(double %vec.coerce) local_unnamed_addr {
+entry:
+  %c2 = alloca <2 x i8>, align 2
+  %0 = bitcast double %vec.coerce to <8 x i8>
+  %1 = shufflevector <8 x i8> %0, <8 x i8> undef, <4 x i32> <i32 2, i32 undef, i32 undef, i32 0>
+  %2 = shufflevector <4 x i8> %1, <4 x i8> undef, <2 x i32> <i32 3, i32 0>
+  store volatile <2 x i8> %2, <2 x i8>* %c2, align 2
+  br label %if.end
+
+if.end:
+  %3 = bitcast <2 x i8> %2 to i16
+  ret void
+}