[CodeGen] Add sources to isVectorClearMaskLegal. NFC.

[HanKuanChen]

Summary

Targets may require source information to know whether a mask is legal.
(e.g., whether sources are from a same vector). But current
isVectorClearMaskLegal only provides mask and value type, which is not
enough.

Because DAG.getBitcast may add use number of operand (which blocks some
optimization), we will remove the bitcast node if the node is not used.

[RKSimon]

Do you have an upstream target that actually requires this?

[HanKuanChen]

In D137635#3914902, @RKSimon wrote:

Do you have an upstream target that actually requires this?

Well, actually I want to modify isShuffleMaskLegal so that RISCV target can recognize same vector_shuffle patterns if the source are from a same vector (or one of input is undef).
But in some targets (X86 and AArch64), isVectorClearMaskLegal will call isShuffleMaskLegal. If we modify isShuffleMaskLegal first, the implementations in isVectorClearMaskLegal will be very strange.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23166

How nasty would it be NOT to add the bitcast here at all and allow the N0/N1 operands into isVectorClearMaskLegal to keep their original types as long as they can safely be bitcasted to IntVT?

We'd still bitcast it as part of the getVectorShuffle call of course.

[spatel]

Would the problem be solved if we pulled this code:
https://github.com/llvm/llvm-project/blob/aa37342b3b548a75a4d8ce330f58a361a9ce22f4/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp#L22893
into a helper that could be used to see the simpler mask before it is actually created in a new shuffle node?

I was thinking about a similar utility in IR, so we don't have to worry as much about non-canonical masks in IR vector passes (this came up in D137341).

