[AArch64][GlobalISel] Enable extending loads combines post-legalization.

[aemerson]

Summary

During legalization we can end up with extends of loads, which in the case of zexts causes us to not hit tablegen imported patterns.

The caveat here is that we don't want anyext load forming, since some variants are illegal.

[aemerson]

Forgot to run update script on test.

[arsenm]

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

102–106

This should just fix the combiner to understand checking for legal loads rather than special casing anyext?

[aemerson]

Check for legality in the combiner helper.

[aemerson]

ping

[paquette]

LGTM

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

 
 class GISelChangeObserver;
 class MachineIRBuilder;
 class MachineRegisterInfo;
 class MachineInstr;
 class MachineOperand;
 class GISelKnownBits;
 class MachineDominatorTree;
+class LegalizerInfo;
 
 struct PreferredTuple {
   LLT Ty;                // The result type of the extend.
   unsigned ExtendOpcode; // G_ANYEXT/G_SEXT/G_ZEXT
   MachineInstr *MI;
 };
 
 struct IndexedLoadStoreMatchInfo {

 
 class CombinerHelper {
 protected:
   MachineIRBuilder &Builder;
   MachineRegisterInfo &MRI;
   GISelChangeObserver &Observer;
   GISelKnownBits *KB;
   MachineDominatorTree *MDT;
+  const LegalizerInfo *LI;
 
 public:
   CombinerHelper(GISelChangeObserver &Observer, MachineIRBuilder &B,
                  GISelKnownBits *KB = nullptr,
-                 MachineDominatorTree *MDT = nullptr);
+                 MachineDominatorTree *MDT = nullptr,
+                 const LegalizerInfo *LI = nullptr);
 
   GISelKnownBits *getKnownBits() const {
     return KB;
   }
 
   /// MachineRegisterInfo::replaceRegWith() and inform the observer of the changes
   void replaceRegWith(MachineRegisterInfo &MRI, Register FromReg, Register ToReg) const;
 
public:
   /// construction, this function returns a conservative result that tracks just
   /// a single basic block.
   bool dominates(const MachineInstr &DefMI, const MachineInstr &UseMI);
 
   /// If \p MI is extend that consumes the result of a load, try to combine it.
   /// Returns true if MI changed.
   bool tryCombineExtendingLoads(MachineInstr &MI);
   bool matchCombineExtendingLoads(MachineInstr &MI, PreferredTuple &MatchInfo);
   void applyCombineExtendingLoads(MachineInstr &MI, PreferredTuple &MatchInfo);
 
   /// Combine \p MI into a pre-indexed or post-indexed load/store operation if
   /// legal and the surrounding code makes it useful.
   bool tryCombineIndexedLoadStore(MachineInstr &MI);
   bool matchCombineIndexedLoadStore(MachineInstr &MI, IndexedLoadStoreMatchInfo &MatchInfo);
   void applyCombineIndexedLoadStore(MachineInstr &MI, IndexedLoadStoreMatchInfo &MatchInfo);
 
   bool matchElideBrByInvertingCond(MachineInstr &MI);
   void applyElideBrByInvertingCond(MachineInstr &MI);
   bool tryElideBrByInvertingCond(MachineInstr &MI);
 
   /// If \p MI is G_CONCAT_VECTORS, try to combine it.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

 //===-- lib/CodeGen/GlobalISel/GICombinerHelper.cpp -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/CodeGen/GlobalISel/CombinerHelper.h"
 #include "llvm/CodeGen/GlobalISel/Combiner.h"
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
 #include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
+#include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"

 static cl::opt<bool>
     ForceLegalIndexing("force-legal-indexing", cl::Hidden, cl::init(false),
                        cl::desc("Force all indexed operations to be "
                                 "legal for the GlobalISel combiner"));
 
 
 CombinerHelper::CombinerHelper(GISelChangeObserver &Observer,
                                MachineIRBuilder &B, GISelKnownBits *KB,
-                               MachineDominatorTree *MDT)
+                               MachineDominatorTree *MDT,
+                               const LegalizerInfo *LI)
     : Builder(B), MRI(Builder.getMF().getRegInfo()), Observer(Observer),
