[GlobalISel] LegalizationArtifactCombiner: Combine aext([asz]ext x) -> [asz]ext x

[volkan]

Summary

Replace aext([asz]ext x) with aext/sext/zext x in order to
reduce the number of instructions generated to clean up some
legalization artifacts.

[dsanders]

What's the reason for not putting this in a post-legalize combiner? It seems to me that we'll have to implement this in the Legalizer and the Combiner

[volkan]

In D54174#1291716, @dsanders wrote:

What's the reason for not putting this in a post-legalize combiner? It seems to me that we'll have to implement this in the Legalizer and the Combiner

The main reason is generating a better code by doing this in the legalizer. If we don't handle this in the legalizer, it would be really hard to combine
in a post-legalize combine as LegalizationArtifactCombiner replaces the pattern with shifts and ands.
Here is an example:

%0:_(s32) = COPY $w0
%1:_(s1) = G_TRUNC %0(s32)
%2:_(s8) = G_SEXT %1(s1)
%3:_(s32) = G_ANYEXT %2(s8)
$w0 = COPY %3(s32)

Currently, we generate:

%0:_(s32) = COPY $w0
%13:_(s32) = G_CONSTANT i32 7
%4:_(s8) = G_TRUNC %13(s32)
%10:_(s32) = COPY %0(s32)
%16:_(s32) = G_CONSTANT i32 255
%17:_(s32) = COPY %13(s32)
%11:_(s32) = G_AND %17, %16
%12:_(s32) = G_SHL %10, %11
%18:_(s32) = G_CONSTANT i32 24
%19:_(s32) = COPY %12(s32)
%20:_(s32) = G_SHL %19, %18
%7:_(s32) = G_ASHR %20, %18
%14:_(s32) = G_CONSTANT i32 255
%15:_(s32) = COPY %13(s32)
%8:_(s32) = G_AND %15, %14
%9:_(s32) = G_ASHR %7, %8
%3:_(s32) = COPY %9(s32)
$w0 = COPY %3(s32)

With D54174 applied, we generate:

%0:_(s32) = COPY $w0
%4:_(s32) = G_CONSTANT i32 31
%5:_(s32) = COPY %0(s32)
%6:_(s32) = G_SHL %5, %4
%3:_(s32) = G_ASHR %6, %4
$w0 = COPY %3(s32)

[aemerson]

So these extends are created by the legalizer, correct? If so, it seems ok to me as that's what these artifact combines were intended for.

[volkan]

In D54174#1299872, @aemerson wrote:

So these extends are created by the legalizer, correct? If so, it seems ok to me as that's what these artifact combines were intended for.

Yes, they are created by the legalizer.

include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h