[HanKuanChen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23166

How nasty would it be NOT to add the bitcast here at all and allow the N0/N1 operands into isVectorClearMaskLegal to keep their original types as long as they can safely be bitcasted to IntVT?

Is the reason we do this because we don't want a DAG.RemoveDeadNode in the end?
I wonder will this make people confused since it violates the vector_shuffle interface (return and inputs should have same type).

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23166

Mainly I'm expecting this to go wrong in the future :-( I'm not sure we should be encouraging creation/deletion of temporary nodes like this just for an analysis call.

llvm/include/llvm/CodeGen/TargetLowering.h

public:
   ///
   /// VT must be a legal type. By default, we optimistically assume most
   /// operations don't trap except for integer divide and remainder.
   virtual bool canOpTrap(unsigned Op, EVT VT) const;
 
   /// Similar to isShuffleMaskLegal. Targets can use this to indicate if there
   /// is a suitable VECTOR_SHUFFLE that can be used to replace a VAND with a
   /// constant pool entry.
-  virtual bool isVectorClearMaskLegal(ArrayRef<int> /*Mask*/,
+  virtual bool isVectorClearMaskLegal(SDValue /*N0*/, SDValue /*N1*/,
+                                      ArrayRef<int> /*Mask*/,
                                       EVT /*VT*/) const {
     return false;
   }
 
   /// How to legalize this custom operation?
   virtual LegalizeAction getCustomOperationAction(SDNode &Op) const {
     return Legal;
   }

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (!IsInLaneMask &&
       if (IntSVT.getSizeInBits() >= IntVT.getScalarSizeInBits()) {
         SDValue ZeroElt = DAG.getConstant(0, DL, IntSVT);
         SDValue AllOnesElt = DAG.getAllOnesConstant(DL, IntSVT);
         SmallVector<SDValue, 16> AndMask(NumElts, DAG.getUNDEF(IntSVT));
         for (int I = 0; I != (int)NumElts; ++I)
           if (0 <= Mask[I])
             AndMask[I] = Mask[I] == I ? AllOnesElt : ZeroElt;
 
+        SDValue BitCastN0 = DAG.getBitcast(IntVT, N0);
+        SDValue ZeroInIntVT = DAG.getConstant(0, DL, IntVT);
         // See if a clear mask is legal instead of going via
         // XformToShuffleWithZero which loses UNDEF mask elements.
-        if (TLI.isVectorClearMaskLegal(ClearMask, IntVT))
-          return DAG.getBitcast(
-              VT, DAG.getVectorShuffle(IntVT, DL, DAG.getBitcast(IntVT, N0),
-                                      DAG.getConstant(0, DL, IntVT), ClearMask));
+        if (TLI.isVectorClearMaskLegal(BitCastN0, ZeroInIntVT, ClearMask,
+                                       IntVT))
+          return DAG.getBitcast(VT,
+                                DAG.getVectorShuffle(IntVT, DL, BitCastN0,
+                                                     ZeroInIntVT, ClearMask));
 
         if (TLI.isOperationLegalOrCustom(ISD::AND, IntVT))
           return DAG.getBitcast(
-              VT, DAG.getNode(ISD::AND, DL, IntVT, DAG.getBitcast(IntVT, N0),
+              VT, DAG.getNode(ISD::AND, DL, IntVT, BitCastN0,
                               DAG.getBuildVector(IntVT, DL, AndMask)));
+
+        if (BitCastN0.use_empty())
+          DAG.RemoveDeadNode(BitCastN0.getNode());
       }
     }
   }
 
   // Attempt to combine a shuffle of 2 inputs of 'scalar sources' -
   // BUILD_VECTOR or SCALAR_TO_VECTOR into a single BUILD_VECTOR.
   if (Level < AfterLegalizeDAG && TLI.isTypeLegal(VT))
     if (SDValue Res = combineShuffleOfScalars(SVN, DAG, TLI))
for (int i = 0; i != NumSubElts; ++i) {
         Indices.push_back(i + NumSubElts);
       else
         return SDValue();
     }
 
     // Let's see if the target supports this vector_shuffle.
     EVT ClearSVT = EVT::getIntegerVT(*DAG.getContext(), NumSubBits);
     EVT ClearVT = EVT::getVectorVT(*DAG.getContext(), ClearSVT, NumSubElts);
-    if (!TLI.isVectorClearMaskLegal(Indices, ClearVT))
+    SDValue BitcastLHS = DAG.getBitcast(ClearVT, LHS);
+    SDValue Zero = DAG.getConstant(0, DL, ClearVT);
+    if (!TLI.isVectorClearMaskLegal(BitcastLHS, Zero, Indices, ClearVT)) {
+      if (BitcastLHS.use_empty())
+        DAG.RemoveDeadNode(BitcastLHS.getNode());
       return SDValue();
+    }
 
-    SDValue Zero = DAG.getConstant(0, DL, ClearVT);
-    return DAG.getBitcast(VT, DAG.getVectorShuffle(ClearVT, DL,
-                                                   DAG.getBitcast(ClearVT, LHS),
-                                                   Zero, Indices));
+    return DAG.getBitcast(
+        VT, DAG.getVectorShuffle(ClearVT, DL, BitcastLHS, Zero, Indices));
   };
 
   // Determine maximum split level (byte level masking).
   int MaxSplit = 1;
   if (RVT.getScalarSizeInBits() % 8 == 0)
     MaxSplit = RVT.getScalarSizeInBits() / 8;
 
   for (int Split = 1; Split <= MaxSplit; ++Split)

llvm/lib/Target/AArch64/AArch64ISelLowering.h

bool isFPImmLegal(const APFloat &Imm, EVT VT,
                     bool ForCodeSize) const override;
 
   /// Return true if the given shuffle mask can be codegen'd directly, or if it
   /// should be stack expanded.
   bool isShuffleMaskLegal(ArrayRef<int> M, EVT VT) const override;
 
   /// Similar to isShuffleMaskLegal. Return true is the given 'select with zero'
   /// shuffle mask can be codegen'd directly.
-  bool isVectorClearMaskLegal(ArrayRef<int> M, EVT VT) const override;
+  bool isVectorClearMaskLegal(SDValue N0, SDValue N1, ArrayRef<int> M,
+                              EVT VT) const override;
 
   /// Return the ISD::SETCC ValueType.
   EVT getSetCCResultType(const DataLayout &DL, LLVMContext &Context,
                          EVT VT) const override;
 
   SDValue ReconstructShuffle(SDValue Op, SelectionDAG &DAG) const;
 
   MachineBasicBlock *EmitF128CSEL(MachineInstr &MI,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

return (ShuffleVectorSDNode::isSplatMask(&M[0], VT) || isREVMask(M, VT, 64) ||
           isZIPMask(M, VT, DummyUnsigned) ||
           isTRN_v_undef_Mask(M, VT, DummyUnsigned) ||
           isUZP_v_undef_Mask(M, VT, DummyUnsigned) ||
           isZIP_v_undef_Mask(M, VT, DummyUnsigned) ||
           isINSMask(M, VT.getVectorNumElements(), DummyBool, DummyInt) ||
           isConcatMask(M, VT, VT.getSizeInBits() == 128));
 }
 
-bool AArch64TargetLowering::isVectorClearMaskLegal(ArrayRef<int> M,
+bool AArch64TargetLowering::isVectorClearMaskLegal(SDValue N0, SDValue N1,
+                                                   ArrayRef<int> M,
                                                    EVT VT) const {
   // Just delegate to the generic legality, clear masks aren't special.
   return isShuffleMaskLegal(M, VT);
 }
 
 /// getVShiftImm - Check if this is a valid build_vector for the immediate
 /// operand of a vector shift operation, where all the elements of the
 /// build_vector must have the same constant integer value.

llvm/lib/Target/X86/X86ISelLowering.h

public:
     /// VECTOR_SHUFFLE operations, those with specific masks. By default, if a
     /// target supports the VECTOR_SHUFFLE node, all mask values are assumed to
     /// be legal.
     bool isShuffleMaskLegal(ArrayRef<int> Mask, EVT VT) const override;
 
     /// Similar to isShuffleMaskLegal. Targets can use this to indicate if there
     /// is a suitable VECTOR_SHUFFLE that can be used to replace a VAND with a
     /// constant pool entry.
-    bool isVectorClearMaskLegal(ArrayRef<int> Mask, EVT VT) const override;
+    bool isVectorClearMaskLegal(SDValue N0, SDValue N1, ArrayRef<int> Mask,
+                                EVT VT) const override;
 
     /// Returns true if lowering to a jump table is allowed.
     bool areJTsAllowed(const Function *Fn) const override;
 
     MVT getPreferredSwitchConditionType(LLVMContext &Context,
                                         EVT ConditionVT) const override;
 
     /// If true, then instruction selection should

llvm/lib/Target/X86/X86ISelLowering.cpp

   if (VT.getSimpleVT().getSizeInBits() == 64)
     return false;
 
   // We only care that the types being shuffled are legal. The lowering can
   // handle any possible shuffle mask that results.
   return isTypeLegal(VT.getSimpleVT());
 }
 
-bool X86TargetLowering::isVectorClearMaskLegal(ArrayRef<int> Mask,
+bool X86TargetLowering::isVectorClearMaskLegal(SDValue N0, SDValue N1,
+                                               ArrayRef<int> Mask,
                                                EVT VT) const {
   // Don't convert an 'and' into a shuffle that we don't directly support.
   // vpblendw and vpshufb for 256-bit vectors are not available on AVX1.
   if (!Subtarget.hasAVX2())
     if (VT == MVT::v32i8 || VT == MVT::v16i16)
       return false;
 
   // Just delegate to the generic legality, clear masks aren't special.