[VP] IR expansion pass for VP strided load and store

[loralb]

Summary

Add IR expansion support for experimental.vp.strided.load and experimental.vp.strided.store intrinsics. If the vp.gather/vp.scatter intrinsics are legal operations for the target, then the strided ones are expanded to them; otherwise, they are expanded to the unpredicated masked.load and masked.store intrinsics.

[loralb]

Changelog:

• Change parent revisions to include both D109584 and D120664

[loralb]

Changelog:

• Add fixed vectors tests

llvm/lib/CodeGen/ExpandVectorPredication.cpp

 #define VPINTERNAL_CASE(X) .Case(#X, VPLegalization::X)
 
 static VPTransform parseOverrideOption(const std::string &TextOpt) {
   return StringSwitch<VPTransform>(TextOpt) VPINTERNAL_VPLEGAL_CASES;
 }
 
 #undef VPINTERNAL_VPLEGAL_CASES
 
+// VP strided load/store -> VP gather/scatter transformation
+static cl::opt<bool> StridedOnlyTransformation(
+    "expandvp-strided-only", cl::init(false), cl::Hidden,
+    cl::desc("If true, ignore TargetTransformInfo and always expand "
+             "experimental.vp.strided.load/store intrinsics to "
+             "vp.gather/scatter ones (Used in testing)."));
+
 // Whether any override options are set.
 static bool anyExpandVPOverridesSet() {
-  return !EVLTransformOverride.empty() || !MaskTransformOverride.empty();
+  return !EVLTransformOverride.empty() || !MaskTransformOverride.empty() ||
+         StridedOnlyTransformation;
 }
 
 #define DEBUG_TYPE "expandvp"
 
 STATISTIC(NumFoldedVL, "Number of folded vector length params");
 STATISTIC(NumLoweredVPOps, "Number of folded vector predication operations");
 
 ///// Helpers {
struct CachingVPExpander {
   /// intrinsic.
   Value *expandPredicationInReduction(IRBuilder<> &Builder,
                                       VPReductionIntrinsic &PI);
 
   /// \brief Lower this VP memory operation to a non-VP intrinsic.
   Value *expandPredicationInMemoryIntrinsic(IRBuilder<> &Builder,
                                             VPIntrinsic &VPI);
 
+  /// \brief Lower this VP strided either to a VP gather/scatter, if legal for
+  /// the target, or to an unpredicated masked gather/scatter.
+  Value *expandPredicationInStridedLoadStore(IRBuilder<> &Builder,
+                                             VPIntrinsic &PI);
+
   /// \brief Query TTI and expand the vector predication in \p P accordingly.
   Value *expandPredication(VPIntrinsic &PI);
 
   /// \brief  Determine how and whether the VPIntrinsic \p VPI shall be
   /// expanded. This overrides TTI with the cl::opts listed at the top of this
   /// file.
   VPLegalization getVPLegalizationStrategy(const VPIntrinsic &VPI) const;
   bool UsingTTIOverrides;
Value *CachingVPExpander::expandPredication(VPIntrinsic &VPI) {
   switch (VPI.getIntrinsicID()) {
   default:
     break;
   case Intrinsic::vp_load:
   case Intrinsic::vp_store:
   case Intrinsic::vp_gather:
   case Intrinsic::vp_scatter:
     return expandPredicationInMemoryIntrinsic(Builder, VPI);
+  case Intrinsic::experimental_vp_strided_load:
+  case Intrinsic::experimental_vp_strided_store:
+    return expandPredicationInStridedLoadStore(Builder, VPI);
   }
 
   return &VPI;
 }
 
 //// } CachingVPExpander
 