bool tryCombineAnyExt(MachineInstr &MI,
     // aext(trunc x) - > aext/copy/trunc x
     unsigned TruncSrc;
     if (mi_match(SrcReg, MRI, m_GTrunc(m_Reg(TruncSrc)))) {
       LLVM_DEBUG(dbgs() << ".. Combine MI: " << MI;);
       Builder.buildAnyExtOrTrunc(DstReg, TruncSrc);
       markInstAndDefDead(MI, *MRI.getVRegDef(SrcReg), DeadInsts);
       return true;
     }
+
+    // aext([asz]ext x) -> [asz]ext x
+    unsigned ExtSrc;
+    MachineInstr *ExtMI;
+    if (mi_match(SrcReg, MRI,
+                 m_all_of(m_MInstr(ExtMI), m_any_of(m_GAnyExt(m_Reg(ExtSrc)),
+                                                    m_GSExt(m_Reg(ExtSrc)),
+                                                    m_GZExt(m_Reg(ExtSrc)))))) {
+      Builder.buildInstr(ExtMI->getOpcode(), DstReg, ExtSrc);
+      markInstAndDefDead(MI, *ExtMI, DeadInsts);
+      return true;
+    }
     return tryFoldImplicitDef(MI, DeadInsts);
   }
 
   bool tryCombineZExt(MachineInstr &MI,
                       SmallVectorImpl<MachineInstr *> &DeadInsts) {
 
     if (MI.getOpcode() != TargetOpcode::G_ZEXT)
       return false;

test/CodeGen/AArch64/GlobalISel/legalize-ext.mir

bb.0.entry:
     ; CHECK: $w0 = COPY [[TRUNC8]](s32)
     ; CHECK: [[C5:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
     ; CHECK: [[TRUNC9:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK: [[SHL2:%[0-9]+]]:_(s32) = G_SHL [[TRUNC9]], [[C5]]
     ; CHECK: [[ASHR2:%[0-9]+]]:_(s32) = G_ASHR [[SHL2]], [[C5]]
     ; CHECK: $w0 = COPY [[ASHR2]](s32)
     ; CHECK: [[C6:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
     ; CHECK: [[TRUNC10:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
-    ; CHECK: [[COPY5:%[0-9]+]]:_(s32) = COPY [[C6]](s32)
-    ; CHECK: [[AND3:%[0-9]+]]:_(s32) = G_AND [[TRUNC10]], [[COPY5]]
-    ; CHECK: [[COPY6:%[0-9]+]]:_(s32) = COPY [[AND3]](s32)
-    ; CHECK: $w0 = COPY [[COPY6]](s32)
+    ; CHECK: [[AND3:%[0-9]+]]:_(s32) = G_AND [[TRUNC10]], [[C6]]
+    ; CHECK: $w0 = COPY [[AND3]](s32)
     ; CHECK: [[TRUNC11:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK: $w0 = COPY [[TRUNC11]](s32)
     ; CHECK: [[TRUNC12:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
     ; CHECK: $w0 = COPY [[TRUNC12]](s32)
     ; CHECK: [[FPEXT:%[0-9]+]]:_(s64) = G_FPEXT [[TRUNC12]](s32)
     ; CHECK: $x0 = COPY [[FPEXT]](s64)
     ; CHECK: [[C7:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
     ; CHECK: $w0 = COPY [[C7]](s32)
bb.0.entry:
     %25:_(s32) = G_ZEXT %24(s16)
     $w0 = COPY %25(s32)
     %26:_(s32) = G_SEXT %24(s16)
     $w0 = COPY %26(s32)
     %27:_(s32) = G_ANYEXT %24(s16)
     $w0 = COPY %27(s32)
 
 ...
+---
+name:            test_anyext_anyext
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: test_anyext_anyext
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY [[COPY]](s32)
+    ; CHECK: $w0 = COPY [[COPY1]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s1) = G_TRUNC %0(s32)
+    %2:_(s8) = G_ANYEXT %1(s1)
+    %3:_(s32) = G_ANYEXT %2(s8)
+    $w0 = COPY %3(s32)
+
+...
+---
+name:            test_anyext_sext
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: test_anyext_sext
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 31
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY [[COPY]](s32)
+    ; CHECK: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY1]], [[C]]
+    ; CHECK: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]]
+    ; CHECK: $w0 = COPY [[ASHR]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s1) = G_TRUNC %0(s32)
+    %2:_(s8) = G_SEXT %1(s1)
+    %3:_(s32) = G_ANYEXT %2(s8)
+    $w0 = COPY %3(s32)
+
+...
+---
+name:            test_anyext_zext
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: test_anyext_zext
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY [[COPY]](s32)
+    ; CHECK: [[AND:%[0-9]+]]:_(s32) = G_AND [[COPY1]], [[C]]
+    ; CHECK: $w0 = COPY [[AND]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s1) = G_TRUNC %0(s32)
+    %2:_(s8) = G_ZEXT %1(s1)
+    %3:_(s32) = G_ANYEXT %2(s8)
+    $w0 = COPY %3(s32)
+
+...