-      KB(KB), MDT(MDT) {
+      KB(KB), MDT(MDT), LI(LI) {
   (void)this->KB;
 }
 
 void CombinerHelper::replaceRegWith(MachineRegisterInfo &MRI, Register FromReg,
                                     Register ToReg) const {
   Observer.changingAllUsesOfReg(MRI, FromReg);
 
   if (MRI.constrainRegAttrs(ToReg, FromReg))
unsigned PreferredOpcode = MI.getOpcode() == TargetOpcode::G_LOAD
                                  ? TargetOpcode::G_ANYEXT
                                  : MI.getOpcode() == TargetOpcode::G_SEXTLOAD
                                        ? TargetOpcode::G_SEXT
                                        : TargetOpcode::G_ZEXT;
   Preferred = {LLT(), PreferredOpcode, nullptr};
   for (auto &UseMI : MRI.use_nodbg_instructions(LoadValue.getReg())) {
     if (UseMI.getOpcode() == TargetOpcode::G_SEXT ||
         UseMI.getOpcode() == TargetOpcode::G_ZEXT ||
-        UseMI.getOpcode() == TargetOpcode::G_ANYEXT) {
+        (UseMI.getOpcode() == TargetOpcode::G_ANYEXT)) {
+      // Check for legality.
+      if (LI) {
+        LegalityQuery::MemDesc MMDesc;
+        const auto &MMO = **MI.memoperands_begin();
+        MMDesc.SizeInBits = MMO.getSizeInBits();
+        MMDesc.AlignInBits = MMO.getAlign().value() * 8;
+        MMDesc.Ordering = MMO.getOrdering();
+        LLT UseTy = MRI.getType(UseMI.getOperand(0).getReg());
+        LLT SrcTy = MRI.getType(MI.getOperand(1).getReg());
+        if (LI->getAction({MI.getOpcode(), {UseTy, SrcTy}, {MMDesc}}).Action !=
+            LegalizeActions::Legal)
+          continue;
+      }
       Preferred = ChoosePreferredUse(Preferred,
                                      MRI.getType(UseMI.getOperand(0).getReg()),
                                      UseMI.getOpcode(), &UseMI);
     }
   }
 
   // There were no extends
   if (!Preferred.MI)

llvm/lib/Target/AArch64/AArch64Combine.td

 
 def AArch64PreLegalizerCombinerHelper: GICombinerHelper<
   "AArch64GenPreLegalizerCombinerHelper", [all_combines,
                                            elide_br_by_inverting_cond,
                                            fconstant_to_constant]> {
   let DisableRuleOption = "aarch64prelegalizercombiner-disable-rule";
 }
 
-def AArch64PostLegalizerCombinerHelper: GICombinerHelper<
-  "AArch64GenPostLegalizerCombinerHelper", [erase_undef_store]> {
+def AArch64PostLegalizerCombinerHelper
+    : GICombinerHelper<"AArch64GenPostLegalizerCombinerHelper",
+                       [erase_undef_store, combines_for_extload]> {
   let DisableRuleOption = "aarch64postlegalizercombiner-disable-rule";
 }

llvm/lib/Target/AArch64/AArch64PostLegalizerCombiner.cpp

public:
 
   virtual bool combine(GISelChangeObserver &Observer, MachineInstr &MI,
                        MachineIRBuilder &B) const override;
 };
 
 bool AArch64PostLegalizerCombinerInfo::combine(GISelChangeObserver &Observer,
                                                MachineInstr &MI,
                                                MachineIRBuilder &B) const {
-  CombinerHelper Helper(Observer, B, KB, MDT);
+  const auto *LI =
+      MI.getParent()->getParent()->getSubtarget().getLegalizerInfo();
+  CombinerHelper Helper(Observer, B, KB, MDT, LI);
   return Generated.tryCombineAll(Observer, MI, B, Helper);
 }
 
 #define AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_CPP
 #include "AArch64GenPostLegalizeGICombiner.inc"
 #undef AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_CPP
 
 class AArch64PostLegalizerCombiner : public MachineFunctionPass {

llvm/test/CodeGen/AArch64/GlobalISel/postlegalizercombiner-extending-loads.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -run-pass=aarch64-postlegalizer-combiner -global-isel -verify-machineinstrs %s -o - | FileCheck %s
+
+--- |
+  target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64--"
+  define void @test_zeroext(i8* %addr) {
+  entry:
+    ret void
+  }
+  define void @test_no_anyext(i8* %addr) {
+  entry:
+    ret void
+  }
+...
+
+---
+name:            test_zeroext
+legalized:       true
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: test_zeroext
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[ZEXTLOAD:%[0-9]+]]:_(s32) = G_ZEXTLOAD [[COPY]](p0) :: (load 1 from %ir.addr)
+    ; CHECK: $w0 = COPY [[ZEXTLOAD]](s32)
+    %0:_(p0) = COPY $x0
+    %1:_(s8) = G_LOAD %0 :: (load 1 from %ir.addr)
+    %2:_(s32) = G_ZEXT %1
+    $w0 = COPY %2
+...
+
+---
+name:            test_no_anyext
+legalized:       true
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; Check that we don't try to do an anyext combine. We don't want to do this
+    ; because an anyexting load like s64 = G_LOAD %p (load 4) isn't legal.
+    ; CHECK-LABEL: name: test_no_anyext
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[LOAD:%[0-9]+]]:_(s32) = G_LOAD [[COPY]](p0) :: (load 4 from %ir.addr)
+    ; CHECK: [[ANYEXT:%[0-9]+]]:_(s64) = G_ANYEXT [[LOAD]](s32)
+    ; CHECK: $x0 = COPY [[ANYEXT]](s64)
+    %0:_(p0) = COPY $x0
+    %1:_(s32) = G_LOAD %0 :: (load 4 from %ir.addr)
+    %2:_(s64) = G_ANYEXT %1
+    $x0 = COPY %2
+...