[AArch64][SME]: Scalarize masked gather/scatter in streaming mode

[hassnaa-arm]

Summary

To generate code compatible to streaming mode:

• Force scalarizing masked gather/scatter.

Testing files:

• expand-masked-gather.ll
• expand-masked-scatter.ll

[paulwalker-arm]

Are code generation tests all that useful for this patch? The change really affects the ScalarizeMaskedMem pass so the code generator never sees a call to gather or scatter intrinsics, which is presumably why they're no existing NEON code generation tests for these intrinsics. Perhaps instead you want to create an AArch64 version of llvm/test/Transforms/ScalarizeMaskedMemIntrin/X86/expand-masked-gather.ll (and expand-masked-scatter.ll)? Given the code change is independent of the vector length I also doubt you need to go so wide with your the number of vector lengths tested. Perhaps a test per element type and vary the vector length across those tests is enough.

There might still be an issue where some other code can create a GATHER ISD node during code generation but none of the current tests test such a scenario and so we can worry about that later. This patch is all about ensuring ScalarizeMaskedMem does what we need when in streaming mode.

[hassnaa-arm]

Add testing files for scalarizing gather/scatter
Remove masked gather/scatter testing files related to code generation.

[david-arm]

LGTM!

[sdesmalen]

llvm/test/Transforms/ScalarizeMaskedMemIntrin/AArch64/expand-masked-gather.ll

1 ↗

