[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 ↗

(On Diff #271520)

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 ↗

(On Diff #271520)

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 ↗

(On Diff #272264)

This should be >=?

592 ↗

(On Diff #272264)

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 ↗

(On Diff #274927)

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 ↗

(On Diff #277495)

getScalarSizeInBits

llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp

448–449 ↗

(On Diff #277495)

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/LegalizationArtifactCombiner.h

bool tryCombineSExt(MachineInstr &MI,
     Register TruncSrc;
     if (mi_match(SrcReg, MRI, m_GTrunc(m_Reg(TruncSrc)))) {
       LLT DstTy = MRI.getType(DstReg);
       if (isInstUnsupported({TargetOpcode::G_SEXT_INREG, {DstTy}}))
         return false;
       LLVM_DEBUG(dbgs() << ".. Combine MI: " << MI;);
       LLT SrcTy = MRI.getType(SrcReg);
       uint64_t SizeInBits = SrcTy.getScalarSizeInBits();
+
+      // If the source is a G_SEXTLOAD from the same bit width, then we don't
+      // need any extend at all, just a truncate.
+      if (auto *LoadMI =
+              getOpcodeDef(TargetOpcode::G_SEXTLOAD, TruncSrc, MRI)) {
+        const auto &MMO = **LoadMI->memoperands_begin();
+        if (MMO.getSizeInBits() == SizeInBits) {
+          Builder.buildAnyExtOrTrunc(DstReg, LoadMI->getOperand(0).getReg());
+          UpdatedDefs.push_back(DstReg);
+          markInstAndDefDead(MI, *MRI.getVRegDef(SrcReg), DeadInsts);
+          return true;
+        }
+      }
+
       Builder.buildInstr(
           TargetOpcode::G_SEXT_INREG, {DstReg},
           {Builder.buildAnyExtOrTrunc(DstTy, TruncSrc), SizeInBits});
       markInstAndDefDead(MI, *MRI.getVRegDef(SrcReg), DeadInsts);
       return true;
     }
 
     // sext(zext x) -> (zext x)

llvm/test/CodeGen/AArch64/GlobalISel/legalize-sext.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -march=aarch64 -run-pass=legalizer %s -o - | FileCheck %s
+# RUN: llc -march=aarch64 -run-pass=legalizer %s -o - -verify-machineinstrs | FileCheck %s
 ---
 name:            test_sext_inreg
 body: |
   bb.0.entry:
     liveins: $w0, $w1
     ; CHECK-LABEL: name: test_sext_inreg
     ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w1
     ; CHECK: [[SEXT_INREG:%[0-9]+]]:_(s32) = G_SEXT_INREG [[COPY]], 7
     ; CHECK: $w0 = COPY [[SEXT_INREG]](s32)
     %0:_(s32) = COPY $w1
     %2:_(s32) = G_SEXT_INREG %0(s32), 7
     $w0 = COPY %2(s32)
 ...
+---
+name:            test_combine_sext_trunc_of_sextload
+body: |
+  bb.0.entry:
+    liveins: $w0, $w1
+    ; CHECK-LABEL: name: test_combine_sext_trunc_of_sextload
+    ; 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: $w0 = COPY [[TRUNC]](s32)
+    %0:_(p0) = COPY $x0
+    %1:_(s64) = G_SEXTLOAD %0:_(p0) :: (load 2)
+    %2:_(s16) = G_TRUNC %1:_(s64)
+    %3:_(s32) = G_SEXT %2:_(s16)
+    $w0 = COPY %3(s32)
+...