[llvm] [GlobalISel] Make sure to check for load barriers when merging G_EXTRACT_VECTOR_ELT into G_LOAD. (PR #82306)

[Owen Anderson via llvm-commits]

https://github.com/resistor created https://github.com/llvm/llvm-project/pull/82306

Fixes https://github.com/llvm/llvm-project/issues/78477

>From 3489398845d2758a6e9ee0b70e44f7a233b9bc13 Mon Sep 17 00:00:00 2001
From: Owen Anderson <resistor at mac.com>
Date: Tue, 20 Feb 2024 00:03:19 -0500
Subject: [PATCH] [GlobalISel] Make sure to check for load barriers when
 merging G_EXTRACT_VECTOR_ELT into G_LOAD.

Fixes https://github.com/llvm/llvm-project/issues/78477
---
 .../lib/CodeGen/GlobalISel/CombinerHelper.cpp | 12 +++
 .../CodeGen/AArch64/extractvector-of-load.mir | 91 +++++++++++++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/extractvector-of-load.mir

diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index b400eb34e2901b..45427fc2466bd1 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -1198,6 +1198,18 @@ bool CombinerHelper::matchCombineExtractedVectorLoad(MachineInstr &MI,
   if (!VecEltTy.isByteSized())
     return false;
 
+  // Check for load fold barriers between the extraction and the load.
+  if (MI.getParent() != LoadMI->getParent())
+    return false;
+  const unsigned MaxIter = 20;
+  unsigned Iter = 0;
+  for (auto II = LoadMI->getIterator(), IE = MI.getIterator(); II != IE; ++II) {
+    if (II->isLoadFoldBarrier())
+      return false;
+    if (Iter++ == MaxIter)
+      return false;
+  }
+
   // Check if the new load that we are going to create is legal
   // if we are in the post-legalization phase.
   MachineMemOperand MMO = LoadMI->getMMO();
diff --git a/llvm/test/CodeGen/AArch64/extractvector-of-load.mir b/llvm/test/CodeGen/AArch64/extractvector-of-load.mir
new file mode 100644
index 00000000000000..d4c28917300760
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/extractvector-of-load.mir
@@ -0,0 +1,91 @@
+# RUN: llc -run-pass=aarch64-prelegalizer-combiner %s -o - | FileCheck %s
+
+--- |
+  ; ModuleID = 'in.ll'
+  source_filename = "in.ll"
+  target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+  
+  define i32 @f(ptr %0) {
+    %2 = load <2 x i32>, ptr %0, align 8
+    store <4 x i32> zeroinitializer, ptr %0, align 16
+    %3 = extractelement <2 x i32> %2, i64 0
+    ret i32 %3
+  }
+
+...
+---
+name:            f
+alignment:       4
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+failedISel:      false
+tracksRegLiveness: true
+hasWinCFI:       false
+callsEHReturn:   false
+callsUnwindInit: false
+hasEHCatchret:   false
+hasEHScopes:     false
+hasEHFunclets:   false
+isOutlined:      false
+debugInstrRef:   false
+failsVerification: false
+tracksDebugUserValues: false
+registers:
+  - { id: 0, class: _, preferred-register: '' }
+  - { id: 1, class: _, preferred-register: '' }
+  - { id: 2, class: _, preferred-register: '' }
+  - { id: 3, class: _, preferred-register: '' }
+  - { id: 4, class: _, preferred-register: '' }
+  - { id: 5, class: _, preferred-register: '' }
+liveins:
+  - { reg: '$x0', virtual-reg: '' }
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   false
+  stackSize:       0
+  offsetAdjustment: 0
+  maxAlignment:    1
+  adjustsStack:    false
+  hasCalls:        false
+  stackProtector:  ''
+  functionContext: ''
+  maxCallFrameSize: 4294967295
+  cvBytesOfCalleeSavedRegisters: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+  hasTailCall:     false
+  localFrameSize:  0
+  savePoint:       ''
+  restorePoint:    ''
+fixedStack:      []
+stack:           []
+entry_values:    []
+callSites:       []
+debugValueSubstitutions: []
+constants:       []
+machineFunctionInfo: {}
+body:             |
+  bb.1 (%ir-block.1):
+    liveins: $x0
+  
+    %0:_(p0) = COPY $x0
+    %3:_(s32) = G_CONSTANT i32 0
+    %2:_(<4 x s32>) = G_BUILD_VECTOR %3(s32), %3(s32), %3(s32), %3(s32)
+    %5:_(s64) = G_CONSTANT i64 0
+    %1:_(<2 x s32>) = G_LOAD %0(p0) :: (load (<2 x s32>) from %ir.0)
+    G_STORE %2(<4 x s32>), %0(p0) :: (store (<4 x s32>) into %ir.0)
+    %4:_(s32) = G_EXTRACT_VECTOR_ELT %1(<2 x s32>), %5(s64)
+    $w0 = COPY %4(s32)
+    RET_ReallyLR implicit $w0
+
+...
+
+# CHECK: bb.0
+# CHECK: G_LOAD
+# CHECK: G_STORE
+# CHECK: G_EXTRACT_VECTOR_ELT