[llvm] [InstCombine] Fold bitcast into vp.load (PR #192173)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Min-Yih Hsu (mshockwave)

<details>
<summary>Changes</summary>

Similar to normal loads, we should be able to fold bitcast into `vp.load` if (1) mask is all-ones (2) either the new vector type has a larger known minimum length than that of the original vector, or you need to make sure the original EVL can be exact divided by the decreasing factor (of the known minimum length).

This patch adds such folding pattern, though it only support cases where the new vector type has a larger known minimum length.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+36) 
- (added) llvm/test/Transforms/InstCombine/fold-vp-load.ll (+61) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index aa3f0dcc3c534..1cfc0290269d9 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3971,6 +3971,42 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     }
     break;
   }
+  case Intrinsic::vp_load: {
+    // Fold away bit casts of the loaded value by loading the desired type,
+    // if the mask is all-ones.
+    Value *Ptr = II->getArgOperand(0);
+    Value *Mask = II->getArgOperand(1);
+    Value *EVL = II->getArgOperand(2);
+    if (!isa<Constant>(Mask) || !cast<Constant>(Mask)->isAllOnesValue() ||
+        !II->hasOneUse())
+      break;
+
+    auto *Cast = dyn_cast<CastInst>(II->user_back());
+    if (!Cast || !Cast->isNoopCast(II->getDataLayout()) ||
+        !isa<VectorType>(Cast->getDestTy()))
+      break;
+    VectorType *OrigVecTy = cast<VectorType>(II->getType());
+    ElementCount OrigVecCnt = OrigVecTy->getElementCount();
+    VectorType *NewVecTy = cast<VectorType>(Cast->getDestTy());
+    ElementCount NewVecCnt = NewVecTy->getElementCount();
+
+    // Right now we only support cases where
+    // the NewVec is longer, because for cases where
+    // it's shorter, we have to be sure that EVL can be
+    // exactly divided, otherwise it might page fault / yield
+    // incorrect results.
+    if (OrigVecCnt.isScalable() == NewVecCnt.isScalable() &&
+        NewVecCnt.hasKnownScalarFactor(OrigVecCnt)) {
+      unsigned Factor = NewVecCnt.getKnownScalarFactor(OrigVecCnt);
+      Value *NewEVL = Builder.CreateNUWMul(EVL, Builder.getInt32(Factor));
+      Value *NewMask = Builder.CreateVectorSplat(NewVecCnt, Builder.getTrue());
+      Value *NewVP = Builder.CreateIntrinsic(NewVecTy, Intrinsic::vp_load,
+                                             {Ptr, NewMask, NewEVL});
+      replaceInstUsesWith(*Cast, NewVP);
+      return eraseInstFromFunction(*Cast);
+    }
+    break;
+  }
   case Intrinsic::experimental_vp_reverse: {
     Value *X;
     Value *Vec = II->getArgOperand(0);
diff --git a/llvm/test/Transforms/InstCombine/fold-vp-load.ll b/llvm/test/Transforms/InstCombine/fold-vp-load.ll
new file mode 100644
index 0000000000000..891d748a10c93
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/fold-vp-load.ll
@@ -0,0 +1,61 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -p instcombine < %s | FileCheck %s
+
+define <vscale x 8 x i16> @bitcast_vp_load_scalable(ptr %p) {
+; CHECK-LABEL: define <vscale x 8 x i16> @bitcast_vp_load_scalable(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:    [[R:%.*]] = call <vscale x 8 x i16> @llvm.vp.load.nxv8i16.p0(ptr [[P]], <vscale x 8 x i1> splat (i1 true), i32 4)
+; CHECK-NEXT:    ret <vscale x 8 x i16> [[R]]
+;
+  %l = call <vscale x 4 x i32> @llvm.vp.load(ptr %p, <vscale x 4 x i1> splat (i1 true), i32 2)
+  %r = bitcast <vscale x 4 x i32> %l to <vscale x 8 x i16>
+  ret <vscale x 8 x i16> %r
+}
+
+define <vscale x 8 x i16> @bitcast_vp_load_scalable_evl(ptr %p, i32 %evl) {
+; CHECK-LABEL: define <vscale x 8 x i16> @bitcast_vp_load_scalable_evl(
+; CHECK-SAME: ptr [[P:%.*]], i32 [[EVL:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw i32 [[EVL]], 1
+; CHECK-NEXT:    [[R:%.*]] = call <vscale x 8 x i16> @llvm.vp.load.nxv8i16.p0(ptr [[P]], <vscale x 8 x i1> splat (i1 true), i32 [[TMP1]])
+; CHECK-NEXT:    ret <vscale x 8 x i16> [[R]]
+;
+  %l = call <vscale x 4 x i32> @llvm.vp.load(ptr %p, <vscale x 4 x i1> splat (i1 true), i32 %evl)
+  %r = bitcast <vscale x 4 x i32> %l to <vscale x 8 x i16>
+  ret <vscale x 8 x i16> %r
+}
+
+define <8 x i16> @bitcast_vp_load_fixed_evl(ptr %p, i32 %evl) {
+; CHECK-LABEL: define <8 x i16> @bitcast_vp_load_fixed_evl(
+; CHECK-SAME: ptr [[P:%.*]], i32 [[EVL:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw i32 [[EVL]], 1
+; CHECK-NEXT:    [[R:%.*]] = call <8 x i16> @llvm.vp.load.v8i16.p0(ptr [[P]], <8 x i1> splat (i1 true), i32 [[TMP1]])
+; CHECK-NEXT:    ret <8 x i16> [[R]]
+;
+  %l = call <4 x i32> @llvm.vp.load(ptr %p, <4 x i1> splat (i1 true), i32 %evl)
+  %r = bitcast <4 x i32> %l to <8 x i16>
+  ret <8 x i16> %r
+}
+
+define <vscale x 2 x i64> @negative_bitcast_vp_load_scalable(ptr %p) {
+; CHECK-LABEL: define <vscale x 2 x i64> @negative_bitcast_vp_load_scalable(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:    [[L:%.*]] = call <vscale x 4 x i32> @llvm.vp.load.nxv4i32.p0(ptr [[P]], <vscale x 4 x i1> splat (i1 true), i32 3)
+; CHECK-NEXT:    [[R:%.*]] = bitcast <vscale x 4 x i32> [[L]] to <vscale x 2 x i64>
+; CHECK-NEXT:    ret <vscale x 2 x i64> [[R]]
+;
+  %l = call <vscale x 4 x i32> @llvm.vp.load(ptr %p, <vscale x 4 x i1> splat (i1 true), i32 3)
+  %r = bitcast <vscale x 4 x i32> %l to <vscale x 2 x i64>
+  ret <vscale x 2 x i64> %r
+}
+
+define <vscale x 8 x i16> @negative_bitcast_vp_load_scalable_mask(ptr %p, <vscale x 4 x i1> %m) {
+; CHECK-LABEL: define <vscale x 8 x i16> @negative_bitcast_vp_load_scalable_mask(
+; CHECK-SAME: ptr [[P:%.*]], <vscale x 4 x i1> [[M:%.*]]) {
+; CHECK-NEXT:    [[L:%.*]] = call <vscale x 4 x i32> @llvm.vp.load.nxv4i32.p0(ptr [[P]], <vscale x 4 x i1> [[M]], i32 2)
+; CHECK-NEXT:    [[R:%.*]] = bitcast <vscale x 4 x i32> [[L]] to <vscale x 8 x i16>
+; CHECK-NEXT:    ret <vscale x 8 x i16> [[R]]
+;
+  %l = call <vscale x 4 x i32> @llvm.vp.load(ptr %p, <vscale x 4 x i1> %m, i32 2)
+  %r = bitcast <vscale x 4 x i32> %l to <vscale x 8 x i16>
+  ret <vscale x 8 x i16> %r
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/192173