[CodeGen][AMDGPU] EXTRACT_VECTOR_ELT: input vector element type can differ from output type

[jmmartinez]

Summary

In function SITargetLowering::performExtractVectorElt,
the output type was not considered which could lead to type mismatches
later.

[jmmartinez]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10817

This condition has to be stengtened. This is valid for smin, smax, umin and umax if the types have the same size or if they are truncated.

[arsenm]

Implicit element conversion strikes again. LGTM but needs tests

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10817

Not sure what you mean, those are handled below?

[jmmartinez]

Corrected binop case

Still missing tests for all of this

[jmmartinez]

In D139943#3992295, @arsenm wrote:

Implicit element conversion strikes again. LGTM but needs tests

Sure sorry! I didn't know phabricatior sent notifications when the patch is marked as draft. I submitted it in draft mode and I hoped to submit it after adding the tests.

[jmmartinez]

• Added test (I found the test a bit awkward but I haven't managed to trigger it otherwise)

[jmmartinez]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10817

I meant that when ResVT != VecEltVT the transformation could not be valid.

For example, if ResVT == i8 but VecEltVT == i16 for the min/max operations.

I wasn't able to trigger this path, so I haven't added the code to do the transformation. Still, I added a check in the if anyways.

llvm/test/CodeGen/AMDGPU/dagcomb-extract-vec-elt-different-sizes.ll

1 ↗

(On Diff #484200)

@arsenm If you have some advice about how to reduce even more this test case I'll be glad.

I tried triggering this case by hand but I didn't succeed. In the end I used llvm-reduce and applied some optimizations to get up to this.

[jmmartinez]

ping

[arsenm]

llvm/test/CodeGen/AMDGPU/dagcomb-extract-vec-elt-different-sizes.ll

2 ↗

(On Diff #484200)

Should drop the -verify-machineinstrs, I've been thinking we should drop most of these in tests

5–7 ↗

(On Diff #484200)

Don't need this stuff

9 ↗

(On Diff #484200)

Should use opaque pointers

93 ↗

(On Diff #484200)

I'm sure you don't need the addrspacecasts here, should just change the address spaces of the arguments (We should have an optimization and have llvm-reduce start doing this too)

110 ↗

(On Diff #484200)

Does it only happen with 1 x vectors, and do you need so many elements?

119 ↗

(On Diff #484200)

Can you hit this with a few more vector types?

[jmmartinez]

• Remove -verify-machineinstrs
• Remove source_filename, datalayout and triple
• Make test pointers opaque
• Remove addresspace cast and move addrspace to the declaration

llvm/test/CodeGen/AMDGPU/dagcomb-extract-vec-elt-different-sizes.ll

110 ↗

(On Diff #484200)

It seems that <1 x i8> is the only case where the bug gets triggered: I tried using directly i8; using <8 x i1>; changing the int type to i16, i32 or float; or changing the number of elements in the vector. I also tried removing some of the store instructions or replacing some of its arguments by another.

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

 }
 
 SDValue SITargetLowering::performExtractVectorEltCombine(
   SDNode *N, DAGCombinerInfo &DCI) const {
   SDValue Vec = N->getOperand(0);
   SelectionDAG &DAG = DCI.DAG;
 
   EVT VecVT = Vec.getValueType();
-  EVT EltVT = VecVT.getVectorElementType();
+  EVT VecEltVT = VecVT.getVectorElementType();
+  EVT ResVT = N->getValueType(0);
+
+  unsigned VecSize = VecVT.getSizeInBits();
+  unsigned VecEltSize = VecEltVT.getSizeInBits();
 
   if ((Vec.getOpcode() == ISD::FNEG ||
        Vec.getOpcode() == ISD::FABS) && allUsesHaveSourceMods(N)) {
     SDLoc SL(N);
-    EVT EltVT = N->getValueType(0);
     SDValue Idx = N->getOperand(1);
-    SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT,
-                              Vec.getOperand(0), Idx);
-    return DAG.getNode(Vec.getOpcode(), SL, EltVT, Elt);
+    SDValue Elt =
+        DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, ResVT, Vec.getOperand(0), Idx);
+    return DAG.getNode(Vec.getOpcode(), SL, ResVT, Elt);
   }
 
   // ScalarRes = EXTRACT_VECTOR_ELT ((vector-BINOP Vec1, Vec2), Idx)
   //    =>
   // Vec1Elt = EXTRACT_VECTOR_ELT(Vec1, Idx)
   // Vec2Elt = EXTRACT_VECTOR_ELT(Vec2, Idx)
   // ScalarRes = scalar-BINOP Vec1Elt, Vec2Elt
   if (Vec.hasOneUse() && DCI.isBeforeLegalize()) {
     SDLoc SL(N);
-    EVT EltVT = N->getValueType(0);
     SDValue Idx = N->getOperand(1);
     unsigned Opc = Vec.getOpcode();
 
+    bool ZeroExtOnTypeMismatch = true;
     switch(Opc) {
     default:
       break;
       // TODO: Support other binary operations.
+    case ISD::SMIN:
+    case ISD::SMAX:
+      ZeroExtOnTypeMismatch = false;
+      LLVM_FALLTHROUGH;
     case ISD::FADD:
     case ISD::FSUB:
     case ISD::FMUL:
     case ISD::ADD:
     case ISD::UMIN:
     case ISD::UMAX:
-    case ISD::SMIN:
-    case ISD::SMAX:
     case ISD::FMAXNUM:
     case ISD::FMINNUM:
     case ISD::FMAXNUM_IEEE:
     case ISD::FMINNUM_IEEE: {
-      SDValue Elt0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT,
+      EVT ExtractEltTy = ResVT.getSizeInBits() <= VecEltSize ? ResVT : VecEltVT;
+      SDValue Elt0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, ExtractEltTy,
                                  Vec.getOperand(0), Idx);
-      SDValue Elt1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT,
+      SDValue Elt1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, ExtractEltTy,
                                  Vec.getOperand(1), Idx);
 
       DCI.AddToWorklist(Elt0.getNode());
       DCI.AddToWorklist(Elt1.getNode());
-      return DAG.getNode(Opc, SL, EltVT, Elt0, Elt1, Vec->getFlags());
+
+      if (ExtractEltTy != ResVT) {
+        auto CastOpc =
+            ZeroExtOnTypeMismatch ? ISD::ZERO_EXTEND : ISD::SIGN_EXTEND;
+        Elt0 = DAG.getNode(CastOpc, SL, ResVT, Elt0);
+        Elt1 = DAG.getNode(CastOpc, SL, ResVT, Elt1);
+        DCI.AddToWorklist(Elt0.getNode());
+        DCI.AddToWorklist(Elt1.getNode());
+      }
+
+      return DAG.getNode(Opc, SL, ResVT, Elt0, Elt1, Vec->getFlags());
     }
     }
   }
-
-  unsigned VecSize = VecVT.getSizeInBits();
-  unsigned EltSize = EltVT.getSizeInBits();
 
   // EXTRACT_VECTOR_ELT (<n x e>, var-idx) => n x select (e, const-idx)
   if (shouldExpandVectorDynExt(N)) {
     SDLoc SL(N);
     SDValue Idx = N->getOperand(1);
     SDValue V;
     for (unsigned I = 0, E = VecVT.getVectorNumElements(); I < E; ++I) {
       SDValue IC = DAG.getVectorIdxConstant(I, SL);
-      SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT, Vec, IC);
+      SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, ResVT, Vec, IC);
       if (I == 0)
         V = Elt;
       else
         V = DAG.getSelectCC(SL, Idx, IC, Elt, V, ISD::SETEQ);
     }
     return V;
   }
 
   if (!DCI.isBeforeLegalize())
     return SDValue();
 
   // Try to turn sub-dword accesses of vectors into accesses of the same 32-bit
   // elements. This exposes more load reduction opportunities by replacing
   // multiple small extract_vector_elements with a single 32-bit extract.
   auto *Idx = dyn_cast<ConstantSDNode>(N->getOperand(1));
-  if (isa<MemSDNode>(Vec) &&
-      EltSize <= 16 &&
-      EltVT.isByteSized() &&
-      VecSize > 32 &&
-      VecSize % 32 == 0 &&
-      Idx) {
+  if (isa<MemSDNode>(Vec) && VecEltSize <= 16 && VecEltVT.isByteSized() &&
+      VecSize > 32 && VecSize % 32 == 0 && Idx) {
     EVT NewVT = getEquivalentMemType(*DAG.getContext(), VecVT);
 
-    unsigned BitIndex = Idx->getZExtValue() * EltSize;
+    unsigned BitIndex = Idx->getZExtValue() * VecEltSize;
     unsigned EltIdx = BitIndex / 32;
     unsigned LeftoverBitIdx = BitIndex % 32;
     SDLoc SL(N);
 
     SDValue Cast = DAG.getNode(ISD::BITCAST, SL, NewVT, Vec);
     DCI.AddToWorklist(Cast.getNode());
 
     SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, MVT::i32, Cast,
                               DAG.getConstant(EltIdx, SL, MVT::i32));
     DCI.AddToWorklist(Elt.getNode());
     SDValue Srl = DAG.getNode(ISD::SRL, SL, MVT::i32, Elt,
                               DAG.getConstant(LeftoverBitIdx, SL, MVT::i32));
     DCI.AddToWorklist(Srl.getNode());
 
-    SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SL, EltVT.changeTypeToInteger(), Srl);
+    EVT VecEltAsIntVT = VecEltVT.changeTypeToInteger();
+    SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SL, VecEltAsIntVT, Srl);
     DCI.AddToWorklist(Trunc.getNode());
-    return DAG.getNode(ISD::BITCAST, SL, EltVT, Trunc);
+
+    if (VecEltVT == ResVT) {
+      return DAG.getNode(ISD::BITCAST, SL, VecEltVT, Trunc);
+    }
+
+    assert(ResVT.isScalarInteger());
+    return DAG.getAnyExtOrTrunc(Trunc, SL, ResVT);
   }
 
   return SDValue();
 }
 
 SDValue
 SITargetLowering::performInsertVectorEltCombine(SDNode *N,
                                                 DAGCombinerInfo &DCI) const {