[RISCV][IA] Support masked.store of deinterleaveN intrinsic

[preames]

This is the masked.store side to the masked.load support added in 881b3fd.

With this change, we support masked.load and masked.store via the intrinsic lowering path used primarily with scalable vectors. An upcoming change will extend the fixed vector (i.a. shuffle vector) paths in the same manner.

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Philip Reames (preames)

Changes

This is the masked.store side to the masked.load support added in 881b3fd.

With this change, we support masked.load and masked.store via the intrinsic lowering path used primarily with scalable vectors. An upcoming change will extend the fixed vector (i.a. shuffle vector) paths in the same manner.

---

Full diff: https://github.com/llvm/llvm-project/pull/149893.diff

3 Files Affected:

• (modified) llvm/lib/CodeGen/InterleavedAccessPass.cpp (+18-8)
• (modified) llvm/lib/Target/RISCV/RISCVInterleavedAccess.cpp (+11)
• (modified) llvm/test/CodeGen/RISCV/rvv/vector-interleave-store.ll (+23)

diff --git a/llvm/lib/CodeGen/InterleavedAccessPass.cpp b/llvm/lib/CodeGen/InterleavedAccessPass.cpp
index c1b84c109af8e..df162fca18d92 100644
--- a/llvm/lib/CodeGen/InterleavedAccessPass.cpp
+++ b/llvm/lib/CodeGen/InterleavedAccessPass.cpp
@@ -661,7 +661,9 @@ bool InterleavedAccessImpl::lowerInterleaveIntrinsic(
   Instruction *StoredBy = dyn_cast<Instruction>(IntII->user_back());
   if (!StoredBy)
     return false;
-  if (!isa<StoreInst, VPIntrinsic>(StoredBy))
+  auto *SI = dyn_cast<StoreInst>(StoredBy);
+  auto *II = dyn_cast<IntrinsicInst>(StoredBy);
+  if (!SI && !II)
     return false;
 
   SmallVector<Value *, 8> InterleaveValues(IntII->args());
@@ -669,20 +671,28 @@ bool InterleavedAccessImpl::lowerInterleaveIntrinsic(
   assert(Factor && "unexpected interleave intrinsic");
 
   Value *Mask = nullptr;
-  if (auto *VPStore = dyn_cast<VPIntrinsic>(StoredBy)) {
-    if (VPStore->getIntrinsicID() != Intrinsic::vp_store)
+  if (II) {
+    // Check mask operand. Handle both all-true/false and interleaved mask.
+    Value *WideMask;
+    switch (II->getIntrinsicID()) {
+    default:
       return false;
-
-    Value *WideMask = VPStore->getOperand(2);
+    case Intrinsic::vp_store:
+      WideMask = II->getOperand(2);
+      break;
+    case Intrinsic::masked_store:
+      WideMask = II->getOperand(3);
+      break;
+    }
     Mask = getMask(WideMask, Factor,
                    cast<VectorType>(InterleaveValues[0]->getType()));
     if (!Mask)
       return false;
 
-    LLVM_DEBUG(dbgs() << "IA: Found a vp.store with interleave intrinsic "
-                      << *IntII << " and factor = " << Factor << "\n");
+    LLVM_DEBUG(dbgs() << "IA: Found a vp.store or masked.store with interleave"
+                      << " intrinsic " << *IntII << " and factor = "
+                      << Factor << "\n");
   } else {
-    auto *SI = cast<StoreInst>(StoredBy);
     if (!SI->isSimple())
       return false;
 
diff --git a/llvm/lib/Target/RISCV/RISCVInterleavedAccess.cpp b/llvm/lib/Target/RISCV/RISCVInterleavedAccess.cpp
index 7d1f78996d118..25817b6d2707f 100644
--- a/llvm/lib/Target/RISCV/RISCVInterleavedAccess.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInterleavedAccess.cpp
@@ -169,6 +169,17 @@ static bool getMemOperands(unsigned Factor, VectorType *VTy, Type *XLenTy,
              : Constant::getAllOnesValue(XLenTy);
     return true;
   }
+  case Intrinsic::masked_store: {
+    Ptr = II->getOperand(1);
+    Alignment = cast<ConstantInt>(II->getArgOperand(2))->getAlignValue();
+
+    assert(Mask && "masked.store needs a mask!");
+
+    VL = isa<FixedVectorType>(VTy)
+             ? Builder.CreateElementCount(XLenTy, VTy->getElementCount())
+             : Constant::getAllOnesValue(XLenTy);
+    return true;
+  }
   }
 }
 
diff --git a/llvm/test/CodeGen/RISCV/rvv/vector-interleave-store.ll b/llvm/test/CodeGen/RISCV/rvv/vector-interleave-store.ll
index af55aaa8fce86..7e7d11e8384dc 100644
--- a/llvm/test/CodeGen/RISCV/rvv/vector-interleave-store.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/vector-interleave-store.ll
@@ -303,3 +303,26 @@ define void @vector_interleave_store_factor8(<vscale x 2 x i32> %a, <vscale x 2
   store <vscale x 16 x i32> %v, ptr %p
   ret void
 }
+
+define void @masked_store_factor3(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b, <vscale x 2 x i32> %c, ptr %p) {
+; CHECK-LABEL: masked_store_factor3:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e32, m1, ta, ma
+; CHECK-NEXT:    vsseg3e32.v v8, (a0)
+; CHECK-NEXT:    ret
+  %v = call <vscale x 6 x i32> @llvm.vector.interleave3(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b, <vscale x 2 x i32> %c)
+  call void @llvm.masked.store(<vscale x 6 x i32> %v, ptr %p, i32 4, <vscale x 6 x i1> splat (i1 true))
+  ret void
+}
+
+define void @masked_store_factor3_masked(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b, <vscale x 2 x i32> %c, ptr %p, <vscale x 2 x i1> %m) {
+; CHECK-LABEL: masked_store_factor3_masked:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli a1, zero, e32, m1, ta, ma
+; CHECK-NEXT:    vsseg3e32.v v8, (a0), v0.t
+; CHECK-NEXT:    ret
+  %interleaved.mask = call <vscale x 6 x i1> @llvm.vector.interleave3(<vscale x 2 x i1> %m, <vscale x 2 x i1> %m, <vscale x 2 x i1> %m)
+  %v = call <vscale x 6 x i32> @llvm.vector.interleave3(<vscale x 2 x i32> %a, <vscale x 2 x i32> %b, <vscale x 2 x i32> %c)
+  call void @llvm.masked.store(<vscale x 6 x i32> %v, ptr %p, i32 4, <vscale x 6 x i1> %interleaved.mask)
+  ret void
+}

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff HEAD~1 HEAD --extensions cpp -- llvm/lib/CodeGen/InterleavedAccessPass.cpp llvm/lib/Target/RISCV/RISCVInterleavedAccess.cpp

View the diff from clang-format here.

diff --git a/llvm/lib/CodeGen/InterleavedAccessPass.cpp b/llvm/lib/CodeGen/InterleavedAccessPass.cpp
index df162fca1..b87295571 100644
--- a/llvm/lib/CodeGen/InterleavedAccessPass.cpp
+++ b/llvm/lib/CodeGen/InterleavedAccessPass.cpp
@@ -690,8 +690,8 @@ bool InterleavedAccessImpl::lowerInterleaveIntrinsic(
       return false;
 
     LLVM_DEBUG(dbgs() << "IA: Found a vp.store or masked.store with interleave"
-                      << " intrinsic " << *IntII << " and factor = "
-                      << Factor << "\n");
+                      << " intrinsic " << *IntII << " and factor = " << Factor
+                      << "\n");
   } else {
     if (!SI->isSimple())
       return false;

[topperc]

LGTM