[AArch64][GlobalISel] Add support for narrowScalar of G_ZEXT

[aemerson]

Summary

We do this by merging the source with the high bits set to 0.

[arsenm]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

622 ↗

(On Diff #214955)

This should be able to handle cases where the result size is any multiple of the narrow size

[aemerson]

Add support for any multiple of the narrow type.

[paquette]

LGTM

[leonardchan]

Hi! After running a bisect, I believe this patch is causing the assertion error in https://bugs.llvm.org/show_bug.cgi?id=43171. Could you look into this? There is a reproducer attached to the bug.

llvm/trunk/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

case TargetOpcode::G_SEXT: {
     // Shift the sign bit of the low register through the high register.
     auto ShiftAmt =
         MIRBuilder.buildConstant(LLT::scalar(64), NarrowTy.getSizeInBits() - 1);
     auto Shift = MIRBuilder.buildAShr(NarrowTy, SrcReg, ShiftAmt);
     MIRBuilder.buildMerge(MI.getOperand(0).getReg(), {SrcReg, Shift.getReg(0)});
     MI.eraseFromParent();
     return Legalized;
   }
+  case TargetOpcode::G_ZEXT: {
+    if (TypeIdx != 0)
+      return UnableToLegalize;
+
+    if (SizeOp0 % NarrowTy.getSizeInBits() != 0)
+      return UnableToLegalize;
+
+    // Generate a merge where the bottom bits are taken from the source, and
+    // zero everything else.
+    Register ZeroReg = MIRBuilder.buildConstant(NarrowTy, 0).getReg(0);
+    unsigned NumParts = SizeOp0 / NarrowTy.getSizeInBits();
+    SmallVector<Register, 4> Srcs = {MI.getOperand(1).getReg()};
+    for (unsigned Part = 1; Part < NumParts; ++Part)
+      Srcs.push_back(ZeroReg);
+    MIRBuilder.buildMerge(MI.getOperand(0).getReg(), Srcs);
+    MI.eraseFromParent();
+    return Legalized;
+  }
 
   case TargetOpcode::G_ADD: {
     // FIXME: add support for when SizeOp0 isn't an exact multiple of
     // NarrowSize.
     if (SizeOp0 % NarrowSize != 0)
       return UnableToLegalize;
     // Expand in terms of carry-setting/consuming G_ADDE instructions.
     int NumParts = SizeOp0 / NarrowTy.getSizeInBits();

llvm/trunk/lib/Target/AArch64/AArch64LegalizerInfo.cpp

getActionDefinitionsBuilder(G_FCMP)
       .clampScalar(1, s32, s64)
       .widenScalarToNextPow2(1);
 
   // Extensions
   auto ExtLegalFunc = [=](const LegalityQuery &Query) {
     unsigned DstSize = Query.Types[0].getSizeInBits();
 
     if (DstSize == 128 && !Query.Types[0].isVector())
-      return false; // Extending to a scalar s128 is not legal.
+      return false; // Extending to a scalar s128 needs narrowing.
     
     // Make sure that we have something that will fit in a register, and
     // make sure it's a power of 2.
     if (DstSize < 8 || DstSize > 128 || !isPowerOf2_32(DstSize))
       return false;
 
     const LLT &SrcTy = Query.Types[1];
 
     // Special case for s1.
     if (SrcTy == s1)
       return true;
 
     // Make sure we fit in a register otherwise. Don't bother checking that
     // the source type is below 128 bits. We shouldn't be allowing anything
     // through which is wider than the destination in the first place.
     unsigned SrcSize = SrcTy.getSizeInBits();
     if (SrcSize < 8 || !isPowerOf2_32(SrcSize))
       return false;
 
     return true;
   };
-  getActionDefinitionsBuilder({G_ZEXT, G_ANYEXT}).legalIf(ExtLegalFunc);
-  getActionDefinitionsBuilder(G_SEXT)
+  getActionDefinitionsBuilder({G_ZEXT, G_SEXT, G_ANYEXT})
       .legalIf(ExtLegalFunc)
       .clampScalar(0, s64, s64); // Just for s128, others are handled above.
 
   getActionDefinitionsBuilder(G_TRUNC).alwaysLegal();
 
   getActionDefinitionsBuilder(G_SEXT_INREG).lower();
 
   // FP conversions

llvm/trunk/lib/Target/X86/X86RegisterBankInfo.cpp

 }
 
 const RegisterBank &X86RegisterBankInfo::getRegBankFromRegClass(
     const TargetRegisterClass &RC) const {
 
   if (X86::GR8RegClass.hasSubClassEq(&RC) ||
       X86::GR16RegClass.hasSubClassEq(&RC) ||
       X86::GR32RegClass.hasSubClassEq(&RC) ||
-      X86::GR64RegClass.hasSubClassEq(&RC))
+      X86::GR64RegClass.hasSubClassEq(&RC) ||
+      X86::LOW32_ADDR_ACCESSRegClass.hasSubClassEq(&RC) ||
+      X86::LOW32_ADDR_ACCESS_RBPRegClass.hasSubClassEq(&RC))
     return getRegBank(X86::GPRRegBankID);
 
   if (X86::FR32XRegClass.hasSubClassEq(&RC) ||
       X86::FR64XRegClass.hasSubClassEq(&RC) ||
       X86::VR128XRegClass.hasSubClassEq(&RC) ||
       X86::VR256XRegClass.hasSubClassEq(&RC) ||
       X86::VR512RegClass.hasSubClassEq(&RC))
     return getRegBank(X86::VECRRegBankID);

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

-# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -march=aarch64 -run-pass=legalizer -verify-machineinstrs %s -o - | FileCheck %s
----
-name:            narrow_s128
-tracksRegLiveness: true
-body:             |
-  bb.1:
-    liveins: $x0, $x1
-
-    ; CHECK-LABEL: name: narrow_s128
-    ; CHECK: liveins: $x0, $x1
-    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
-    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x1
-    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 63
-    ; CHECK: [[ASHR:%[0-9]+]]:_(s64) = G_ASHR [[COPY]], [[C]](s64)
-    ; CHECK: [[MV:%[0-9]+]]:_(s128) = G_MERGE_VALUES [[COPY]](s64), [[ASHR]](s64)
-    ; CHECK: G_STORE [[MV]](s128), [[COPY1]](p0) :: (store 16)
-    ; CHECK: RET_ReallyLR
-    %0:_(s64) = COPY $x0
-    %1:_(p0) = COPY $x1
-    %2:_(s128) = G_SEXT %0(s64)
-    G_STORE %2(s128), %1(p0) :: (store 16)
-    RET_ReallyLR
-
-...

llvm/trunk/test/CodeGen/AArch64/GlobalISel/legalize-sext-zext-128.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=aarch64 -run-pass=legalizer -verify-machineinstrs %s -o - | FileCheck %s
+---
+name:            narrow_sext_s128
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: narrow_sext_s128
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x1
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 63
+    ; CHECK: [[ASHR:%[0-9]+]]:_(s64) = G_ASHR [[COPY]], [[C]](s64)
+    ; CHECK: [[MV:%[0-9]+]]:_(s128) = G_MERGE_VALUES [[COPY]](s64), [[ASHR]](s64)
+    ; CHECK: G_STORE [[MV]](s128), [[COPY1]](p0) :: (store 16)
+    ; CHECK: RET_ReallyLR
+    %0:_(s64) = COPY $x0
+    %1:_(p0) = COPY $x1
+    %2:_(s128) = G_SEXT %0(s64)
+    G_STORE %2(s128), %1(p0) :: (store 16)
+    RET_ReallyLR
+
+...
+---
+name:            narrow_zext_s128
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: narrow_zext_s128
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x1
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+    ; CHECK: [[MV:%[0-9]+]]:_(s128) = G_MERGE_VALUES [[COPY]](s64), [[C]](s64)
+    ; CHECK: G_STORE [[MV]](s128), [[COPY1]](p0) :: (store 16)
+    ; CHECK: RET_ReallyLR
+    %0:_(s64) = COPY $x0
+    %1:_(p0) = COPY $x1
+    %2:_(s128) = G_ZEXT %0(s64)
+    G_STORE %2(s128), %1(p0) :: (store 16)
+    RET_ReallyLR
+
+...
+---
+name:            narrow_zext_s192
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: narrow_zext_s192
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x1
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+    ; CHECK: [[MV:%[0-9]+]]:_(s192) = G_MERGE_VALUES [[COPY]](s64), [[C]](s64), [[C]](s64)
+    ; CHECK: [[UV:%[0-9]+]]:_(s64), [[UV1:%[0-9]+]]:_(s64), [[UV2:%[0-9]+]]:_(s64) = G_UNMERGE_VALUES [[MV]](s192)
+    ; CHECK: G_STORE [[UV]](s64), [[COPY1]](p0) :: (store 8)
+    ; CHECK: RET_ReallyLR
+    %0:_(s64) = COPY $x0
+    %1:_(p0) = COPY $x1
+    %2:_(s192) = G_ZEXT %0(s64)
+    %3:_(s64), %4:_(s64), %5:_(s64) = G_UNMERGE_VALUES %2(s192)
+    G_STORE %3, %1(p0) :: (store 8)
+    RET_ReallyLR
+
+...

llvm/trunk/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir

 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
 # DEBUG-NEXT: G_VASTART (opcode {{[0-9]+}}): 1 type index, 0 imm indices
 # DEBUG-NEXT: .. the first uncovered type index: 1, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
 # DEBUG-NEXT: G_VAARG (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. the first uncovered type index: 2, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
 # DEBUG-NEXT: G_SEXT (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
+# DEBUG-NEXT: .. opcode {{[0-9]+}} is aliased to {{[0-9]+}}
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_SEXT_INREG (opcode {{[0-9]+}}): 1 type index, 1 imm index
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_ZEXT (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected