[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

14263

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

14268

ShuffleOp1 isn't needed?

[samtebbs]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14263

We don't in fact, will remove.

14268

It isn't :)

[dmgreen]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14267

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

14275

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.

14277

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

14275

Agreed.

14277

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

14267

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

14279

just cast is fine.

14282

cast is fine.

[samparker]

LGTM. But please address comments before committing.

llvm/lib/Target/ARM/ARMISelLowering.cpp

14275

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

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

2

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

14276

Make sure you run clang-format.

[samtebbs]

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

2

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()) {
+    switch (I->getOpcode()) {
+    case Instruction::Add:
+    case Instruction::Mul: {
+      int Op = 0;
+      if (!isa<ShuffleVectorInst>(I->getOperand(Op)))
+        Op = 1;
+      if (match(I->getOperand(Op),
+                m_ShuffleVector(
+                    m_InsertElement(m_Undef(), m_Value(), m_ZeroInt()),
+                    m_Undef(), m_Zero()))) {
+        Instruction *Shuffle = dyn_cast<Instruction>(I->getOperand(Op));
+        Ops.push_back(&Shuffle->getOperandUse(0));
+        Ops.push_back(&I->getOperandUse(Op));
+        return true;
+      }
+    }
+    }
+  }
+
   return false;
 }
 
 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;

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

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; 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
+}