[DAG] Add functionality for masked truncating store actions

[DavidTruby]

Summary

This adds functionality for setting and getting the relevant action
for masked truncating stores in the same way as for standard truncating
stores. There should be no change to code generation for any architecture
at this stage.

This support is needed to correctly implement masked truncating stores for
SVE on AArch64.

[dmgreen]

Is this needed for SVE support of masked truncating stores? My point in D108115 was that we could probably (until someone came along with an opposite position) treat masked truncating stores and truncating stores as symmetric. This will be true for MVE and SVE I think, and is true for X86 as far as I understand. (A second opinion on that would be good though). I imagine RiscV is the same with how their vector architecture works, not sure about hexagon. Sorry for not being very clear on that. Perhaps it's best to add anyway, but the equivalent code for zext/sext masked loads uses isLoadExtLegalOrCustom for the two and hasn't needed to split them into normal/masked variants: https://github.com/llvm/llvm-project/blob/e42f5d4b488e78ebf5b756e1e76422c7458ba81c/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp#L10917

My only comment on D108115 was a minor nitpick with formatting and a question about how the AArch64 code was working. RKSimon had a comment about the difficult to update tests too.

On a purely technical note, masked loads/stores only operate on vector types as far as I understand, and TruncMStoreActions doesn't appear to be initialized anywhere.

[DavidTruby]

Add initialisation for TruncMStoreActions

[DavidTruby]

In D112536#3089509, @dmgreen wrote:

Is this needed for SVE support of masked truncating stores? My point in D108115 was that we could probably (until someone came along with an opposite position) treat masked truncating stores and truncating stores as symmetric. This will be true for MVE and SVE I think, and is true for X86 as far as I understand. (A second opinion on that would be good though). I imagine RiscV is the same with how their vector architecture works, not sure about hexagon. Sorry for not being very clear on that. Perhaps it's best to add anyway, but the equivalent code for zext/sext masked loads uses isLoadExtLegalOrCustom for the two and hasn't needed to split them into normal/masked variants: https://github.com/llvm/llvm-project/blob/e42f5d4b488e78ebf5b756e1e76422c7458ba81c/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp#L10917

My only comment on D108115 was a minor nitpick with formatting and a question about how the AArch64 code was working. RKSimon had a comment about the difficult to update tests too.

On a purely technical note, masked loads/stores only operate on vector types as far as I understand, and TruncMStoreActions doesn't appear to be initialized anywhere.

Ah ok, well this still feels like the correct thing to do to me. With regards to the extends, I was also thinking of looking at that and splitting the extend logic in the same way in future as there's no logical reason that masked and non-masked versions of these should be symmetric either.

[dmgreen]

OK. We have got this far without needing a MLoad version, so SVE doesn't seem like it would necessitate adding an TruncMStore. Perhaps it is worth adding at any rate. It is conceptually cleaner, but does involve keeping extra state in TargetLowering. I will defer to the what others on review think.

[peterwaller-arm]

LGTM, would it be a good idea to give this patch an [NFCI] tag before committing?

[paulwalker-arm]

llvm/include/llvm/CodeGen/TargetLowering.h

1302–1303

This comment does not make sense to me. What does "has solution" mean?

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9849–9855

Can you explain why these new DAG combines are required for what is supposed to be an NFC patch?

[paulwalker-arm]

Taking another look I would expect this patch to also change LegalizeDAG to query TruncMStoreActions?

[DavidTruby]

This needs a large rework that I haven't had time to work on yet+

[DavidTruby]

When we come around to doing this, it should be done sufficiently differently to this anyway, so abandoning for now.

llvm/include/llvm/CodeGen/TargetLowering.h

#include "llvm/IR/ConstrainedOps.def"
   virtual bool canCombineTruncStore(EVT ValVT, EVT MemVT,
                                     bool LegalOnly) const {
     if (LegalOnly)
       return isTruncStoreLegal(ValVT, MemVT);
 
     return isTruncStoreLegalOrCustom(ValVT, MemVT);
   }
 
