[AArch64] [MachineCombiner] Verify load instructions in gather pattern cannot possibly overlap

Author: jcohen-apple

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -22,6 +22,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/CodeGen/CFIInstBuilder.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
@@ -86,6 +87,11 @@ static cl::opt<unsigned>
     BDisplacementBits("aarch64-b-offset-bits", cl::Hidden, cl::init(26),
                       cl::desc("Restrict range of B instructions (DEBUG)"));
 
+#define DEBUG_TYPE "aarch64-machine-combine"
+STATISTIC(NumGathersMatched, "Number of `gather`-like patterns matched");
+STATISTIC(NumGathersDroppedAliasing, "Number of `gather`-like patterns dropped "
+                                     "due to potential pointer aliasing");
+
 AArch64InstrInfo::AArch64InstrInfo(const AArch64Subtarget &STI)
     : AArch64GenInstrInfo(AArch64::ADJCALLSTACKDOWN, AArch64::ADJCALLSTACKUP,
                           AArch64::CATCHRET),
@@ -7416,14 +7422,21 @@ static bool getGatherPattern(MachineInstr &Root,
   // 1. It has a single non-debug use (since we will be replacing the virtual
   // register)
   // 2. That the addressing mode only uses a single offset register.
+  // 3. The address operand does not have any users that are a COPY operation to
+  // a physical reg.
+  //    This could indicate that it is copied as part of an ABI of a function
+  //    call, which means that it may be modified in unexpected ways, see: <link
+  //    to github>
   auto *CurrInstr = MRI.getUniqueVRegDef(Root.getOperand(1).getReg());
   auto Range = llvm::seq<unsigned>(1, NumLanes - 1);
-  SmallSet<unsigned, 4> RemainingLanes(Range.begin(), Range.end());
+  SmallSet<unsigned, 16> RemainingLanes(Range.begin(), Range.end());
+  SmallSet<const MachineInstr *, 16> LoadInstrs = {};
   while (!RemainingLanes.empty() && CurrInstr &&
          CurrInstr->getOpcode() == LoadLaneOpCode &&
          MRI.hasOneNonDBGUse(CurrInstr->getOperand(0).getReg()) &&
          CurrInstr->getNumOperands() == 4) {
     RemainingLanes.erase(CurrInstr->getOperand(2).getImm());
+    LoadInstrs.insert(CurrInstr);
     CurrInstr = MRI.getUniqueVRegDef(CurrInstr->getOperand(1).getReg());
   }
 
@@ -7444,6 +7457,15 @@ static bool getGatherPattern(MachineInstr &Root,
   if (!MRI.hasOneNonDBGUse(Lane0LoadReg))
     return false;
 
+  LoadInstrs.insert(MRI.getUniqueVRegDef(Lane0LoadReg));
+
+  // Conservative check that we can
+  const TargetInstrInfo *TII = MF->getSubtarget().getInstrInfo();
+  for (auto LoadA = LoadInstrs.begin(); LoadA != LoadInstrs.end(); ++LoadA)
+    for (auto LoadB = ++LoadA; LoadB != LoadInstrs.end(); ++LoadB)
+      if (!TII->areMemAccessesTriviallyDisjoint(**LoadA, **LoadB))
+        return false;
+
   switch (NumLanes) {
   case 4:
     Patterns.push_back(AArch64MachineCombinerPattern::GATHER_LANE_i32);

llvm/test/CodeGen/AArch64/aarch64-combine-gather-lanes-with-call.mir

@@ -0,0 +1,45 @@
+# RUN: llc -run-pass=machine-combiner -mcpu=neoverse-n2 -mtriple=aarch64-none-linux-gnu -verify-machineinstrs %s -o - | FileCheck %s
+
+
+--- |
+  @external_func = external global i32
+  define void @negative_pattern_offset_reg_copied_to_physical(i64 %arg0, i64 %arg1, i64 %arg2, i64 %arg3, i64 %arg4) {
+  entry:  
+    ret void
+  }
+...
+---
+name:            negative_pattern_offset_reg_copied_to_physical
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1, $x2, $x3
+
+    ; CHECK-LABEL: name: negative_pattern_offset_reg_copied_to_physical
+    ; CHECK: [[BASE_REG:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK-NEXT: [[PTR_1:%[0-9]+]]:gpr64common = COPY $x1
+    ; CHECK-NEXT: [[PTR_2:%[0-9]+]]:gpr64common = COPY $x2
+    ; CHECK-NEXT: [[PTR_3:%[0-9]+]]:gpr64common = COPY $x3
+    ; CHECK-NEXT: [[LD_i32:%[0-9]+]]:fpr32 = LDRSroX [[BASE_REG]], killed [[PTR_1]], 0, 1
+    ; CHECK-NEXT: [[LD_LANE_0:%[0-9]+]]:fpr128 = SUBREG_TO_REG 0, killed [[LD_i32]], %subreg.ssub
+    ; CHECK-NEXT: [[LD_LANE_1:%[0-9]+]]:fpr128 = LD1i32 [[LD_LANE_0]], 1, [[PTR_2]]
+    ; CHECK-NEXT: $x0 = COPY [[PTR_2]]
+    ; CHECK-NEXT: BL @external_func, csr_aarch64_aapcs, implicit-def $lr, implicit $x0, implicit-def $x0
+    ; CHECK-NEXT: [[LD_LANE_2:%[0-9]+]]:fpr128 = LD1i32 [[LD_LANE_1]], 2, killed [[PTR_2]]
+    ; CHECK-NEXT: [[LD_LANE_3:%[0-9]+]]:fpr128 = LD1i32 [[LD_LANE_2]], 3, killed [[PTR_3]]
+    ; CHECK-NEXT: [[RESULT:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK-NEXT: $q0 = COPY [[LD_LANE_3]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $q0
+    %0:gpr64common = COPY $x0
+    %1:gpr64common = COPY $x1
+    %2:gpr64common = COPY $x2
+    %3:gpr64common = COPY $x3
+    %5:fpr32 = LDRSroX %0, killed %1, 0, 1
+    %6:fpr128 = SUBREG_TO_REG 0, killed %5, %subreg.ssub
+    %7:fpr128 = LD1i32 %6, 1, %2
+    $x0 = COPY %2
+    BL @external_func, csr_aarch64_aapcs, implicit-def $lr, implicit $x0, implicit-def $x0
+    %8:fpr128 = LD1i32 %7, 2, killed %2
+    %9:fpr128 = LD1i32 %8, 3, killed %3
+    %10:gpr64common = COPY $x0
+    $q0 = COPY %9
+    RET_ReallyLR implicit $q0

llvm/test/CodeGen/AArch64/aarch64-combine-gather-with-alias.mir

@@ -0,0 +1,42 @@
+# RUN: llc -run-pass=machine-combiner -mcpu=neoverse-n2 -mtriple=aarch64-none-linux-gnu -verify-machineinstrs %s -o - | FileCheck %s
+
+--- |
+  define void @aliasing_store_between_vector_loads(i64 %arg0, i64 %arg1, i64 %arg2, i64 %arg3) {
+  entry:
+    ret void
+  }
+
+...
+---
+name:            aliasing_store_between_vector_loads
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1, $x2, $x3
+
+    ; CHECK-LABEL: name: aliasing_store_between_vector_loads
+    ; CHECK: [[BASE_PTR:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK-NEXT: [[OFFSET_PTR:%[0-9]+]]:gpr64common = COPY $x1
+    ; CHECK-NEXT: [[ALIAS_ADDR:%[0-9]+]]:gpr64common = COPY $x2
+    ; CHECK-NEXT: [[OTHER_ADDR:%[0-9]+]]:gpr64common = COPY $x3
+    ; CHECK-NEXT: [[VEC1:%[0-9]+]]:fpr128 = LD1i32 %{{[0-9]+}}, 1, [[ALIAS_ADDR]]
+    ; CHECK-NEXT: [[CONST:%[0-9]+]]:gpr32 = MOVi32 99
+    ; CHECK-NEXT: STRWui [[CONST]], [[ALIAS_ADDR]], 0
+    ; CHECK-NEXT: [[VEC2:%[0-9]+]]:fpr128 = LD1i32 [[VEC1]], 2, killed [[ALIAS_ADDR]]
+    ; CHECK-NEXT: [[VEC3:%[0-9]+]]:fpr128 = LD1i32 [[VEC2]], 3, killed [[OTHER_ADDR]]
+    ; CHECK-NEXT: $q0 = COPY [[VEC3]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $q0
+    %0:gpr64common = COPY $x0
+    %1:gpr64common = COPY $x1
+    %2:gpr64common = COPY $x2
+    %3:gpr64common = COPY $x3
+    %5:fpr32 = LDRSroX %0, killed %1, 0, 1
+    %6:fpr128 = SUBREG_TO_REG 0, killed %5, %subreg.ssub
+    %7:fpr128 = LD1i32 %6, 1, %2
+    %10:gpr32 = MOVi32imm 99
+    STRWui %10, %2, 0
+    %8:fpr128 = LD1i32 %7, 2, killed %2
+    %9:fpr128 = LD1i32 %8, 3, killed %3
+    $q0 = COPY %9
+    RET_ReallyLR implicit $q0