[AArch64][SVE][InstCombine] Combine contiguous gather/scatter to load/store

[peterwaller-arm]

Summary

Contiguous gather => masked load:

(sve.ld1.gather.index Mask BasePtr (sve.index IndexBase 1))
=> (masked.load (gep BasePtr IndexBase) Align Mask undef)

Contiguous scatter => masked store:

(sve.ld1.scatter.index Value Mask BasePtr (sve.index IndexBase 1))
=> (masked.store Value (gep BasePtr IndexBase) Align Mask)

Tests with <vscale x 2 x double>:

[Gather, Scatter] for each [Positive test (index=1), Negative test (index=2), Alignment propagation].

[david-arm]

Hi @peterwaller-arm, looks like a nice improvement, although I haven't reviewed it yet! Just one thought - do you think there is any value in adding a similar instcombine for normal IR, i.e.

%indices = call <vscale x 4 x i64> @llvm.experimental.stepvector()
%geps = getelementptr %base, %indices
%val = call <vscale x 4 x i32> @llvm.masked.gather(<vscale x 4 x i32*> %geps, ...)

?

[peterwaller-arm]

In D112076#3072951, @david-arm wrote:

do you think there is any value in adding a similar instcombine for normal IR

Yes, I had a quick go at this but it requires a bit more time than I have, so I'll defer that for now.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

874

The SVE gather intrinsics explicitly zero inactive lanes and so the replacement masked load must do likewise.

[peterwaller-arm]