+  /// Return how this store with truncation should be treated: either it is
+  /// legal, needs to be promoted to a larger size, needs to be expanded to some
+  /// other code sequence, or the target has a custom expander for it.
+  LegalizeAction getTruncMStoreAction(EVT ValVT, EVT MemVT) const {
+    if (ValVT.isExtended() || MemVT.isExtended())
+      return Expand;
+    unsigned ValI = (unsigned)ValVT.getSimpleVT().SimpleTy;
+    unsigned MemI = (unsigned)MemVT.getSimpleVT().SimpleTy;
+    assert(ValI < MVT::VALUETYPE_SIZE && MemI < MVT::VALUETYPE_SIZE &&
+           "Table isn't big enough!");
+    return TruncMStoreActions[ValI][MemI];
+  }
+
+  /// Return true if the specified store with truncation is legal on this
+  /// target.
+  bool isTruncMStoreLegal(EVT ValVT, EVT MemVT) const {
+    return isTypeLegal(ValVT) && getTruncMStoreAction(ValVT, MemVT) == Legal;
+  }
+
+  /// Return true if the specified store with truncation has solution on this
+  /// target.
+  bool isTruncMStoreLegalOrCustom(EVT ValVT, EVT MemVT) const {
+    return isTypeLegal(ValVT) && (getTruncMStoreAction(ValVT, MemVT) == Legal ||
+                                  getTruncMStoreAction(ValVT, MemVT) == Custom);
+  }
+
+  virtual bool canCombineTruncMStore(EVT ValVT, EVT MemVT,
+                                     bool LegalOnly) const {
+    if (LegalOnly)
+      return isTruncMStoreLegal(ValVT, MemVT);
+
+    return isTruncMStoreLegalOrCustom(ValVT, MemVT);
+  }
+
   /// Return how the indexed load should be treated: either it is legal, needs
   /// to be promoted to a larger size, needs to be expanded to some other code
   /// sequence, or the target has a custom expander for it.
   LegalizeAction getIndexedLoadAction(unsigned IdxMode, MVT VT) const {
     return getIndexedModeAction(IdxMode, VT, IMAB_Load);
   }
 
   /// Return true if the specified indexed load is legal on this target.
protected:
 
   /// Indicate that the specified truncating store does not work with the
   /// specified type and indicate what to do about it.
   void setTruncStoreAction(MVT ValVT, MVT MemVT, LegalizeAction Action) {
     assert(ValVT.isValid() && MemVT.isValid() && "Table isn't big enough!");
     TruncStoreActions[(unsigned)ValVT.SimpleTy][MemVT.SimpleTy] = Action;
   }
 
+  /// Indicate that the specified masked truncating store does not work with the
+  /// specified type and indicate what to do about it.
+  void setTruncMStoreAction(MVT ValVT, MVT MemVT, LegalizeAction Action) {
+    assert(ValVT.isValid() && MemVT.isValid() && "Table isn't big enough!");
+    TruncMStoreActions[(unsigned)ValVT.SimpleTy][MemVT.SimpleTy] = Action;
+  }
+
   /// Indicate that the specified indexed load does or does not work with the
   /// specified type and indicate what to do abort it.
   ///
   /// NOTE: All indexed mode loads are initialized to Expand in
   /// TargetLowering.cpp
   void setIndexedLoadAction(unsigned IdxMode, MVT VT, LegalizeAction Action) {
     setIndexedModeAction(IdxMode, VT, IMAB_Load, Action);
   }
private:
   /// specific value type and extension type. Uses 4-bits to store the action
   /// for each of the 4 load ext types.
   uint16_t LoadExtActions[MVT::VALUETYPE_SIZE][MVT::VALUETYPE_SIZE];
 
   /// For each value type pair keep a LegalizeAction that indicates whether a
   /// truncating store of a specific value type and truncating type is legal.
   LegalizeAction TruncStoreActions[MVT::VALUETYPE_SIZE][MVT::VALUETYPE_SIZE];
 