(On Diff #478618)

Is it worth moving both tests to a file called streaming-compatible-expand-masked-gather-scatter.ll ?

33–64 ↗

(On Diff #478618)

These tests don't add much value in testing the expansion of the intrinsic with the code you've written, so can just as well be removed.

llvm/test/Transforms/ScalarizeMaskedMemIntrin/AArch64/expand-masked-scatter.ll

4 ↗

(On Diff #478618)

can you use opaque pointers, so <2 x ptr> ?

31–62 ↗

(On Diff #478618)

These tests don't add much value in testing the expansion of the intrinsic with the code you've written, so can just as well be removed.

[hassnaa-arm]

Use opaque PTRs, Remove not needed tests.

[hassnaa-arm]

Merge testing files gather & scatter into one gather-scatter file

[sdesmalen]

llvm/test/Transforms/ScalarizeMaskedMemIntrin/AArch64/expand-masked-gather.ll

1 ↗

(On Diff #478946)

For this patch, I don't think you have to add any files in a precursory patch. There's not much value in that since we're not looking at a subtle differences in assembly.
In that case, there is also no need to 'delete' any files.

[hassnaa-arm]

Drop the precursory patch

[sdesmalen]

Thanks!

[DavidSpickett]

This test is failing on our bot that only builds the ARM target - https://lab.llvm.org/buildbot/#/builders/245/builds/1637

Is that command line option only available when you have AArch64?

[david-arm]

In D138588#3962977, @DavidSpickett wrote:

This test is failing on our bot that only builds the ARM target - https://lab.llvm.org/buildbot/#/builders/245/builds/1637

Is that command line option only available when you have AArch64?

Hi @DavidSpickett, hopefully I've fixed this with commit a043fcf92fa3

[DavidSpickett]

That did the trick, thanks.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

bool isLegalMaskedLoad(Type *DataType, Align Alignment) {
     return isLegalMaskedLoadStore(DataType, Alignment);
   }
 
   bool isLegalMaskedStore(Type *DataType, Align Alignment) {
     return isLegalMaskedLoadStore(DataType, Alignment);
   }
 
   bool isLegalMaskedGatherScatter(Type *DataType) const {
-    if (!ST->hasSVE())
+    if (!ST->hasSVE() || ST->forceStreamingCompatibleSVE())
       return false;
 
     // For fixed vectors, scalarize if not using SVE for them.
     auto *DataTypeFVTy = dyn_cast<FixedVectorType>(DataType);
     if (DataTypeFVTy && (!ST->useSVEForFixedLengthVectors() ||
                          DataTypeFVTy->getNumElements() < 2))
       return false;
 

llvm/test/Transforms/ScalarizeMaskedMemIntrin/AArch64/streaming-compatible-expand-masked-gather-scatter.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -enable-new-pm=0 -S %s -scalarize-masked-mem-intrin -mtriple=aarch64-linux-gnu -mattr=+sve -force-streaming-compatible-sve | FileCheck %s
+
+define <2 x i32> @scalarize_v2i32(<2 x ptr> %p, <2 x i1> %mask, <2 x i32> %passthru) {
+; CHECK-LABEL: @scalarize_v2i32(
+; CHECK-NEXT:    [[SCALAR_MASK:%.*]] = bitcast <2 x i1> [[MASK:%.*]] to i2
+; CHECK-NEXT:    [[TMP1:%.*]] = and i2 [[SCALAR_MASK]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i2 [[TMP1]], 0
+; CHECK-NEXT:    br i1 [[TMP2]], label [[COND_LOAD:%.*]], label [[ELSE:%.*]]
+; CHECK:       cond.load:
+; CHECK-NEXT:    [[PTR0:%.*]] = extractelement <2 x ptr> [[P:%.*]], i64 0
+; CHECK-NEXT:    [[LOAD0:%.*]] = load i32, ptr [[PTR0]], align 8
+; CHECK-NEXT:    [[RES0:%.*]] = insertelement <2 x i32> [[PASSTHRU:%.*]], i32 [[LOAD0]], i64 0
+; CHECK-NEXT:    br label [[ELSE]]
+; CHECK:       else:
+; CHECK-NEXT:    [[RES_PHI_ELSE:%.*]] = phi <2 x i32> [ [[RES0]], [[COND_LOAD]] ], [ [[PASSTHRU]], [[TMP0:%.*]] ]
+; CHECK-NEXT:    [[TMP3:%.*]] = and i2 [[SCALAR_MASK]], -2
+; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i2 [[TMP3]], 0
+; CHECK-NEXT:    br i1 [[TMP4]], label [[COND_LOAD1:%.*]], label [[ELSE2:%.*]]
+; CHECK:       cond.load1:
+; CHECK-NEXT:    [[PTR1:%.*]] = extractelement <2 x ptr> [[P]], i64 1
+; CHECK-NEXT:    [[LOAD1:%.*]] = load i32, ptr [[PTR1]], align 8
+; CHECK-NEXT:    [[RES1:%.*]] = insertelement <2 x i32> [[RES_PHI_ELSE]], i32 [[LOAD1]], i64 1
+; CHECK-NEXT:    br label [[ELSE2]]
+; CHECK:       else2:
+; CHECK-NEXT:    [[RES_PHI_ELSE3:%.*]] = phi <2 x i32> [ [[RES1]], [[COND_LOAD1]] ], [ [[RES_PHI_ELSE]], [[ELSE]] ]
+; CHECK-NEXT:    ret <2 x i32> [[RES_PHI_ELSE3]]
+;
+  %ret = call <2 x i32> @llvm.masked.gather.v2i32.v2p0i32(<2 x ptr> %p, i32 8, <2 x i1> %mask, <2 x i32> %passthru)
+  ret <2 x i32> %ret
+}
+
+define void @scalarize_v2i64(<2 x ptr> %p, <2 x i1> %mask, <2 x i64> %value) {
+; CHECK-LABEL: @scalarize_v2i64(
+; CHECK-NEXT:    [[SCALAR_MASK:%.*]] = bitcast <2 x i1> [[MASK:%.*]] to i2
+; CHECK-NEXT:    [[TMP1:%.*]] = and i2 [[SCALAR_MASK]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i2 [[TMP1]], 0
+; CHECK-NEXT:    br i1 [[TMP2]], label [[COND_STORE:%.*]], label [[ELSE:%.*]]
+; CHECK:       cond.store:
+; CHECK-NEXT:    [[ELT0:%.*]] = extractelement <2 x i64> [[VALUE:%.*]], i64 0
+; CHECK-NEXT:    [[PTR0:%.*]] = extractelement <2 x ptr> [[P:%.*]], i64 0
+; CHECK-NEXT:    store i64 [[ELT0]], ptr [[PTR0]], align 8
+; CHECK-NEXT:    br label [[ELSE]]
+; CHECK:       else:
+; CHECK-NEXT:    [[TMP3:%.*]] = and i2 [[SCALAR_MASK]], -2
+; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i2 [[TMP3]], 0
+; CHECK-NEXT:    br i1 [[TMP4]], label [[COND_STORE1:%.*]], label [[ELSE2:%.*]]
+; CHECK:       cond.store1:
+; CHECK-NEXT:    [[ELT1:%.*]] = extractelement <2 x i64> [[VALUE]], i64 1
+; CHECK-NEXT:    [[PTR1:%.*]] = extractelement <2 x ptr> [[P]], i64 1
+; CHECK-NEXT:    store i64 [[ELT1]], ptr [[PTR1]], align 8
+; CHECK-NEXT:    br label [[ELSE2]]
+; CHECK:       else2:
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.masked.scatter.v2i64.v2p0i64(<2 x i64> %value, <2 x ptr> %p, i32 8, <2 x i1> %mask)
+  ret void
+}
+
+declare <2 x i32> @llvm.masked.gather.v2i32.v2p0i32(<2 x i32*>, i32, <2 x i1>, <2 x i32>)
+declare void @llvm.masked.scatter.v2i64.v2p0i64(<2 x i64>, <2 x i64*>, i32, <2 x i1>)