[AArch64][SVE] Combine bitcasts to predicate types with vector inserts of loads

[bsmith]

Summary

An insert subvector that is inserting the result of a vector predicate
sized load into undef at index 0, whose result is casted to a predicate
type, can be combined into a direct predicate load.

The purpose of this optimization is to clean up cases that will be
introduced in a later patch where casts to/from predicate types from i8
types will use insert subvector, rather than going through memory early.

This optimization is done in SVEIntrinsicOpts rather than InstCombine to
re-introduce scalable loads as late as possible, to give other
optimizations the best chance possible to do a good job.

[efriedma]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

294

Do you need to check getMaxSVEVectorSizeInBits() somewhere?

This code could probably use some comments explaining why it's checking various conditions.

[junparser]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

361

shall we need to handle vector_extract as well?

[bsmith]

• Add equivalent optimization for vector extract
• Add comment to clarify what is being checked
• Ensure max vector length is also exactly the size we want
• Check the attribute rather than the subtarget options

[junparser]

The patch looks good to me. However, It seems that we can also handle SVE vector extract/insert here as well rather than just predicate vector.

[junparser]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

297

Why do we disable differing vscale min/max? I donot see any difference between vscale_range(4,0) and vscale_range(4,6).

[peterwaller-arm]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

365

Nit: s/loads/stores/

[bsmith]

In D106549#2906237, @junparser wrote:

The patch looks good to me. However, It seems that we can also handle SVE vector extract/insert here as well rather than just predicate vector.

In principal we could, however we haven't actually seen a point where this would be useful. Clang doesn't generate the same kinds of bitcast for other types as it does for predicates. It's something we can add later if it does turn out to be useful.

[bsmith]

• Don't allow a maximum vector size of 0
• Fixed clang format issues
• Fix copy-paste issue in comment

[paulwalker-arm]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

286

Is this name representative? I think optimizePredicateStore is more in keeping with what is going on.

295

These are really MinVScale/MaxVScale.

297

I think this needs || MinSVEVectorSize == 0 to cover the vscale_range(0,0) case. Perhaps worth adding a comment along the lines of "The transform needs to know the exact runtime length of scalable vectors".

To bring this point home perhaps it's worth introducing unsigned SVEVectorSize = MinVScale*128;, but I'll leave you to decide if it's worth it.

315–316

I've got a feeling this needs to be more specific otherwise we'll introduce endianness issues. Specially we know predicate load/store instructions are byte based and so we should only allow the transform when the fixed length load/store instructions are also byte based.

Given this is all very specific I imagine you'll end up with something like if (Store->getOperand(0)->getType() == PrecalculatedTy)

323

Is this restriction strictly necessary? I'm thinking we might have multiple stores of the same value so unless there is a real affect on code quality I'd rather not artificially restrict the transform.

327–328

cast<ConstantInt>(IntrI->getOperand(1))->isZero()?

359

I don't see any test to ensure the old bitcast is dead? When combined with my previous comment you'll likely want:

if (IntrI->getNumUses() == 0)
  IntrI->eraseFromParent();
if (BitCast->getNumUses() == 0)
  BitCast->eraseFromParent();

369

I haven't looked but I imagine this'll have similar issues as raise for optimizePredicateVectorExtract.

[bsmith]

• Avoid iterating over the instructions twice
• Rename function to better reflect what they are doing
• Disallow Vscale == 0 case
• Only allow optimization when the fixed type involved is of i8 type
• Relax only one use requirements
• Only remove instructions when there are no more uses

[paulwalker-arm]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

287

Not that bothered but you could pass this into the function given optimizeInstructions already knows it. If you do keep the lookup then can you use I->getFunction().

305

To keep things simple (pun intended) I suggest adding || !Store->isSimple()) so that we only transform ordinary stores.

309

Is the dyn_cast necessary?

353

As with optimizePredicateStore.

369–370

Although possible, is this worth caring about?

390