+  /// For each value type pair keep a LegalizeAction that indicates whether a
+  /// masked truncating store of a specific value type and truncating type is
+  /// legal.
+  LegalizeAction TruncMStoreActions[MVT::VALUETYPE_SIZE][MVT::VALUETYPE_SIZE];
+
   /// For each indexed mode and each value type, keep a quad of LegalizeAction
   /// that indicates how instruction selection should deal with the load /
   /// store / maskedload / maskedstore.
   ///
   /// The first dimension is the value_type for the reference. The second
   /// dimension represents the various modes for load store.
   uint16_t IndexedModeActions[MVT::VALUETYPE_SIZE][ISD::LAST_INDEXED_MODE];
 

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitMSCATTER(SDNode *N) {
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitMSTORE(SDNode *N) {
   MaskedStoreSDNode *MST = cast<MaskedStoreSDNode>(N);
   SDValue Mask = MST->getMask();
   SDValue Chain = MST->getChain();
+  SDValue Value = MST->getValue();
+  SDValue Ptr = MST->getBasePtr();
   SDLoc DL(N);
 
   // Zap masked stores with a zero mask.
   if (ISD::isConstantSplatVectorAllZeros(Mask.getNode()))
     return Chain;
 
   // If this is a masked load with an all ones mask, we can use a unmasked load.
   // FIXME: Can we do this for indexed, compressing, or truncating stores?
   if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) &&
       MST->isUnindexed() && !MST->isCompressingStore() &&
       !MST->isTruncatingStore())
     return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(),
                         MST->getBasePtr(), MST->getMemOperand());
 
