[AArch64][GlobalISel] Add post-legalize combine for sext(trunc(sextload)) -> trunc/copy

[aemerson]

Summary

On AArch64 we generate redundant G_SEXTs or G_SEXT_INREGs because of this.

[arsenm]

This doesn't seem necessary for correctness? I think we do need some optimizing combines in the legalizer itself, but I assumed we would start putting those in a distinct place

[aemerson]

In D81993#2097971, @arsenm wrote:

This doesn't seem necessary for correctness? I think we do need some optimizing combines in the legalizer itself, but I assumed we would start putting those in a distinct place

True, I suppose we need to be stricter about these combines. I'll move this into the post legalizer combiner.

[aemerson]

Move the combine to post legalizer, now matching what ends up coming out of the legalizer for AArch64 which is a G_SEXT_INREG.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

588

This is just a special case of checking computeNumSignBits? Do we recognize the number of sign bits in a sextload yet?

[aemerson]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

588

We don't, but using computeNumSignBits doesn't tell us how to find the original def in the chain that has redundantly been sign extended. E.g. in this case we won't know how to look through the truncate.

[aemerson]

Use computeNumSignBits() since we can tolerate extra copies with SEXT_INREG.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

586

This should be >=?

592

I don't think introducing a new AnyExt would be right?

[aemerson]

Address comments. Use copy instead of buildAnyExtOrTrunc.

[aemerson]

Ping

[arsenm]

llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp

447

I think this needs to be careful about a vector sextload

[aemerson]

Don't try to handle vector SEXTLOAD in KB.

[arsenm]

LGTM with nits for future vector handling

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

584

getScalarSizeInBits

llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp

448–449

Should add TODO to handle vectors

[aemerson]

I reverted this because it turns out that computeNumSignBits() doesn't actually do what we need here. Its definition of sign bits are any top-most bits which are known to be identical. Because of that it matches G_ZEXTLOAD with the known zero bits at the top returns as "sign bits". This behavior seems to be consistent with the SDAG KB implementation too.

I committed the original version of this patch in 3b10e42ba1a3ff97346bc51f13195ed9595800f4, with the vector bail out added. Let me know if I should change it post commit.

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

public:
   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 matchSextAlreadyExtended(MachineInstr &MI);
+  bool applySextAlreadyExtended(MachineInstr &MI);
+
   bool matchElideBrByInvertingCond(MachineInstr &MI);
   void applyElideBrByInvertingCond(MachineInstr &MI);
   bool tryElideBrByInvertingCond(MachineInstr &MI);
 
   /// If \p MI is G_CONCAT_VECTORS, try to combine it.
   /// Returns true if MI changed.
   /// Right now, we support:
   /// - concat_vector(undef, undef) => undef