• Update tests per review comments, use ptr param instead of alloca.
• Use zeroinitializer instead of undef for PassThru to preserve gathers' semantics.
• Fix store pattern (doesn't take undef, takes Value), also in commit message.

[peterwaller-arm]

I had my wires crossed for this comment: "Update tests per review comments, use ptr param instead of alloca." which belongs to a different patch, please ignore.

[david-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

880

nit: Is it worth reversing the logic to remove indentation and avoid unnecessary creation of variables above? For example, something like this:

Value *Index = II.getOperand(2);
Value *IndexBase;
if (!match(Index, m_Intrinsic<Intrinsic::aarch64_sve_index>(
    m_Value(IndexBase), m_SpecificInt(1))))
  return None;

Value *Mask = II.getOperand(0);
Value *BasePtr = II.getOperand(1);
...

897

Should we also call II.eraseFromParent(); here too like the scatter case?

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

880

As with many other combines, I've structured it in this way to enable further combines to be added without having to touch my logic.

(I'm also a fan of early return in principle *tips hat*, but I think it might not be best for this case).

897

Nice nitpick, but my belief is that the IC.replaceInstUsesWith-related logic erases it in the load case because loads return a value. When these uses are removed the load becomes trivially dead and deleted. Whereas the stores do not have any uses, and therefore the erasure logic doesn't kick in.

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

880

I'm also making a loose assumption here that those variable assignments will optimize to be free or very cheap in the case where they are not used. Additionally, that those variable assignments are likely to find use in hypothetical future combines. These points are all debatable, I think!

[david-arm]

LGTM! OK, fair enough and thanks for the explanations. :)

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

885–888

With the advice below this can become:

Align Alignment = BasePtr->getPointerAlignment(II.getModule()->getDataLayout());

894–895

Please use Builder.CreateMaskedLoad here.

921–924

As above.

931–932

Please use Builder.CreateMaskedStore here.

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-gatherscatter.ll

11

Choose to ignore if you disagree but I don't like seeing . in test function names and would prefer _ to be used.

[peterwaller-arm]

Address review comments:

• Use Builder.Create* functions
• Use _ for test names

[peterwaller-arm]

This patch triggers a read of uninitialized memory on sanitizer-x86_64-linux-fast. I intend to land a fix within the next hour or two, please revert if it is causing you a problem and I will reland it, otherwise fix will appear soon.

[peterwaller-arm]

Relanded in 7a34145f407e4ec652111f5e9483a36f816a6a3a
Reland "[AArch64][SVE][InstCombine] Combine contiguous gather/scatter to load/store"

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

if (match(II.getArgOperand(0),
       match(II.getArgOperand(1), m_Intrinsic<Intrinsic::aarch64_sve_uzp2>(
                                      m_Specific(A), m_Specific(B))))
     return IC.replaceInstUsesWith(
         II, (II.getIntrinsicID() == Intrinsic::aarch64_sve_zip1 ? A : B));
 
   return None;
 }
 
+static Optional<Instruction *> instCombineLD1GatherIndex(InstCombiner &IC,
+                                                         IntrinsicInst &II) {
+  Value *Mask = II.getOperand(0);
+  Value *BasePtr = II.getOperand(1);
+  Value *Index = II.getOperand(2);
+  Type *Ty = II.getType();
+  Type *BasePtrTy = BasePtr->getType();
+  Value *PassThru = ConstantAggregateZero::get(Ty);
+
+  // Contiguous gather => masked load.
+  // (sve.ld1.gather.index Mask BasePtr (sve.index IndexBase 1))
+  // => (masked.load (gep BasePtr IndexBase) Align Mask zeroinitializer)
+  Value *IndexBase;
+  if (match(Index, m_Intrinsic<Intrinsic::aarch64_sve_index>(
+                       m_Value(IndexBase), m_SpecificInt(1)))) {
+    IRBuilder<> Builder(II.getContext());
+    Builder.SetInsertPoint(&II);
+
+    Align Alignment =
+        BasePtr->getPointerAlignment(II.getModule()->getDataLayout());
+
+    Type *VecPtrTy = PointerType::getUnqual(Ty);
+    Value *Ptr = Builder.CreateGEP(BasePtrTy->getPointerElementType(), BasePtr,
+                                   IndexBase);
+    Ptr = Builder.CreateBitCast(Ptr, VecPtrTy);
+    CallInst *MaskedLoad =
+        Builder.CreateMaskedLoad(Ty, Ptr, Alignment, Mask, PassThru);
+    MaskedLoad->takeName(&II);
+    return IC.replaceInstUsesWith(II, MaskedLoad);
+  }
+
+  return None;
+}
+
+static Optional<Instruction *> instCombineST1ScatterIndex(InstCombiner &IC,
+                                                          IntrinsicInst &II) {
+  Value *Val = II.getOperand(0);
+  Value *Mask = II.getOperand(1);
+  Value *BasePtr = II.getOperand(2);
+  Value *Index = II.getOperand(3);
+  Type *Ty = Val->getType();
+  Type *BasePtrTy = BasePtr->getType();
+
+  // Contiguous scatter => masked store.
+  // (sve.ld1.scatter.index Value Mask BasePtr (sve.index IndexBase 1))
+  // => (masked.store Value (gep BasePtr IndexBase) Align Mask)
+  Value *IndexBase;
+  if (match(Index, m_Intrinsic<Intrinsic::aarch64_sve_index>(
+                       m_Value(IndexBase), m_SpecificInt(1)))) {
+    IRBuilder<> Builder(II.getContext());
+    Builder.SetInsertPoint(&II);
+
+    Align Alignment =
+        BasePtr->getPointerAlignment(II.getModule()->getDataLayout());
+
+    Value *Ptr = Builder.CreateGEP(BasePtrTy->getPointerElementType(), BasePtr,
+                                   IndexBase);
+    Type *VecPtrTy = PointerType::getUnqual(Ty);
+    Ptr = Builder.CreateBitCast(Ptr, VecPtrTy);
+
+    CallInst *MaskedStore =
+        Builder.CreateMaskedStore(Val, Ptr, Alignment, Mask);
+    MaskedStore->takeName(&II);
+    II.eraseFromParent();
+
+    return IC.replaceInstUsesWith(II, MaskedStore);
+  }
+
+  return None;
+}
+
 Optional<Instruction *>
 AArch64TTIImpl::instCombineIntrinsic(InstCombiner &IC,
                                      IntrinsicInst &II) const {
   Intrinsic::ID IID = II.getIntrinsicID();
   switch (IID) {
   default:
     break;
   case Intrinsic::aarch64_sve_convert_from_svbool:
AArch64TTIImpl::instCombineIntrinsic(InstCombiner &IC,
   case Intrinsic::aarch64_sve_sunpkhi:
   case Intrinsic::aarch64_sve_sunpklo:
     return instCombineSVEUnpack(IC, II);
   case Intrinsic::aarch64_sve_tuple_get:
     return instCombineSVETupleGet(IC, II);
   case Intrinsic::aarch64_sve_zip1:
   case Intrinsic::aarch64_sve_zip2:
     return instCombineSVEZip(IC, II);
+  case Intrinsic::aarch64_sve_ld1_gather_index:
+    return instCombineLD1GatherIndex(IC, II);
+  case Intrinsic::aarch64_sve_st1_scatter_index:
+    return instCombineST1ScatterIndex(IC, II);
   }
 
   return None;
 }
 
 bool AArch64TTIImpl::isWideningInstruction(Type *DstTy, unsigned Opcode,
                                            ArrayRef<const Value *> Args) {
 

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-gatherscatter.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -instcombine -dce < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+;;
+;; Gathers.
+;;
+
+define <vscale x 2 x double> @test_ld1_gather_index_nxv2f64_stride1(<vscale x 2 x i1> %pred, double* %x, i64 %base) #0 {
+; CHECK-LABEL: @test_ld1_gather_index_nxv2f64_stride1(
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr double, double* [[X:%.*]], i64 [[BASE:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast double* [[TMP1]] to <vscale x 2 x double>*
+; CHECK-NEXT:    [[LD:%.*]] = call <vscale x 2 x double> @llvm.masked.load.nxv2f64.p0nxv2f64(<vscale x 2 x double>* [[TMP2]], i32 1, <vscale x 2 x i1> [[PRED:%.*]], <vscale x 2 x double> zeroinitializer)
+; CHECK-NEXT:    ret <vscale x 2 x double> [[LD]]
+;
+  %idx = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %base, i64 1)
+  %ld = tail call <vscale x 2 x double> @llvm.aarch64.sve.ld1.gather.index.nxv2f64(<vscale x 2 x i1> %pred, double* %x, <vscale x 2 x i64> %idx)
+  ret <vscale x 2 x double> %ld
+}
+
+define <vscale x 2 x double> @test_ld1_gather_index_nxv2f64_stride2_negtest(<vscale x 2 x i1> %pred, double* %x, i64 %base) #0 {
+; CHECK-LABEL: @test_ld1_gather_index_nxv2f64_stride2_negtest(
+; CHECK-NEXT:    [[IDX:%.*]] = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 [[BASE:%.*]], i64 2)
+; CHECK-NEXT:    [[LD:%.*]] = tail call <vscale x 2 x double> @llvm.aarch64.sve.ld1.gather.index.nxv2f64(<vscale x 2 x i1> [[PRED:%.*]], double* [[X:%.*]], <vscale x 2 x i64> [[IDX]])
+; CHECK-NEXT:    ret <vscale x 2 x double> [[LD]]
+;
+  %idx = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %base, i64 2)
+  %ld = tail call <vscale x 2 x double> @llvm.aarch64.sve.ld1.gather.index.nxv2f64(<vscale x 2 x i1> %pred, double* %x, <vscale x 2 x i64> %idx)
+  ret <vscale x 2 x double> %ld
+}
+
+define <vscale x 2 x double> @test_ld1_gather_index_nxv2f64_stride1_align8(<vscale x 2 x i1> %pred, double* align 8 %x, i64 %base) #0 {
+; CHECK-LABEL: @test_ld1_gather_index_nxv2f64_stride1_align8(
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr double, double* [[X:%.*]], i64 [[BASE:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast double* [[TMP1]] to <vscale x 2 x double>*
+; CHECK-NEXT:    [[LD:%.*]] = call <vscale x 2 x double> @llvm.masked.load.nxv2f64.p0nxv2f64(<vscale x 2 x double>* [[TMP2]], i32 8, <vscale x 2 x i1> [[PRED:%.*]], <vscale x 2 x double> zeroinitializer)
+; CHECK-NEXT:    ret <vscale x 2 x double> [[LD]]
+;
+  %idx = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %base, i64 1)
+  %ld = tail call <vscale x 2 x double> @llvm.aarch64.sve.ld1.gather.index.nxv2f64(<vscale x 2 x i1> %pred, double* %x, <vscale x 2 x i64> %idx)
+  ret <vscale x 2 x double> %ld
+}
+
+;;
+;; Scatters.
+;;
+
+define void @test_st1_scatter_index_nxv2f64_stride1(<vscale x 2 x i1> %pred, double* %x, i64 %base, <vscale x 2 x double> %val) #0 {
+; CHECK-LABEL: @test_st1_scatter_index_nxv2f64_stride1(
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr double, double* [[X:%.*]], i64 [[BASE:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast double* [[TMP1]] to <vscale x 2 x double>*
+; CHECK-NEXT:    call void @llvm.masked.store.nxv2f64.p0nxv2f64(<vscale x 2 x double> [[VAL:%.*]], <vscale x 2 x double>* [[TMP2]], i32 1, <vscale x 2 x i1> [[PRED:%.*]])
+; CHECK-NEXT:    ret void
+;
+  %idx = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %base, i64 1)
+  tail call void @llvm.aarch64.sve.st1.scatter.index.nxv2f64(<vscale x 2 x double> %val, <vscale x 2 x i1> %pred, double* %x, <vscale x 2 x i64> %idx)
+  ret void
+}
+
+define void @test_st1_scatter_index_nxv2f64_stride2_negtest(<vscale x 2 x i1> %pred, double* %x, i64 %base, <vscale x 2 x double> %val) #0 {
+; CHECK-LABEL: @test_st1_scatter_index_nxv2f64_stride2_negtest(
+; CHECK-NEXT:    [[IDX:%.*]] = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 [[BASE:%.*]], i64 2)
+; CHECK-NEXT:    tail call void @llvm.aarch64.sve.st1.scatter.index.nxv2f64(<vscale x 2 x double> [[VAL:%.*]], <vscale x 2 x i1> [[PRED:%.*]], double* [[X:%.*]], <vscale x 2 x i64> [[IDX]])
+; CHECK-NEXT:    ret void
+;
+  %idx = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %base, i64 2)
+  tail call void @llvm.aarch64.sve.st1.scatter.index.nxv2f64(<vscale x 2 x double> %val, <vscale x 2 x i1> %pred, double* %x, <vscale x 2 x i64> %idx)
+  ret void
+}
+
+define void @test_st1_scatter_index_nxv2f64_stride1_align8(<vscale x 2 x i1> %pred, double* align 8 %x, i64 %base, <vscale x 2 x double> %val) #0 {
+; CHECK-LABEL: @test_st1_scatter_index_nxv2f64_stride1_align8(
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr double, double* [[X:%.*]], i64 [[BASE:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast double* [[TMP1]] to <vscale x 2 x double>*
+; CHECK-NEXT:    call void @llvm.masked.store.nxv2f64.p0nxv2f64(<vscale x 2 x double> [[VAL:%.*]], <vscale x 2 x double>* [[TMP2]], i32 8, <vscale x 2 x i1> [[PRED:%.*]])
+; CHECK-NEXT:    ret void
+;
+  %idx = tail call <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64 %base, i64 1)
+  tail call void @llvm.aarch64.sve.st1.scatter.index.nxv2f64(<vscale x 2 x double> %val, <vscale x 2 x i1> %pred, double* %x, <vscale x 2 x i64> %idx)
+  ret void
+}
+
+declare <vscale x 2 x i64> @llvm.aarch64.sve.index.nxv2i64(i64, i64)
+declare <vscale x 2 x double> @llvm.aarch64.sve.ld1.gather.index.nxv2f64(<vscale x 2 x i1>, double*, <vscale x 2 x i64>)
+declare void @llvm.aarch64.sve.st1.scatter.index.nxv2f64(<vscale x 2 x double>, <vscale x 2 x i1>, double*, <vscale x 2 x i64>)
+
+attributes #0 = { "target-features"="+sve" }