[ARM] Sink add/mul(shufflevector(insertelement(...), ...), ...) for MVE instruction selection

[samtebbs]

Summary

This patch sinks add/mul(shufflevector(insertelement(...), ...), ...) into the basic block in which they are used so that they can then be selected together. This is useful for various MVE instructions, such as vmla and others that take R registers.

Loop tests have been added to the vmla test file to make sure vmlas are generated in loops.

[dmgreen]

Looks useful. What happens if there are multiple uses of the splat? Do things get handled correctly then?

llvm/lib/Target/ARM/ARMISelLowering.cpp

14338

Do we need to check that this is an instruction, or will the match handle that for us?

14343

ShuffleOp1 isn't needed?

[samtebbs]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14338

We don't in fact, will remove.

14343

It isn't :)

[dmgreen]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14342

Is it possible to come up with a way to not have to repeat which instructions can be sunk to? Possibly considering things like vsub, which can only be sunk if the operand is the second instruction (although that is not your problem here, it will likely be someones problem soon enough).

llvm/test/Transforms/CodeGenPrepare/ARM/sink-add-mul-shufflevector.ll

130

Can you change this to the first operand of the sub. There is technically an MVE VSUB instruction that takes a grp as the second operand, and I presume eventually we will be looking at making that perform the same trick we have here for add and mul.

[dmgreen]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14350