As with the store case, we should only allow ordinary loads.

394

Is the dyn_cast necessary?

398

The insertion point needs to be where the original load is to ensure the new load maintains the same load/store ordering.

llvm/test/CodeGen/AArch64/sve-insert-vector-to-predicate-load.ll

11

Given the new load placement issue I mentioned above it's worth having the load in a separate basic block in order to validate that fix.

[bsmith]

• Only perform optimization on simple loads/stores
• Remove unnecessary dyn_casts
• Move insertion point for load case to the load
• Add testcase for load insertion point

[paulwalker-arm]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

339

Not sure why I've only just spotted this but does this do anything? Can there be any uses of a store?

llvm/test/CodeGen/AArch64/sve-insert-vector-to-predicate-load.ll

62–65

Not really relevant for this patch but it occurs to me that this IR is provably bogus if we wanted to add the necessary hooks into the verifier.

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

 //     ; (%1 can be replaced with a reinterpret of %2)
 //
 // - optimizes ptest intrinsics where the operands are being needlessly
 //   converted to and from svbool_t.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AArch64.h"
+#include "AArch64Subtarget.h"
 #include "Utils/AArch64BaseInfo.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsAArch64.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PatternMatch.h"
struct SVEIntrinsicOpts : public ModulePass {
 
   bool runOnModule(Module &M) override;
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
 private:
   bool coalescePTrueIntrinsicCalls(BasicBlock &BB,
                                    SmallSetVector<IntrinsicInst *, 4> &PTrues);
   bool optimizePTrueIntrinsicCalls(SmallSetVector<Function *, 4> &Functions);
+  bool optimizePredicateVectorInsert(SmallSetVector<Function *, 4> &Functions);
 
   /// Operates at the function-scope. I.e., optimizations are applied local to
   /// the functions themselves.
   bool optimizeFunctions(SmallSetVector<Function *, 4> &Functions);
+
+  const TargetMachine *TM;
 };
 } // end anonymous namespace
 
 void SVEIntrinsicOpts::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<DominatorTreeWrapperPass>();
   AU.setPreservesCFG();
 }
 
for (auto &BB : *F) {
       Changed |= coalescePTrueIntrinsicCalls(BB, SVAllPTrues);
       Changed |= coalescePTrueIntrinsicCalls(BB, SVPow2PTrues);
     }
   }
 
   return Changed;
 }
 
