[SDAG] Refine MMO size when converting masked load/store to normal load/store

[dmgreen]

Summary

After D113888 / 32b6c17b29079e7d the MMO size of a masked loads/store is unknown. When we are converting back to a standard load/store because the mask is known all ones, we can refine that to the correct size from the size of the vector being loaded/stored.

[RKSimon]

@dmgreen please can you rebase?

[dmgreen]

Rebase over formatting.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10126–10127

Can we just one of the other signatures of getLoad and let it rebuild the MMO? Like this one

SDValue SelectionDAG::getLoad(EVT VT, const SDLoc &dl, SDValue Chain,            
                              SDValue Ptr, MachinePointerInfo PtrInfo,           
                              MaybeAlign Alignment,                              
                              MachineMemOperand::Flags MMOFlags,                 
                              const AAMDNodes &AAInfo, const MDNode *Ranges)

Or is there a risk that MemVT and VT don't match?

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10126–10127

Yeah sounds good. We check that the masked load is non-extending, so I believe it should be OK.

[dmgreen]

Update to use a different getLoad/setStore signature.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10068

(style) drop the braces?

[dmgreen]

Remove superfluous brackets. Thanks

[RKSimon]

LGTM

[craig.topper]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitMSTORE(SDNode *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());
+                        MST->getBasePtr(), MST->getPointerInfo(),
+                        MST->getOriginalAlign(), MachineMemOperand::MOStore,
+                        MST->getAAInfo());
 
   // Try transforming N to an indexed store.
   if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N))
     return SDValue(N, 0);
 
   return SDValue();
 }
 
SDValue DAGCombiner::visitMLOAD(SDNode *N) {
 
   // Zap masked loads with a zero mask.
   if (ISD::isConstantSplatVectorAllZeros(Mask.getNode()))
     return CombineTo(N, MLD->getPassThru(), MLD->getChain());
 
   // If this is a masked load with an all ones mask, we can use a unmasked load.
   // FIXME: Can we do this for indexed, expanding, or extending loads?
   if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) && MLD->isUnindexed() &&
       !MLD->isExpandingLoad() && MLD->getExtensionType() == ISD::NON_EXTLOAD) {
-    SDValue NewLd = DAG.getLoad(N->getValueType(0), SDLoc(N), MLD->getChain(),
-                                MLD->getBasePtr(), MLD->getMemOperand());
+    SDValue NewLd = DAG.getLoad(
+        N->getValueType(0), SDLoc(N), MLD->getChain(), MLD->getBasePtr(),
+        MLD->getPointerInfo(), MLD->getOriginalAlign(),
+        MachineMemOperand::MOLoad, MLD->getAAInfo(), MLD->getRanges());
     return CombineTo(N, NewLd, NewLd.getValue(1));
   }
 
   // Try transforming N to an indexed load.
   if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N))
     return SDValue(N, 0);
 
   return SDValue();

llvm/test/CodeGen/X86/masked_loadstore_split.ll

+; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-apple-darwin -mattr=avx2 -stop-after=finalize-isel | FileCheck %s
+
+define void @split_masked_store(<8 x double>* %0) {
+  ; CHECK-LABEL: name: split_masked_store
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   liveins: $rdi
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gr64 = COPY $rdi
+  ; CHECK-NEXT:   [[AVX_SET0_:%[0-9]+]]:vr256 = AVX_SET0
+  ; CHECK-NEXT:   [[VMOVAPSYrm:%[0-9]+]]:vr256 = VMOVAPSYrm $rip, 1, $noreg, %const.0, $noreg :: (load (s256) from constant-pool)
+  ; CHECK-NEXT:   VMASKMOVPDYmr [[COPY]], 1, $noreg, 32, $noreg, killed [[VMOVAPSYrm]], [[AVX_SET0_]] :: (store unknown-size into %ir.0 + 32, align 8)
+  ; CHECK-NEXT:   VMOVUPDYmr [[COPY]], 1, $noreg, 0, $noreg, [[AVX_SET0_]] :: (store (s256) into %ir.0, align 8)
+  ; CHECK-NEXT:   RET 0
+entry:
+  call void @llvm.masked.store.v8f64.p0v8f64(<8 x double> zeroinitializer, <8 x double>* %0, i32 8, <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false, i1 false>)
+  ret void
+}
+
+define <8 x double> @split_masked_load(<8 x double>* %0) {
+  ; CHECK-LABEL: name: split_masked_load
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   liveins: $rdi
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gr64 = COPY $rdi
+  ; CHECK-NEXT:   [[VMOVAPSYrm:%[0-9]+]]:vr256 = VMOVAPSYrm $rip, 1, $noreg, %const.0, $noreg :: (load (s256) from constant-pool)
+  ; CHECK-NEXT:   [[VMASKMOVPDYrm:%[0-9]+]]:vr256 = VMASKMOVPDYrm killed [[VMOVAPSYrm]], [[COPY]], 1, $noreg, 32, $noreg :: (load unknown-size from %ir.0 + 32, align 8)
+  ; CHECK-NEXT:   [[VMOVUPDYrm:%[0-9]+]]:vr256 = VMOVUPDYrm [[COPY]], 1, $noreg, 0, $noreg :: (load (s256) from %ir.0, align 8)
+  ; CHECK-NEXT:   $ymm0 = COPY [[VMOVUPDYrm]]
+  ; CHECK-NEXT:   $ymm1 = COPY [[VMASKMOVPDYrm]]
+  ; CHECK-NEXT:   RET 0, $ymm0, $ymm1
+entry:
+  %x = call <8 x double> @llvm.masked.load.v8f64.p0v8f64(<8 x double>* %0, i32 8, <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 false, i1 false>, <8 x double> poison)
+  ret <8 x double> %x
+}
+
+declare <8 x double> @llvm.masked.load.v8f64.p0v8f64(<8 x double>*, i32, <8 x i1>, <8 x double>)
+declare void @llvm.masked.store.v8f64.p0v8f64(<8 x double>, <8 x double>*, i32, <8 x i1>)