 struct TransformJob {
CachingVPExpander::getVPLegalizationStrategy(const VPIntrinsic &VPI) const {
   auto VPStrat = TTI.getVPLegalizationStrategy(VPI);
   if (LLVM_LIKELY(!UsingTTIOverrides)) {
     // No overrides - we are in production.
     return VPStrat;
   }
 
   // Overrides set - we are in testing, the following does not need to be
   // efficient.
+
+  // If StridedOnlyTransformation is true, EVLTransformOverride and
+  // MaskTransformOverride are ignored
+  if (StridedOnlyTransformation) {
+    VPStrat.EVLParamStrategy = VPLegalization::Legal;
+    switch (VPI.getIntrinsicID()) {
+    default:
+      VPStrat.OpStrategy = VPLegalization::Legal;
+      break;
+    case Intrinsic::experimental_vp_strided_load:
+    case Intrinsic::experimental_vp_strided_store:
+      VPStrat.OpStrategy = VPLegalization::Convert;
+      break;
+    }
+    return VPStrat;
+  }
+
   VPStrat.EVLParamStrategy = parseOverrideOption(EVLTransformOverride);
   VPStrat.OpStrategy = parseOverrideOption(MaskTransformOverride);
   return VPStrat;
 }
 
+Value *
+CachingVPExpander::expandPredicationInStridedLoadStore(IRBuilder<> &Builder,
+                                                       VPIntrinsic &PI) {
+  auto CreateGatherScatterPointer = [&](VectorType *VT,
+                                        Value *Stride) -> Value * {
+    Value *MemPtr = PI.getMemoryPointerParam();
+    ElementCount EC = VT->getElementCount();
+    // Bitcast the MemPtr to a pointer to i8
+    Value *Cast = Builder.CreateBitCast(MemPtr, Builder.getInt8PtrTy());
+    // Create a vector of pointers %addrs in the form:
+    // %addrs = <%addr, %addr + %stride, %addr + 2 * %stride, ...>
+    Value *Splat = Builder.CreateVectorSplat(EC, Stride);
+    Value *StepVector =
+        Builder.CreateStepVector(VectorType::get(Builder.getInt64Ty(), EC));
+    Value *Indexes =
+        Builder.CreateBinOp(Instruction::BinaryOps::Mul, StepVector, Splat);
+    Value *GEP = Builder.CreateGEP(Builder.getInt8Ty(), Cast, Indexes);
+    // Bitcast the generated array of pointers to i8 to pointers to the original
+    // element type
+    VectorType *DstType =
+        VectorType::get(VT->getElementType()->getPointerTo(), EC);
+    return Builder.CreateBitCast(GEP, DstType);
+  };
+
+  unsigned IntrNo = Intrinsic::not_intrinsic;
+  SmallVector<Type *, 4> Types;
+  SmallVector<Value *, 4> Args;
+  switch (PI.getIntrinsicID()) {
+  default:
+    llvm_unreachable("Not a vp.strided.* intrinsic");
+  case Intrinsic::experimental_vp_strided_load: {
+    IntrNo = Intrinsic::vp_gather;
+    VectorType *VT = cast<VectorType>(PI.getType());
+    Value *VecOfPtrs = CreateGatherScatterPointer(VT, PI.getOperand(1));
+    Types = {VT, VecOfPtrs->getType()};
+    Args = {VecOfPtrs, PI.getMaskParam(), PI.getVectorLengthParam()};
+    break;
+  }
+  case Intrinsic::experimental_vp_strided_store: {
+    IntrNo = Intrinsic::vp_scatter;
+    VectorType *VT = cast<VectorType>(PI.getMemoryDataParam()->getType());
+    Value *VecOfPtrs = CreateGatherScatterPointer(VT, PI.getOperand(2));
+    Types = {VT, VecOfPtrs->getType()};
+    Args = {PI.getMemoryDataParam(), VecOfPtrs, PI.getMaskParam(),
+            PI.getVectorLengthParam()};
+    break;
+  }
+  }
+
+  auto *Intr = Builder.CreateIntrinsic(IntrNo, Types, Args);
+  replaceOperation(*Intr, PI);
+
+  VPIntrinsic *VPIntr = cast<VPIntrinsic>(Intr);
+  auto VPStrat = getVPLegalizationStrategy(*VPIntr);
+  sanitizeStrategy(*Intr, VPStrat);
+  if (!VPStrat.shouldDoNothing()) {
+    Builder.SetInsertPoint(VPIntr);
+    return expandPredicationInMemoryIntrinsic(Builder, *VPIntr);
+  }
+
+  return Intr;
+}
+
 /// \brief Expand llvm.vp.* intrinsics as requested by \p TTI.
 bool CachingVPExpander::expandVectorPredication() {
   SmallVector<TransformJob, 16> Worklist;
 
   // Collect all VPIntrinsics that need expansion and determine their expansion
   // strategy.
   for (auto &I : instructions(F)) {
     auto *VPI = dyn_cast<VPIntrinsic>(&I);
case VPLegalization::Discard:
       break;
     case VPLegalization::Convert:
       if (foldEVLIntoMask(*Job.PI))
         ++NumFoldedVL;
       break;
     }
     Job.Strategy.EVLParamStrategy = VPLegalization::Legal;
 
-    // Replace with a non-predicated operation.
+    // Replace with a non-predicated operation or another legal VP operation.
     switch (Job.Strategy.OpStrategy) {
     case VPLegalization::Legal:
       break;
     case VPLegalization::Discard:
       llvm_unreachable("Invalid strategy for operators.");
     case VPLegalization::Convert:
       expandPredication(*Job.PI);
       ++NumLoweredVPOps;

llvm/test/CodeGen/Generic/expand-vp-strided-accesses.ll

+; RUN: opt --expandvp --expandvp-strided-only -S < %s | FileCheck %s --check-prefixes=CHECK,VP_TO_VP
+; RUN: opt --expandvp -S < %s | FileCheck %s --check-prefixes=CHECK,VP_TO_NOT_VP
+
+; Fixed vectors
+define <4 x float> @strided_load_v4f32(float* %ptr, i64 %stride, <4 x i1> %mask, i32 zeroext %evl) {
+; CHECK-LABEL: @strided_load_v4f32(
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <4 x i32> poison, i32 [[EVL:%.*]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <4 x i32> [[DOTSPLATINSERT]], <4 x i32> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <4 x i32> <i32 0, i32 1, i32 2, i32 3>, [[DOTSPLAT]]
+; CHECK-NEXT:    [[TMP2:%.*]] = and <4 x i1> [[TMP1]], [[MASK:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast float* [[PTR:%.*]] to i8*
+; CHECK-NEXT:    [[DOTSPLATINSERT1:%.*]] = insertelement <4 x i64> poison, i64 [[STRIDE:%.*]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT2:%.*]] = shufflevector <4 x i64> [[DOTSPLATINSERT1]], <4 x i64> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP4:%.*]] = mul <4 x i64> <i64 0, i64 1, i64 2, i64 3>, [[DOTSPLAT2]]
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i8, i8* [[TMP3]], <4 x i64> [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = bitcast <4 x i8*> [[TMP5]] to <4 x float*>
+;
+; VP_TO_VP:        [[TMP7:%.*]] = call <4 x float> @llvm.vp.gather.v4f32.v4p0f32(<4 x float*> [[TMP6]], <4 x i1> [[TMP2]], i32 4)
+; VP_TO_NOT_VP:    [[TMP7:%.*]] = call <4 x float> @llvm.masked.gather.v4f32.v4p0f32(<4 x float*> [[TMP6]], i32 4, <4 x i1> [[TMP2]], <4 x float> undef)
+;
+  %v = call <4 x float> @llvm.experimental.vp.strided.load.v4f32.p0f32.i64(float* %ptr, i64 %stride, <4 x i1> %mask, i32 %evl)
+  ret <4 x float> %v
+}
+
+define void @strided_store_v4f32(<4 x float> %v, float *%ptr, i64 %stride, <4 x i1> %mask, i32 zeroext %evl) {
+; CHECK-LABEL: @strided_store_v4f32(
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <4 x i32> poison, i32 [[EVL:%.*]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <4 x i32> [[DOTSPLATINSERT]], <4 x i32> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <4 x i32> <i32 0, i32 1, i32 2, i32 3>, [[DOTSPLAT]]
+; CHECK-NEXT:    [[TMP2:%.*]] = and <4 x i1> [[TMP1]], [[MASK:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast float* [[PTR:%.*]] to i8*
+; CHECK-NEXT:    [[DOTSPLATINSERT1:%.*]] = insertelement <4 x i64> poison, i64 [[STRIDE:%.*]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT2:%.*]] = shufflevector <4 x i64> [[DOTSPLATINSERT1]], <4 x i64> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP4:%.*]] = mul <4 x i64> <i64 0, i64 1, i64 2, i64 3>, [[DOTSPLAT2]]
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i8, i8* [[TMP3]], <4 x i64> [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = bitcast <4 x i8*> [[TMP5]] to <4 x float*>
+;
+; VP_TO_VP:        call void @llvm.vp.scatter.v4f32.v4p0f32(<4 x float> [[V:%.*]], <4 x float*> [[TMP6]], <4 x i1> [[TMP2]], i32 4)
+; VP_TO_NOT_VP:    call void @llvm.masked.scatter.v4f32.v4p0f32(<4 x float> [[V:%.*]], <4 x float*> [[TMP6]], i32 4, <4 x i1> [[TMP2]])
+;
+  call void @llvm.experimental.vp.strided.store.v4f32.p0f32.i64(<4 x float> %v, float* %ptr, i64 %stride, <4 x i1> %mask, i32 %evl)
+  ret void
+}
+
+; Scalable vectors
+define <vscale x 2 x float> @strided_load_nxv2f32(float* %ptr, i64 %stride, <vscale x 2 x i1> %mask, i32 zeroext %evl) {
+; CHECK-LABEL: @strided_load_nxv2f32(
+; CHECK-NEXT:    [[TMP1:%.*]] = call <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1.i32(i32 0, i32 [[EVL:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = and <vscale x 2 x i1> [[TMP1]], [[MASK:%.*]]
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[SCALABLE_SIZE:%.*]] = mul nuw i32 [[VSCALE]], 2
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast float* [[PTR:%.*]] to i8*
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 2 x i64> poison, i64 [[STRIDE:%.*]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 2 x i64> [[DOTSPLATINSERT]], <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP4:%.*]] = call <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
+; CHECK-NEXT:    [[TMP5:%.*]] = mul <vscale x 2 x i64> [[TMP4]], [[DOTSPLAT]]
+; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr i8, i8* [[TMP3]], <vscale x 2 x i64> [[TMP5]]
+; CHECK-NEXT:    [[TMP7:%.*]] = bitcast <vscale x 2 x i8*> [[TMP6]] to <vscale x 2 x float*>
+;
+; VP_TO_VP:        [[TMP8:%.*]] = call <vscale x 2 x float> @llvm.vp.gather.nxv2f32.nxv2p0f32(<vscale x 2 x float*> [[TMP7]], <vscale x 2 x i1> [[TMP2]], i32 [[SCALABLE_SIZE]])
+; VP_TO_NOT_VP:    [[TMP8:%.*]] = call <vscale x 2 x float> @llvm.masked.gather.nxv2f32.nxv2p0f32(<vscale x 2 x float*> [[TMP7]], i32 4, <vscale x 2 x i1> [[TMP2]], <vscale x 2 x float> undef)
+;
+  %v = call <vscale x 2 x float> @llvm.experimental.vp.strided.load.nxv2f32.p0f32.i64(float* %ptr, i64 %stride, <vscale x 2 x i1> %mask, i32 %evl)
+  ret <vscale x 2 x float> %v
+}
+
+define void @strided_store_nxv2f32(<vscale x 2 x float> %v, float *%ptr, i64 %stride, <vscale x 2 x i1> %mask, i32 zeroext %evl) {
+; CHECK-LABEL: @strided_store_nxv2f32(
+; CHECK-NEXT:    [[TMP1:%.*]] = call <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1.i32(i32 0, i32 [[EVL:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = and <vscale x 2 x i1> [[TMP1]], [[MASK:%.*]]
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[SCALABLE_SIZE:%.*]] = mul nuw i32 [[VSCALE]], 2
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast float* [[PTR:%.*]] to i8*
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 2 x i64> poison, i64 [[STRIDE:%.*]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 2 x i64> [[DOTSPLATINSERT]], <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP4:%.*]] = call <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
+; CHECK-NEXT:    [[TMP5:%.*]] = mul <vscale x 2 x i64> [[TMP4]], [[DOTSPLAT]]
+; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr i8, i8* [[TMP3]], <vscale x 2 x i64> [[TMP5]]
+; CHECK-NEXT:    [[TMP7:%.*]] = bitcast <vscale x 2 x i8*> [[TMP6]] to <vscale x 2 x float*>
+;
+; VP_TO_VP:        call void @llvm.vp.scatter.nxv2f32.nxv2p0f32(<vscale x 2 x float> [[V:%.*]], <vscale x 2 x float*> [[TMP7]], <vscale x 2 x i1> [[TMP2]], i32 [[SCALABLE_SIZE]])
+; VP_TO_NOT_VP:    call void @llvm.masked.scatter.nxv2f32.nxv2p0f32(<vscale x 2 x float> [[V:%.*]], <vscale x 2 x float*> [[TMP7]], i32 4, <vscale x 2 x i1> [[TMP2]])
+;
+  call void @llvm.experimental.vp.strided.store.nxv2f32.p0f32.i64(<vscale x 2 x float> %v, float* %ptr, i64 %stride, <vscale x 2 x i1> %mask, i32 %evl)
+  ret void
+}
+
+declare <4 x float> @llvm.experimental.vp.strided.load.v4f32.p0f32.i64(float*, i64, <4 x i1>, i32)
+declare void @llvm.experimental.vp.strided.store.v4f32.p0f32.i64(<4 x float>, float*, i64, <4 x i1>, i32)
+
+declare <vscale x 2 x float> @llvm.experimental.vp.strided.load.nxv2f32.p0f32.i64(float*, i64, <vscale x 2 x i1>, i32)
+declare void @llvm.experimental.vp.strided.store.nxv2f32.p0f32.i64(<vscale x 2 x float>, float*, i64, <vscale x 2 x i1>, i32)