[AArch64] Allow loads with imp-def to be handled in getMemOpBaseRegImmOfsWidth()

[junbuml]

Summary

This change will allow loads with imp-def to be clustered in machine-scheduler pass.
areMemAccessesTriviallyDisjoint() can also handle loads with imp-def.

[t.p.northover]

Hi Jun,

I think you should probably be checking MI->getNumExplicitOperands() instead. Other than that it looks like a reasonable change to me.

Tim.

[junbuml]

Use getNumExplicitOperands() as Tim suggested.
Thanks Tim for the review.

[t.p.northover]

Looks good to me. Thanks for updating!

Tim.

lib/Target/AArch64/AArch64InstrInfo.cpp

bool AArch64InstrInfo::getMemOpBaseRegImmOfs(
   };
 }
 
 bool AArch64InstrInfo::getMemOpBaseRegImmOfsWidth(
     MachineInstr *LdSt, unsigned &BaseReg, int64_t &Offset, unsigned &Width,
     const TargetRegisterInfo *TRI) const {
   assert(LdSt->mayLoadOrStore() && "Expected a memory operation.");
   // Handle only loads/stores with base register followed by immediate offset.
-  if (LdSt->getNumOperands() != 3)
+  if (LdSt->getNumOperands() != 3 && LdSt->getNumOperands() != 4)
     return false;
   if (!LdSt->getOperand(1).isReg() || !LdSt->getOperand(2).isImm())
     return false;
 
   // Offset is calculated as the immediate operand multiplied by the scaling factor.
   // Unscaled instructions have scaling factor set to 1.
   unsigned Scale = 0;
   switch (LdSt->getOpcode()) {

test/CodeGen/AArch64/arm64-abi-varargs.ll

; CHECK: ldr {{q[0-9]+}}, [x[[ADDR]]]
   %1 = va_arg i8** %args, <4 x i32>
   store <4 x i32> %1, <4 x i32>* %vv, align 16
   ret void
 }
 
 define void @bar(i32 %x, <4 x i32> %y) nounwind {
 entry:
 ; CHECK-LABEL: bar:
-; CHECK: str {{q[0-9]+}}, [sp, #16]
+; CHECK: stp {{q[0-9]+}}, {{q[0-9]+}}, [sp, #16]
 ; CHECK: str {{x[0-9]+}}, [sp]
   %x.addr = alloca i32, align 4
   %y.addr = alloca <4 x i32>, align 16
   store i32 %x, i32* %x.addr, align 4
   store <4 x i32> %y, <4 x i32>* %y.addr, align 16
   %0 = load i32, i32* %x.addr, align 4
   %1 = load <4 x i32>, <4 x i32>* %y.addr, align 16
   call void (i8*, ...) @foo(i8* getelementptr inbounds ([4 x i8], [4 x i8]* @.str, i32 0, i32 0), i32 %0, <4 x i32> %1)

test/CodeGen/MIR/AArch64/machine-scheduler.mir

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass machine-scheduler  -verify-machineinstrs  -o /dev/null %s 2>&1 | FileCheck %s
+
+--- |
+  define i64 @load_imp-def(i64* nocapture %P, i32 %v) {
+  entry:
+    %0 = bitcast i64* %P to i32*
+    %1 = load i32, i32* %0
+    %conv = zext i32 %1 to i64
+    %arrayidx19 = getelementptr inbounds i64, i64* %P, i64 1
+    %arrayidx1 = bitcast i64* %arrayidx19 to i32*
+    store i32 %v, i32* %arrayidx1
+    %2 = load i64, i64* %arrayidx19
+    %and = and i64 %2, 4294967295
+    %add = add nuw nsw i64 %and, %conv
+    ret i64 %add
+  }
+...
+---
+# CHECK-LABEL: name: load_imp-def
+# CHECK: bb.0.entry:
+# CHECK: LDRWui %x0, 0
+# CHECK: LDRWui %x0, 1
+# CHECK: STRWui %w1, %x0, 2
+name:            load_imp-def
+isSSA:           true
+body:             |
+  bb.0.entry:
+    liveins: %w1, %x0
+    %w8 = LDRWui %x0, 1, implicit-def %x8  :: (load 4 from %ir.0)
+    STRWui killed %w1, %x0, 2 :: (store 4 into %ir.arrayidx1)
+    %w9 = LDRWui killed %x0, 0, implicit-def %x9  :: (load 4 from %ir.arrayidx19, align 8)
+    %x0 = ADDXrr killed %x9, killed %x8
+    RET_ReallyLR implicit %x0
+...
+