Can you add a comment saying that we want all uses to be sunk (else we'll just end up with the same value in both gpr and vector regs!)

Also if this loop is moved to after the match check below, we won't need to do the O(n) loop if the O(1) match fails. Also negating the condition and returning early is apparently a good thing in llvm.

14352

I think this should be checking the operand of the use, which may not be the same for the original instruction.

[samtebbs]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14350

Agreed.

14352

The Op value here isn't actually used by the lambda and is only passed as a reference so that it can be set.

[samparker]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14342

Hoist this condition and exit early if we don't have MVE.

14354

just cast is fine.

14357

cast is fine.

[samparker]

LGTM. But please address comments before committing.

llvm/lib/Target/ARM/ARMISelLowering.cpp

14331

But still use llvm style cast... = cast<Instruction>

llvm/test/Transforms/CodeGenPrepare/ARM/sink-add-mul-shufflevector.ll

1

Since we have codegen support for vmla, could you also add an llc line here to test that vmlas are generated in loops.

[dmgreen]

Yeah. This looks good.

We should wait to commit until we have the instructions that make use of it though. I think Oliver is working on a few patterns for VADD, VMUL and VSUB, which would cover the cases here.

Also can you add a couple of tests for sub (on there own), that the first operand is not sunk but the second operand is.

llvm/lib/Target/ARM/ARMISelLowering.cpp

14326

Make sure you run clang-format.

[samtebbs]

llvm/test/Transforms/CodeGenPrepare/ARM/sink-add-mul-shufflevector.ll

1

I've added loop tests to the vmla test file

[dmgreen]

Sure. LGTM

[dmgreen]

Very nice. Oliver went and put together a few patches for mul, add and sub, in https://reviews.llvm.org/D67268 and related. So we now (or soon will) produce the other instruction this can effect. This is good to go I think.

[samtebbs]

In D66295#1660956, @dmgreen wrote:

Very nice. Oliver went and put together a few patches for mul, add and sub, in https://reviews.llvm.org/D67268 and related. So we now (or soon will) produce the other instruction this can effect. This is good to go I think.

Thanks!

llvm/lib/Target/ARM/ARMISelLowering.cpp

static bool areExtractExts(Value *Ext1, Value *Ext2) {
   return true;
 }
 
 /// Check if sinking \p I's operands to I's basic block is profitable, because
 /// the operands can be folded into a target instruction, e.g.
 /// sext/zext can be folded into vsubl.
 bool ARMTargetLowering::shouldSinkOperands(Instruction *I,
                                            SmallVectorImpl<Use *> &Ops) const {
-  if (!Subtarget->hasNEON() || !I->getType()->isVectorTy())
+  if (!I->getType()->isVectorTy())
     return false;
 
+  if (Subtarget->hasNEON()) {
     switch (I->getOpcode()) {
     case Instruction::Sub:
     case Instruction::Add: {
       if (!areExtractExts(I->getOperand(0), I->getOperand(1)))
         return false;
       Ops.push_back(&I->getOperandUse(0));
       Ops.push_back(&I->getOperandUse(1));
       return true;
     }
     default:
       return false;
     }
+  }
+
+  if (!Subtarget->hasMVEIntegerOps())
     return false;
+
+
+  auto IsSinker = [](Instruction *I, int Operand) {
+    switch (I->getOpcode()) {
+    case Instruction::Add:
+    case Instruction::Mul:
+      return true;
+    case Instruction::Sub:
+      return Operand == 1;
+    default:
+      return false;
+    }
+  };
+
+  int Op = 0;
+  if (!isa<ShuffleVectorInst>(I->getOperand(Op)))
+    Op = 1;
+  if (!IsSinker(I, Op))
+    return false;
+  if (!match(I->getOperand(Op),
+            m_ShuffleVector(
+                m_InsertElement(m_Undef(), m_Value(), m_ZeroInt()),
+                m_Undef(), m_Zero()))) {
+    return false;
+  }
+  Instruction *Shuffle = (Instruction*) I->getOperand(Op);
+  // All uses of the shuffle should be sunk to avoid duplicating it across gpr and vector registers
+  for (Use &U : Shuffle->uses()) {
+    Instruction* Insn = (Instruction*) U.getUser();
+    if (!IsSinker(Insn, U.getOperandNo()))
+      return false;
+  }
+  Ops.push_back(&Shuffle->getOperandUse(0));
+  Ops.push_back(&I->getOperandUse(Op));
+  return true;
 }
 
 bool ARMTargetLowering::isVectorLoadExtDesirable(SDValue ExtVal) const {
   EVT VT = ExtVal.getValueType();
 
   if (!isTypeLegal(VT))
     return false;
 
   // Don't create a loadext if we can fold the extension into a wide/long
   // instruction.
   // If there's more than one user instruction, the loadext is desirable no
   // matter what.  There can be two uses by the same instruction.
   if (ExtVal->use_empty() ||
       !ExtVal->use_begin()->isOnlyUserOf(ExtVal.getNode()))
     return true;
 
   SDNode *U = *ExtVal->use_begin();
   if ((U->getOpcode() == ISD::ADD || U->getOpcode() == ISD::SUB ||
        U->getOpcode() == ISD::SHL || U->getOpcode() == ARMISD::VSHLIMM))
     return false;
 
   return true;
 }
 
 bool ARMTargetLowering::allowTruncateForTailCall(Type *Ty1, Type *Ty2) const {

llvm/test/Transforms/CodeGenPrepare/ARM/sink-add-mul-shufflevector.ll

+; RUN: opt -mtriple=thumbv8.1m.main-arm-none-eabi -mattr=+mve.fp < %s -codegenprepare -S | FileCheck -check-prefix=CHECK %s
+
+define void @sink_add_mul(i32* %s1, i32 %x, i32* %d, i32 %n) {
+; CHECK-LABEL: @sink_add_mul(
+; CHECK:    vector.ph:
+; CHECK-NOT:  [[BROADCAST_SPLATINSERT8:%.*]] = insertelement <4 x i32> undef, i32 [[X:%.*]], i32 0
+; CHECK-NOT:  [[BROADCAST_SPLAT9:%.*]] = shufflevector <4 x i32> [[BROADCAST_SPLATINSERT8]], <4 x i32> undef, <4 x i32> zeroinitializer
+; CHECK:    vector.body:
+; CHECK:      [[TMP2:%.*]] = insertelement <4 x i32> undef, i32 [[X:%.*]], i32 0
+; CHECK:      [[TMP3:%.*]] = shufflevector <4 x i32> [[TMP2]], <4 x i32> undef, <4 x i32> zeroinitializer
+;
+entry:
+  %cmp6 = icmp sgt i32 %n, 0
+  br i1 %cmp6, label %vector.ph, label %for.cond.cleanup
+
+vector.ph:                                        ; preds = %for.body.preheader
+  %n.vec = and i32 %n, -4
+  %broadcast.splatinsert8 = insertelement <4 x i32> undef, i32 %x, i32 0
+  %broadcast.splat9 = shufflevector <4 x i32> %broadcast.splatinsert8, <4 x i32> undef, <4 x i32> zeroinitializer
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %0 = getelementptr inbounds i32, i32* %s1, i32 %index
+  %1 = bitcast i32* %0 to <4 x i32>*
+  %wide.load = load <4 x i32>, <4 x i32>* %1, align 4
+  %2 = mul nsw <4 x i32> %wide.load, %broadcast.splat9
+  %3 = getelementptr inbounds i32, i32* %d, i32 %index
+  %4 = bitcast i32* %3 to <4 x i32>*
+  %wide.load10 = load <4 x i32>, <4 x i32>* %4, align 4
+  %5 = add nsw <4 x i32> %wide.load10, %2
+  %6 = bitcast i32* %3 to <4 x i32>*
+  store <4 x i32> %5, <4 x i32>* %6, align 4
+  %index.next = add i32 %index, 4
+  %7 = icmp eq i32 %index.next, %n.vec
+  br i1 %7, label %for.cond.cleanup, label %vector.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %middle.block, %entry
+  ret void
+}
+
+define void @sink_add_mul_multiple(i32* %s1, i32* %s2, i32 %x, i32* %d, i32* %d2, i32 %n) {
+; CHECK-LABEL: @sink_add_mul_multiple(
+; CHECK:    vector.ph:
+; CHECK-NOT:  [[BROADCAST_SPLATINSERT8:%.*]] = insertelement <4 x i32> undef, i32 [[X:%.*]], i32 0
+; CHECK-NOT:  [[BROADCAST_SPLAT9:%.*]] = shufflevector <4 x i32> [[BROADCAST_SPLATINSERT8]], <4 x i32> undef, <4 x i32> zeroinitializer
+; CHECK:    vector.body:
+; CHECK:      [[TMP2:%.*]] = insertelement <4 x i32> undef, i32 [[X:%.*]], i32 0
+; CHECK:      [[TMP3:%.*]] = shufflevector <4 x i32> [[TMP2]], <4 x i32> undef, <4 x i32> zeroinitializer
+; CHECK:      [[TMP11:%.*]] = shufflevector <4 x i32> [[TMP2]], <4 x i32> undef, <4 x i32> zeroinitializer
+;
+entry:
+  %cmp13 = icmp sgt i32 %n, 0
+  br i1 %cmp13, label %vector.ph, label %for.cond.cleanup
+
+vector.ph:                                        ; preds = %for.body.preheader
+  %n.vec = and i32 %n, -4
+  %broadcast.splatinsert15 = insertelement <4 x i32> undef, i32 %x, i32 0
+  %broadcast.splat16 = shufflevector <4 x i32> %broadcast.splatinsert15, <4 x i32> undef, <4 x i32> zeroinitializer
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %0 = getelementptr inbounds i32, i32* %s1, i32 %index
+  %1 = bitcast i32* %0 to <4 x i32>*
+  %wide.load = load <4 x i32>, <4 x i32>* %1, align 4
+  %2 = mul nsw <4 x i32> %wide.load, %broadcast.splat16
+  %3 = getelementptr inbounds i32, i32* %d, i32 %index
+  %4 = bitcast i32* %3 to <4 x i32>*
+  %wide.load17 = load <4 x i32>, <4 x i32>* %4, align 4
+  %5 = add nsw <4 x i32> %wide.load17, %2
+  %6 = bitcast i32* %3 to <4 x i32>*
+  store <4 x i32> %5, <4 x i32>* %6, align 4
+  %7 = getelementptr inbounds i32, i32* %s2, i32 %index
+  %8 = bitcast i32* %7 to <4 x i32>*
+  %wide.load18 = load <4 x i32>, <4 x i32>* %8, align 4
+  %9 = mul nsw <4 x i32> %wide.load18, %broadcast.splat16
+  %10 = getelementptr inbounds i32, i32* %d2, i32 %index
+  %11 = bitcast i32* %10 to <4 x i32>*
+  %wide.load19 = load <4 x i32>, <4 x i32>* %11, align 4
+  %12 = add nsw <4 x i32> %wide.load19, %9
+  %13 = bitcast i32* %10 to <4 x i32>*
+  store <4 x i32> %12, <4 x i32>* %13, align 4
+  %index.next = add i32 %index, 4
+  %14 = icmp eq i32 %index.next, %n.vec
+  br i1 %14, label %for.cond.cleanup, label %vector.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %middle.block, %entry
+  ret void
+}
+
+
+define void @sink_add_sub_unsinkable(i32* %s1, i32* %s2, i32 %x, i32* %d, i32* %d2, i32 %n) {
+; CHECK-LABEL: @sink_add_sub_unsinkable(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP13:%.*]] = icmp sgt i32 [[N:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP13]], label [[VECTOR_PH:%.*]], label [[FOR_COND_CLEANUP:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    [[N_VEC:%.*]] = and i32 [[N]], -4
+; CHECK-NEXT:    [[BROADCAST_SPLATINSERT15:%.*]] = insertelement <4 x i32> undef, i32 [[X:%.*]], i32 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT16:%.*]] = shufflevector <4 x i32> [[BROADCAST_SPLATINSERT15]], <4 x i32> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
+;
+entry:
+  %cmp13 = icmp sgt i32 %n, 0
+  br i1 %cmp13, label %vector.ph, label %for.cond.cleanup
+
+vector.ph:                                        ; preds = %for.body.preheader
+  %n.vec = and i32 %n, -4
+  %broadcast.splatinsert15 = insertelement <4 x i32> undef, i32 %x, i32 0
+  %broadcast.splat16 = shufflevector <4 x i32> %broadcast.splatinsert15, <4 x i32> undef, <4 x i32> zeroinitializer
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i32 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %0 = getelementptr inbounds i32, i32* %s1, i32 %index
+  %1 = bitcast i32* %0 to <4 x i32>*
+  %wide.load = load <4 x i32>, <4 x i32>* %1, align 4
+  %2 = mul nsw <4 x i32> %wide.load, %broadcast.splat16
+  %3 = getelementptr inbounds i32, i32* %d, i32 %index
+  %4 = bitcast i32* %3 to <4 x i32>*
+  %wide.load17 = load <4 x i32>, <4 x i32>* %4, align 4
+  %5 = add nsw <4 x i32> %wide.load17, %2
+  %6 = bitcast i32* %3 to <4 x i32>*
+  store <4 x i32> %5, <4 x i32>* %6, align 4
+  %7 = getelementptr inbounds i32, i32* %s2, i32 %index
+  %8 = bitcast i32* %7 to <4 x i32>*
+  %wide.load18 = load <4 x i32>, <4 x i32>* %8, align 4
+  %9 = sub nsw <4 x i32> %broadcast.splat16, %wide.load18
+  %10 = getelementptr inbounds i32, i32* %d2, i32 %index
+  %11 = bitcast i32* %10 to <4 x i32>*
+  %wide.load19 = load <4 x i32>, <4 x i32>* %11, align 4
+  %12 = add nsw <4 x i32> %wide.load19, %9
+  %13 = bitcast i32* %10 to <4 x i32>*
+  store <4 x i32> %12, <4 x i32>* %13, align 4
+  %index.next = add i32 %index, 4
+  %14 = icmp eq i32 %index.next, %n.vec
+  br i1 %14, label %for.cond.cleanup, label %vector.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %middle.block, %entry
+  ret void
+}