+// This is done in SVEIntrinsicOpts rather than InstCombine so that we introduce
+// scalable loads as late as possible
+bool SVEIntrinsicOpts::optimizePredicateVectorInsert(
+    SmallSetVector<Function *, 4> &Functions) {
+  bool Changed = false;
+
+  for (auto *F : Functions) {
+    auto &ST = TM->getSubtarget<AArch64Subtarget>(*F);
+    if (!ST.hasSVE())
+      continue;
+    unsigned MinSVEVectorSize = ST.getMinSVEVectorSizeInBits() / 128;
+    if (MinSVEVectorSize == 0)
+      continue;
+
+    auto *PredType =
+        ScalableVectorType::get(Type::getInt1Ty(F->getContext()), 16);
+
+    for (auto &BB : *F) {
+      SmallSetVector<BitCastInst *, 4> InterestingBitCasts;
+
+      for (Instruction &I : BB) {
+        if (I.use_empty())
+          continue;
+
+        auto *BitCast = dyn_cast<BitCastInst>(&I);
+        if (!BitCast || BitCast->getType() != PredType)
+          continue;
+
+        auto *IntrI = dyn_cast<IntrinsicInst>(BitCast->getOperand(0));
+        if (!IntrI ||
+            IntrI->getIntrinsicID() != Intrinsic::experimental_vector_insert ||
+            !IntrI->hasOneUse())
+          continue;
+
+        auto Idx = cast<ConstantInt>(IntrI->getOperand(2))->getZExtValue();
+        if (!isa<UndefValue>(IntrI->getOperand(0)) || Idx != 0)
+          continue;
+
+        auto *Load = dyn_cast<LoadInst>(IntrI->getOperand(1));
+        if (!Load || !Load->hasOneUse())
+          continue;
+
+        auto *FixedLoadType = dyn_cast<FixedVectorType>(Load->getType());
+        if (!FixedLoadType ||
+            FixedLoadType->getPrimitiveSizeInBits() != 16 * MinSVEVectorSize)
+          continue;
+
+        InterestingBitCasts.insert(BitCast);
+      }
+
+      for (auto *BitCast : InterestingBitCasts) {
+        IRBuilder<> Builder(F->getContext());
+        Builder.SetInsertPoint(&BB, BitCast->getIterator());
+
+        auto *IntrI = cast<IntrinsicInst>(BitCast->getOperand(0));
+        auto *Load = cast<LoadInst>(IntrI->getOperand(1));
+        auto *PtrBitCast = Builder.CreateBitCast(
+            Load->getPointerOperand(),
+            PredType->getPointerTo(Load->getPointerAddressSpace()));
+        auto *LoadPred = Builder.CreateLoad(PredType, PtrBitCast);
+
+        BitCast->replaceAllUsesWith(LoadPred);
+        BitCast->eraseFromParent();
+        IntrI->eraseFromParent();
+        Load->eraseFromParent();
+      }
+    }
+  }
+
+  return Changed;
+}
+
 bool SVEIntrinsicOpts::optimizeFunctions(
     SmallSetVector<Function *, 4> &Functions) {
   bool Changed = false;
 
   Changed |= optimizePTrueIntrinsicCalls(Functions);
+  Changed |= optimizePredicateVectorInsert(Functions);
 
   return Changed;
 }
 
 bool SVEIntrinsicOpts::runOnModule(Module &M) {
   bool Changed = false;
   SmallSetVector<Function *, 4> Functions;
 
+  auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
+  assert(TPC && "Expected a TargetPassConfig");
+  TM = &TPC->getTM<TargetMachine>();
+
   // Check for SVE intrinsic declarations first so that we only iterate over
   // relevant functions. Where an appropriate declaration is found, store the
   // function(s) where it is used so we can target these only.
   for (auto &F : M.getFunctionList()) {
     if (!F.isDeclaration())
       continue;
 
     switch (F.getIntrinsicID()) {
+    case Intrinsic::experimental_vector_insert:
     case Intrinsic::aarch64_sve_ptrue:
       for (User *U : F.users())
         Functions.insert(cast<Instruction>(U)->getFunction());
       break;
     default:
       break;
     }
   }
 
   if (!Functions.empty())
     Changed |= optimizeFunctions(Functions);
 
   return Changed;
 }

llvm/test/CodeGen/AArch64/sve-insert-vector-to-predicate-load.ll