+  if (MST->isTruncatingStore() && MST->isUnindexed() &&
+      Value.getValueType().isInteger() &&
+      (!isa<ConstantSDNode>(Value) ||
+       !cast<ConstantSDNode>(Value)->isOpaque())) {
+    APInt TruncDemandedBits =
+        APInt::getLowBitsSet(Value.getScalarValueSizeInBits(),
+                             MST->getMemoryVT().getScalarSizeInBits());
+
+    // See if we can simplify the operation with
+    // SimplifyDemandedBits, which only works if the value has a single use.
+    if (SimplifyDemandedBits(Value, TruncDemandedBits)) {
+      // Re-visit the store if anything changed and the store hasn't been merged
+      // with another node (N is deleted) SimplifyDemandedBits will add Value's
+      // node back to the worklist if necessary, but we also need to re-visit
+      // the Store node itself.
+      if (N->getOpcode() != ISD::DELETED_NODE)
+        AddToWorklist(N);
+      return SDValue(N, 0);
+    }
+  }
+
+  // If this is an FP_ROUND or TRUNC followed by a store, fold this into a
+  // truncating store.  We can do this even if this is already a truncstore.
+  if ((Value.getOpcode() == ISD::FP_ROUND ||
+       Value.getOpcode() == ISD::TRUNCATE) &&
+      Value.getNode()->hasOneUse() && MST->isUnindexed() &&
+      TLI.canCombineTruncMStore(Value.getOperand(0).getValueType(),
+                                MST->getMemoryVT(), LegalOperations)) {
+    return DAG.getMaskedStore(Chain, SDLoc(N), Value.getOperand(0), Ptr,
+                              MST->getOffset(), MST->getMask(),
+                              MST->getMemoryVT(), MST->getMemOperand(),
+                              MST->getAddressingMode(), /*IsTruncating=*/true);
+  }
+
   // Try transforming N to an indexed store.
   if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N))
     return SDValue(N, 0);
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitMGATHER(SDNode *N) {

llvm/lib/CodeGen/TargetLoweringBase.cpp

TargetLoweringBase::TargetLoweringBase(const TargetMachine &tm) : TM(tm) {
   InitCmpLibcallCCs(CmpLibcallCCs);
 }
 
 void TargetLoweringBase::initActions() {
   // All operations default to being supported.
   memset(OpActions, 0, sizeof(OpActions));
   memset(LoadExtActions, 0, sizeof(LoadExtActions));
   memset(TruncStoreActions, 0, sizeof(TruncStoreActions));
+  memset(TruncMStoreActions, 0, sizeof(TruncMStoreActions));
   memset(IndexedModeActions, 0, sizeof(IndexedModeActions));
   memset(CondCodeActions, 0, sizeof(CondCodeActions));
   std::fill(std::begin(RegClassForVT), std::end(RegClassForVT), nullptr);
   std::fill(std::begin(TargetDAGCombineArray),
             std::end(TargetDAGCombineArray), 0);
 
   for (MVT VT : MVT::fp_valuetypes()) {
     MVT IntVT = MVT::getIntegerVT(VT.getFixedSizeInBits());

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (MVT VT : MVT::fixedlen_vector_valuetypes()) {
       setOperationAction(ISD::SMUL_LOHI, VT, Expand);
       setOperationAction(ISD::UMUL_LOHI, VT, Expand);
 
       setOperationAction(ISD::BSWAP, VT, Expand);
       setOperationAction(ISD::CTTZ, VT, Expand);
 
       for (MVT InnerVT : MVT::fixedlen_vector_valuetypes()) {
         setTruncStoreAction(VT, InnerVT, Expand);
+        setTruncMStoreAction(VT, InnerVT, Expand);
         setLoadExtAction(ISD::SEXTLOAD, VT, InnerVT, Expand);
         setLoadExtAction(ISD::ZEXTLOAD, VT, InnerVT, Expand);
         setLoadExtAction(ISD::EXTLOAD, VT, InnerVT, Expand);
       }
     }
 
     // AArch64 has implementations of a lot of rounding-like FP operations.
     for (MVT Ty : {MVT::v2f32, MVT::v4f32, MVT::v2f64}) {
for (auto VT : {MVT::v4f16, MVT::v8f16, MVT::v2f32, MVT::v4f32, MVT::v1f64,
       setOperationAction(ISD::MSCATTER, VT, Custom);
     }
 
     for (MVT VT : MVT::fp_scalable_vector_valuetypes()) {
       for (MVT InnerVT : MVT::fp_scalable_vector_valuetypes()) {
         // Avoid marking truncating FP stores as legal to prevent the
         // DAGCombiner from creating unsupported truncating stores.
         setTruncStoreAction(VT, InnerVT, Expand);
+        setTruncMStoreAction(VT, InnerVT, Expand);
         // SVE does not have floating-point extending loads.
         setLoadExtAction(ISD::SEXTLOAD, VT, InnerVT, Expand);
         setLoadExtAction(ISD::ZEXTLOAD, VT, InnerVT, Expand);
         setLoadExtAction(ISD::EXTLOAD, VT, InnerVT, Expand);
       }
     }
 
     // SVE supports truncating stores of 64 and 128-bit vectors

llvm/lib/Target/Hexagon/HexagonISelLowering.cpp

HexagonTargetLowering::HexagonTargetLowering(const TargetMachine &TM,
   // No extending loads from i32.
   for (MVT VT : MVT::integer_valuetypes()) {
     setLoadExtAction(ISD::ZEXTLOAD, VT, MVT::i32, Expand);
     setLoadExtAction(ISD::SEXTLOAD, VT, MVT::i32, Expand);
     setLoadExtAction(ISD::EXTLOAD,  VT, MVT::i32, Expand);
   }
   // Turn FP truncstore into trunc + store.
   setTruncStoreAction(MVT::f64, MVT::f32, Expand);
+  setTruncMStoreAction(MVT::f64, MVT::f32, Expand);
   // Turn FP extload into load/fpextend.
   for (MVT VT : MVT::fp_valuetypes())
     setLoadExtAction(ISD::EXTLOAD, VT, MVT::f32, Expand);
 
   // Expand BR_CC and SELECT_CC for all integer and fp types.
   for (MVT VT : MVT::integer_valuetypes()) {
     setOperationAction(ISD::BR_CC,     VT, Expand);
     setOperationAction(ISD::SELECT_CC, VT, Expand);
for (MVT VT : MVT::fixedlen_vector_valuetypes()) {
     // Expand all extending loads and truncating stores:
     for (MVT TargetVT : MVT::fixedlen_vector_valuetypes()) {
       if (TargetVT == VT)
         continue;
       setLoadExtAction(ISD::EXTLOAD, TargetVT, VT, Expand);
       setLoadExtAction(ISD::ZEXTLOAD, TargetVT, VT, Expand);
       setLoadExtAction(ISD::SEXTLOAD, TargetVT, VT, Expand);
       setTruncStoreAction(VT, TargetVT, Expand);
+      setTruncMStoreAction(VT, TargetVT, Expand);
     }
 
     // Normalize all inputs to SELECT to be vectors of i32.
     if (VT.getVectorElementType() != MVT::i32) {
       MVT VT32 = MVT::getVectorVT(MVT::i32, VT.getSizeInBits()/32);
       setOperationAction(ISD::SELECT, VT, Promote);
       AddPromotedToType(ISD::SELECT, VT, VT32);
     }

llvm/lib/Target/X86/X86ISelLowering.cpp

X86TargetLowering::X86TargetLowering(const X86TargetMachine &TM,
 
   // We don't accept any truncstore of integer registers.
   setTruncStoreAction(MVT::i64, MVT::i32, Expand);
   setTruncStoreAction(MVT::i64, MVT::i16, Expand);
   setTruncStoreAction(MVT::i64, MVT::i8 , Expand);
   setTruncStoreAction(MVT::i32, MVT::i16, Expand);
   setTruncStoreAction(MVT::i32, MVT::i8 , Expand);
   setTruncStoreAction(MVT::i16, MVT::i8,  Expand);
+  setTruncMStoreAction(MVT::i64, MVT::i32, Expand);
+  setTruncMStoreAction(MVT::i64, MVT::i16, Expand);
+  setTruncMStoreAction(MVT::i64, MVT::i8, Expand);
+  setTruncMStoreAction(MVT::i32, MVT::i16, Expand);
+  setTruncMStoreAction(MVT::i32, MVT::i8, Expand);
+  setTruncMStoreAction(MVT::i16, MVT::i8, Expand);
 
   setTruncStoreAction(MVT::f64, MVT::f32, Expand);
 
   // SETOEQ and SETUNE require checking two conditions.
   for (auto VT : {MVT::f32, MVT::f64, MVT::f80}) {
     setCondCodeAction(ISD::SETOEQ, VT, Expand);
     setCondCodeAction(ISD::SETUNE, VT, Expand);
   }
for (MVT VT : MVT::fixedlen_vector_valuetypes()) {
     setOperationAction(ISD::SIGN_EXTEND_INREG, VT,Expand);
     setOperationAction(ISD::TRUNCATE, VT, Expand);
     setOperationAction(ISD::SIGN_EXTEND, VT, Expand);
     setOperationAction(ISD::ZERO_EXTEND, VT, Expand);
     setOperationAction(ISD::ANY_EXTEND, VT, Expand);
     setOperationAction(ISD::SELECT_CC, VT, Expand);
     for (MVT InnerVT : MVT::fixedlen_vector_valuetypes()) {
       setTruncStoreAction(InnerVT, VT, Expand);
+      setTruncMStoreAction(InnerVT, VT, Expand);
 
       setLoadExtAction(ISD::SEXTLOAD, InnerVT, VT, Expand);
       setLoadExtAction(ISD::ZEXTLOAD, InnerVT, VT, Expand);
 
       // N.b. ISD::EXTLOAD legality is basically ignored except for i1-like
       // types, we have to deal with them whether we ask for Expansion or not.
       // Setting Expand causes its own optimisation problems though, so leave
       // them legal.