llvm/include/llvm/Target/GlobalISel/Combine.td

 
 def extending_loads : GICombineRule<
   (defs root:$root, extending_load_matchdata:$matchinfo),
   (match (wip_match_opcode G_LOAD, G_SEXTLOAD, G_ZEXTLOAD):$root,
          [{ return Helper.matchCombineExtendingLoads(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyCombineExtendingLoads(*${root}, ${matchinfo}); }])>;
 def combines_for_extload: GICombineGroup<[extending_loads]>;
 
+def sext_already_extended : GICombineRule<
+  (defs root:$d),
+  (match (wip_match_opcode G_SEXT_INREG):$d,
+         [{ return Helper.matchSextAlreadyExtended(*${d}); }]),
+  (apply [{ Helper.applySextAlreadyExtended(*${d}); }])>;
+
 def combine_indexed_load_store : GICombineRule<
   (defs root:$root, indexed_load_store_matchdata:$matchinfo),
   (match (wip_match_opcode G_LOAD, G_SEXTLOAD, G_ZEXTLOAD, G_STORE):$root,
          [{ return Helper.matchCombineIndexedLoadStore(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyCombineIndexedLoadStore(*${root}, ${matchinfo}); }])>;
 
 // FIXME: Is there a reason this wasn't in tryCombine? I've left it out of
 //        all_combines because it wasn't there.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

bool CombinerHelper::dominates(const MachineInstr &DefMI,
   if (MDT)
     return MDT->dominates(&DefMI, &UseMI);
   else if (DefMI.getParent() != UseMI.getParent())
     return false;
 
   return isPredecessor(DefMI, UseMI);
 }
 
+bool CombinerHelper::matchSextAlreadyExtended(MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::G_SEXT_INREG);
+  Register SrcReg = MI.getOperand(1).getReg();
+  unsigned SrcSignBits = KB->computeNumSignBits(SrcReg);
+  unsigned NumSextBits =
+      MRI.getType(MI.getOperand(0).getReg()).getScalarSizeInBits() -
+      MI.getOperand(2).getImm();
+  return SrcSignBits >= NumSextBits;
+}
+
+bool CombinerHelper::applySextAlreadyExtended(MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::G_SEXT_INREG);
+  MachineIRBuilder MIB(MI);
+  MIB.buildCopy(MI.getOperand(0).getReg(), MI.getOperand(1).getReg());
+  MI.eraseFromParent();
+  return true;
+}
+
 bool CombinerHelper::findPostIndexCandidate(MachineInstr &MI, Register &Addr,
                                             Register &Base, Register &Offset) {
   auto &MF = *MI.getParent()->getParent();
   const auto &TLI = *MF.getSubtarget().getTargetLowering();
 
 #ifndef NDEBUG
   unsigned Opcode = MI.getOpcode();
   assert(Opcode == TargetOpcode::G_LOAD || Opcode == TargetOpcode::G_SEXTLOAD ||

llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp

 //===- lib/CodeGen/GlobalISel/GISelKnownBits.cpp --------------*- C++ *-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 /// Provides analysis for querying information about KnownBits during GISel
 /// passes.
 //
 //===------------------
 #include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 
 #define DEBUG_TYPE "gisel-known-bits"
case TargetOpcode::COPY: {
     return 1;
   }
   case TargetOpcode::G_SEXT: {
     Register Src = MI.getOperand(1).getReg();
     LLT SrcTy = MRI.getType(Src);
     unsigned Tmp = DstTy.getScalarSizeInBits() - SrcTy.getScalarSizeInBits();
     return computeNumSignBits(Src, DemandedElts, Depth + 1) + Tmp;
   }
+  case TargetOpcode::G_SEXTLOAD: {
+    Register Dst = MI.getOperand(0).getReg();
+    LLT Ty = MRI.getType(Dst);
+    // TODO: add vector support
+    if (Ty.isVector())
+      break;
+    if (MI.hasOneMemOperand())
+      return Ty.getSizeInBits() - (*MI.memoperands_begin())->getSizeInBits();
+    break;
+  }
   case TargetOpcode::G_TRUNC: {
     Register Src = MI.getOperand(1).getReg();
     LLT SrcTy = MRI.getType(Src);
 
     // Check if the sign bits of source go down as far as the truncated value.
     unsigned DstTyBits = DstTy.getScalarSizeInBits();
     unsigned NumSrcBits = SrcTy.getScalarSizeInBits();
     unsigned NumSrcSignBits = computeNumSignBits(Src, DemandedElts, Depth + 1);

llvm/lib/Target/AArch64/AArch64Combine.td

 
 // Combines which replace a G_SHUFFLE_VECTOR with a target-specific pseudo
 // instruction.
 def shuffle_vector_pseudos : GICombineGroup<[dup, rev, ext, zip, uzp, trn]>;
 
 def AArch64PostLegalizerCombinerHelper
     : GICombinerHelper<"AArch64GenPostLegalizerCombinerHelper",
                        [erase_undef_store, combines_for_extload,
-                        shuffle_vector_pseudos]> {
+                        sext_already_extended, shuffle_vector_pseudos]> {
   let DisableRuleOption = "aarch64postlegalizercombiner-disable-rule";
 }

llvm/test/CodeGen/AArch64/GlobalISel/combine-sext-trunc-sextload.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple aarch64 -run-pass=aarch64-postlegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+---
+name:            test_combine_sext_trunc_of_sextload
+legalized:       true
+tracksRegLiveness: true
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: test_combine_sext_trunc_of_sextload
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[SEXTLOAD:%[0-9]+]]:_(s64) = G_SEXTLOAD [[COPY]](p0) :: (load 2)
+    ; CHECK: [[TRUNC:%[0-9]+]]:_(s32) = G_TRUNC [[SEXTLOAD]](s64)
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY [[TRUNC]](s32)
+    ; CHECK: $w0 = COPY [[COPY1]](s32)
+    %0:_(p0) = COPY $x0
+    %1:_(s64) = G_SEXTLOAD %0:_(p0) :: (load 2)
+    %2:_(s32) = G_TRUNC %1:_(s64)
+    %3:_(s32) = G_SEXT_INREG %2:_(s32), 16
+    $w0 = COPY %3(s32)
+...
+---
+name:            test_combine_sext_of_sextload
+legalized:       true
+tracksRegLiveness: true
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: test_combine_sext_of_sextload
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[SEXTLOAD:%[0-9]+]]:_(s32) = G_SEXTLOAD [[COPY]](p0) :: (load 2)
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY [[SEXTLOAD]](s32)
+    ; CHECK: [[COPY2:%[0-9]+]]:_(s32) = COPY [[COPY1]](s32)
+    ; CHECK: $w0 = COPY [[COPY2]](s32)
+    %0:_(p0) = COPY $x0
+    %1:_(s32) = G_SEXTLOAD %0:_(p0) :: (load 2)
+    %2:_(s32) = COPY %1:_(s32)
+    %3:_(s32) = G_SEXT_INREG %2:_(s32), 16
+    $w0 = COPY %3(s32)
+...
+---
+name:            test_combine_sext_of_sextload_not_matching
+legalized:       true
+tracksRegLiveness: true
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; Here we're trying to extend from a larger width than was extended in the load.
+    ; CHECK-LABEL: name: test_combine_sext_of_sextload_not_matching
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[SEXTLOAD:%[0-9]+]]:_(s32) = G_SEXTLOAD [[COPY]](p0) :: (load 2)
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY [[SEXTLOAD]](s32)
+    ; CHECK: $w0 = COPY [[COPY1]](s32)
+    %0:_(p0) = COPY $x0
+    %1:_(s32) = G_SEXTLOAD %0:_(p0) :: (load 2)
+    %2:_(s32) = G_SEXT_INREG %1:_(s32), 24
+    $w0 = COPY %2(s32)
+...
+---
+name:            test_combine_sext_of_sextload_not_enough_src_signbits
+legalized:       true
+tracksRegLiveness: true
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; Here we're trying to extend from a smaller width than was extended in the load.
+    ; Don't perform the combine.
+    ; CHECK-LABEL: name: test_combine_sext_of_sextload_not_enough_src_signbits
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[SEXTLOAD:%[0-9]+]]:_(s32) = G_SEXTLOAD [[COPY]](p0) :: (load 2)
+    ; CHECK: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[SEXTLOAD]], 8
+    ; CHECK: $w0 = COPY [[SEXT_INREG]](s32)
+    %0:_(p0) = COPY $x0
+    %1:_(s32) = G_SEXTLOAD %0:_(p0) :: (load 2)
+    %2:_(s32) = G_SEXT_INREG %1:_(s32), 8
+    $w0 = COPY %2(s32)
+...