+; RUN: opt -S -aarch64-sve-intrinsic-opts < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+define <vscale x 16 x i1> @pred_load_v2i8(<2 x i8>* %addr) #0 {
+; CHECK-LABEL: @pred_load_v2i8(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <2 x i8>* %addr to <vscale x 16 x i1>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <vscale x 16 x i1>, <vscale x 16 x i1>* [[TMP1]]
+; CHECK-NEXT:    ret <vscale x 16 x i1> [[TMP2]]
+  %load = load <2 x i8>, <2 x i8>* %addr, align 4
+  %insert = tail call <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v2i8(<vscale x 2 x i8> undef, <2 x i8> %load, i64 0)
+  %ret = bitcast <vscale x 2 x i8> %insert to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %ret
+}
+
+
+define <vscale x 16 x i1> @pred_load_v4i8(<4 x i8>* %addr) #1 {
+; CHECK-LABEL: @pred_load_v4i8(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <4 x i8>* %addr to <vscale x 16 x i1>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <vscale x 16 x i1>, <vscale x 16 x i1>* [[TMP1]]
+; CHECK-NEXT:    ret <vscale x 16 x i1> [[TMP2]]
+  %load = load <4 x i8>, <4 x i8>* %addr, align 4
+  %insert = tail call <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v4i8(<vscale x 2 x i8> undef, <4 x i8> %load, i64 0)
+  %ret = bitcast <vscale x 2 x i8> %insert to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %ret
+}
+
+define <vscale x 16 x i1> @pred_load_v8i8(<8 x i8>* %addr) #2 {
+; CHECK-LABEL: @pred_load_v8i8(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <8 x i8>* %addr to <vscale x 16 x i1>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <vscale x 16 x i1>, <vscale x 16 x i1>* [[TMP1]]
+; CHECK-NEXT:    ret <vscale x 16 x i1> [[TMP2]]
+  %load = load <8 x i8>, <8 x i8>* %addr, align 4
+  %insert = tail call <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v8i8(<vscale x 2 x i8> undef, <8 x i8> %load, i64 0)
+  %ret = bitcast <vscale x 2 x i8> %insert to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %ret
+}
+
+; Check that too small of a vscale prevents optimization
+define <vscale x 16 x i1> @pred_load_neg1(<4 x i8>* %addr) #0 {
+; CHECK-LABEL: @pred_load_neg1(
+; CHECK:         llvm.experimental.vector.insert
+  %load = load <4 x i8>, <4 x i8>* %addr, align 4
+  %insert = tail call <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v4i8(<vscale x 2 x i8> undef, <4 x i8> %load, i64 0)
+  %ret = bitcast <vscale x 2 x i8> %insert to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %ret
+}
+
+; Check that too large of a vscale prevents optimization
+define <vscale x 16 x i1> @pred_load_neg2(<4 x i8>* %addr) #2 {
+; CHECK-LABEL: @pred_load_neg2(
+; CHECK:         llvm.experimental.vector.insert
+  %load = load <4 x i8>, <4 x i8>* %addr, align 4
+  %insert = tail call <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v4i8(<vscale x 2 x i8> undef, <4 x i8> %load, i64 0)
+  %ret = bitcast <vscale x 2 x i8> %insert to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %ret
+}
+
+; Check that a non-zero index prevents optimization
+define <vscale x 16 x i1> @pred_load_neg3(<4 x i8>* %addr) #1 {
+; CHECK-LABEL: @pred_load_neg3(
+; CHECK:         llvm.experimental.vector.insert
+  %load = load <4 x i8>, <4 x i8>* %addr, align 4
+  %insert = tail call <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v4i8(<vscale x 2 x i8> undef, <4 x i8> %load, i64 4)
+  %ret = bitcast <vscale x 2 x i8> %insert to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %ret
+}
+
+; Check that insertion into a non-undef vector prevents optimization
+define <vscale x 16 x i1> @pred_load_neg4(<4 x i8>* %addr, <vscale x 2 x i8> %passthru) #1 {
+; CHECK-LABEL: @pred_load_neg4(
+; CHECK:         llvm.experimental.vector.insert
+  %load = load <4 x i8>, <4 x i8>* %addr, align 4
+  %insert = tail call <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v4i8(<vscale x 2 x i8> %passthru, <4 x i8> %load, i64 0)
+  %ret = bitcast <vscale x 2 x i8> %insert to <vscale x 16 x i1>
+  ret <vscale x 16 x i1> %ret
+}
+
+declare <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v2i8(<vscale x 2 x i8>, <2 x i8>, i64)
+declare <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v4i8(<vscale x 2 x i8>, <4 x i8>, i64)
+declare <vscale x 2 x i8> @llvm.experimental.vector.insert.nxv2i8.v8i8(<vscale x 2 x i8>, <8 x i8>, i64)
+
+attributes #0 = { "target-features"="+sve" vscale_range(1,1) }
+attributes #1 = { "target-features"="+sve" vscale_range(2,2) }
+attributes #2 = { "target-features"="+sve" vscale_range(4,8) }