[GlobalISel][Utils] Teach getConstantVRegVal how to look through trunc and z|sext

[qcolombet]

Summary

getConstantVRegVal used to only look for G_CONSTANT when looking at
unboxing the value of a vreg. However, constants are sometimes not
directly used and are hidden behind trunc or s|zext chain of computation.

In particular this may be introduced by the legalization process that
doesn't want simplify these patterns because it can lead to infine loop
when legalizing a constant.

[arsenm]

It seems slightly odd to me for this to look through this. Why wouldn’t the combiner fold casts of constants? Do any of these testcases hit the legalization loop problem you mentioned?

[qcolombet]

Hi Matt,

Why wouldn’t the combiner fold casts of constants?

The problem is, let say you have G_CONSTANT being legal only for {s16, s32}
Now consider:

%0 = G_CONSTANT i14

The legalizer will then transform that into:

%tmp = G_CONSTANT i16
%0 = G_TRUNC i16 %tmp to i14

If you were to do the constant folding within the MachineIRBuilder, this would then be transformed back into:

%0 = G_CONSTANT i14

And now you have your loop.

We could do that in the combiner and then we don't need to look through this, However, in the middle of legalization, we will still have the truncates and extensions that may not be cleaned up at all given point in the legalization and that's when we would miss to get the constant value.

Do any of these testcases hit the legalization loop problem you mentioned?

No, they just happen to exercise the code I added.

The problem I have arises during legalization with an out-of-tree target.
Essentially, we promote some illegal operation, then end-up with extensions in the way of the constants and some custom legalization fails (without this patch) because we don't support non-constant patterns.

[volkan]

LGTM.

[qcolombet]

I'm going to hold on that one for a little longer.
I wonder how SDAG handle constant during legalization, because I was not expecting that we would not issue the constant for the new type period.
I suspect it works because constants are just not materialized via instructions and maybe we should just ignore them during legalization. (If a constant is illegal, at the end of the legalization process, it should have any users, right?)

[qcolombet]

• Introduce a different variant of getConstantVRegVal instead of overriding the existing one (after talking to @aditya_nandakumar we decided we may want to have more control on what happens when we call this function).
• Add a few comments.

Longer term, I believe we don't want to legalize G_CONSTANT, but for this to happen we need to make sure we combine in the builder the extensions, otherwise we will have them hanging around and causing problem.

[arsenm]

In D59227#1426992, @qcolombet wrote:

• Introduce a different variant of getConstantVRegVal instead of overriding the existing one (after talking to @aditya_nandakumar we decided we may want to have more control on what happens when we call this function).
• Add a few comments.

Longer term, I believe we don't want to legalize G_CONSTANT, but for this to happen we need to make sure we combine in the builder the extensions, otherwise we will have them hanging around and causing problem.

Why wouldn't G_CONSTANT be legalized? This doesn't make any sense to me

[qcolombet]

Why wouldn't G_CONSTANT be legalized? This doesn't make any sense to me

Because they are just place holder.
When doing selection, only the ones that won't be folded into immediate should be materialized and at this point, we just have to materialize load immediate into valid registers.

The only reason the constant show up as instructions right now is because we wanted to handle machine operand uniformly via vreg. Otherwise, they would just be immediate operands. If you imagine that setting, I guess it is more obvious why we don't legalize them.

Does it make sense?

[qcolombet]

Update a mips test that I forgot.

[arsenm]

In D59227#1427643, @qcolombet wrote:

Why wouldn't G_CONSTANT be legalized? This doesn't make any sense to me

Because they are just place holder.
When doing selection, only the ones that won't be folded into immediate should be materialized and at this point, we just have to materialize load immediate into valid registers.

The only reason the constant show up as instructions right now is because we wanted to handle machine operand uniformly via vreg. Otherwise, they would just be immediate operands. If you imagine that setting, I guess it is more obvious why we don't legalize them.

Does it make sense?

This isn't how I think of G_CONSTANT. I think of it as a mov setting a legal result register that is a legitimate instruction in its own right and not just a placeholder. The immediate folding rules for AMDGPU are somewhat complex, so an arbitrary constant is always going to be selected as a materializing move. We wouldn't want G_* instructions to always fold constant operands since that would create more work. A directly selected instruction would be wrong by default. We have a later target pass to fold these based on the selected instruction rules (which can also account for the number of uses of materialized constants etc.). For operands that truly need to be constants that shouldn't be legalized, I don't think G_CONSTANT should be used. This is part of the motivation behind D58232. In SelectionDAG there is the TargetConstant vs. Constant distinction, but I think that makes less sense here, so weird constants should just never be materialized.

[qcolombet]

Disclaimer: this is just my opinion and in particular I haven't talked with anybody about constants being place holders. So just see my comment as what they are: an individual throwing ideas :).

This isn't how I think of G_CONSTANT. I think of it as a mov setting a legal result register that is a legitimate instruction in its own right and not just a placeholder. The immediate folding rules for AMDGPU are somewhat complex, so an arbitrary constant is always going to be selected as a materializing move. We wouldn't want G_* instructions to always fold constant operands since that would create more work.

I agree that the G_CONSTANT are a mov to a legal result register. However, I don't think they should be legalized on there own. Essentially, if we truncate/extend the constants directly during the legalization of their uses, they naturally become legal (the illegal ones become dead). I.e., we don't have to handle G_CONSTANT explicitly. This is true only if the builders don't create thing like ext(cst), like they do today, because then we indeed have to legalize those immediate moves.

To go back to constant selecting to thing that could be expansive, the fact that we decide to model them as instructions was also to account for that. However, the way we effectively deal with them is not what I had in mind:

• A given constant should appear only once (this is workaround with the CSE builder right now, but that's not how I envisioned thing)
• Constant should be selected last: at this point we would see if they have any users and thus if they need to be materialized into a register

The second point was supposed to come from free as in my mind constants would have been uniquely represented and placed in the entry block, which is proceeded last (i.e., all users have been seen).

Anyway, I am not planning on changing that anytime soon (if ever at all) and because of that this patch is unfortunately required. Bottom line, this patch is a workaround a more deeper problem in the way we handle constants IMHO.

[qcolombet]

For operands that truly need to be constants that shouldn't be legalized, I don't think G_CONSTANT should be used. This is part of the motivation behind D58232 https://reviews.llvm.org/D58232. In SelectionDAG there is the TargetConstant vs. Constant distinction, but I think that makes less sense here, so weird constants should just never be materialized.

If we were to go down the road of allowing constants as MachineOperand, I would rather kill the G_CONSTANT and be smart about how we materialize them (unlike selection DAG). Now, the code churn to deal with that would be so bad (teaching every place that a MachineOperand may not be a vreg) that I don't think it makes sense.
I am somewhat sensible for the intrinsic argument since the generic passes don't look at them.

[arsenm]

In D59227#1427718, @qcolombet wrote:

If we were to go down the road of allowing constants as MachineOperand, I would rather kill the G_CONSTANT and be smart about how we materialize them (unlike selection DAG). Now, the code churn to deal with that would be so bad (teaching every place that a MachineOperand may not be a vreg) that I don't think it makes sense.
I am somewhat sensible for the intrinsic argument since the generic passes don't look at them.

Right, I just want this for intrinsic arguments, and G_CONSTANT for anything else. Code looking at arbitrary instructions already need to worry about the possibility of some operand not being a register for G_PHI and DBG_VALUE anyway.

[arsenm]

In D59227#1427702, @qcolombet wrote:

I agree that the G_CONSTANT are a mov to a legal result register. However, I don't think they should be legalized on there own. Essentially, if we truncate/extend the constants directly during the legalization of their uses, they naturally become legal (the illegal ones become dead). I.e., we don't have to handle G_CONSTANT explicitly. This is true only if the builders don't create thing like ext(cst), like they do today, because then we indeed have to legalize those immediate moves.

I'm not sure why we would want special rules for constant handling. I think it would be cleaner to just keep them treated as any other normal instruction. I'm assuming you mean the legalizer by builder. If you're going to try to avoid the intermediate legalization artifacts by looking at the uses, I don't see why you wouldn't change the legalizer worklist to do this for every instruction.

[qcolombet]

Prototyped a retry legalization process in D59339.
The problem is that this patch is still needed to access the constant, because the legalization artifacts combiner inserts copies instead of updating the users.

Bottom line, @arsenm may I land this, while we figure out what we want to do for the broader problem?

[arsenm]

LGTM

[qcolombet]

Landed in r356116.

include/llvm/CodeGen/GlobalISel/Utils.h

void reportGISelFailure(MachineFunction &MF, const TargetPassConfig &TPC,
                         MachineOptimizationRemarkEmitter &MORE,
                         MachineOptimizationRemarkMissed &R);
 
 void reportGISelFailure(MachineFunction &MF, const TargetPassConfig &TPC,
                         MachineOptimizationRemarkEmitter &MORE,
                         const char *PassName, StringRef Msg,
                         const MachineInstr &MI);
 
+/// If \p VReg is defined by a G_CONSTANT fits in int64_t
+/// returns it.
 Optional<int64_t> getConstantVRegVal(unsigned VReg,
                                      const MachineRegisterInfo &MRI);
+/// Simple struct used to hold a constant integer value and a virtual
+/// register.
+struct ValueAndVReg {
+  int64_t Value;
+  unsigned VReg;
+};
+/// If \p VReg is defined by a statically evaluable chain of
+/// instructions rooted on a G_CONSTANT (\p LookThroughInstrs == true)
+/// and that constant fits in int64_t, returns its value as well as
+/// the virtual register defined by this G_CONSTANT.
+/// When \p LookThroughInstrs == false, this function behaves like
+/// getConstantVRegVal.
+Optional<ValueAndVReg>
+getConstantVRegValWithLookThrough(unsigned VReg, const MachineRegisterInfo &MRI,
+                                  bool LookThroughInstrs = true);
 const ConstantFP* getConstantFPVRegVal(unsigned VReg,
                                        const MachineRegisterInfo &MRI);
 
 /// See if Reg is defined by an single def instruction that is
 /// Opcode. Also try to do trivial folding if it's a COPY with
 /// same types. Returns null otherwise.
 MachineInstr *getOpcodeDef(unsigned Opcode, unsigned Reg,
                            const MachineRegisterInfo &MRI);

lib/CodeGen/GlobalISel/InstructionSelector.cpp

bool InstructionSelector::constrainOperandRegToRegClass(
   return
       constrainRegToClass(MRI, TII, RBI, I, I.getOperand(OpIdx).getReg(), RC);
 }
 
 bool InstructionSelector::isOperandImmEqual(
     const MachineOperand &MO, int64_t Value,
     const MachineRegisterInfo &MRI) const {
   if (MO.isReg() && MO.getReg())
-    if (auto VRegVal = getConstantVRegVal(MO.getReg(), MRI))
-      return *VRegVal == Value;
+    if (auto VRegVal = getConstantVRegValWithLookThrough(MO.getReg(), MRI))
+      return VRegVal->Value == Value;
   return false;
 }
 
 bool InstructionSelector::isBaseWithConstantOffset(
     const MachineOperand &Root, const MachineRegisterInfo &MRI) const {
   if (!Root.isReg())
     return false;
 

lib/CodeGen/GlobalISel/Utils.cpp

void llvm::reportGISelFailure(MachineFunction &MF, const TargetPassConfig &TPC,
   // Printing MI is expensive;  only do it if expensive remarks are enabled.
   if (TPC.isGlobalISelAbortEnabled() || MORE.allowExtraAnalysis(PassName))
     R << ": " << ore::MNV("Inst", MI);
   reportGISelFailure(MF, TPC, MORE, R);
 }
 
 Optional<int64_t> llvm::getConstantVRegVal(unsigned VReg,
                                            const MachineRegisterInfo &MRI) {
-  MachineInstr *MI = MRI.getVRegDef(VReg);
-  if (MI->getOpcode() != TargetOpcode::G_CONSTANT)
+  Optional<ValueAndVReg> ValAndVReg =
+      getConstantVRegValWithLookThrough(VReg, MRI, /*LookThroughInstrs*/ false);
+  assert((!ValAndVReg || ValAndVReg->VReg == VReg) &&
+         "Value found while looking through instrs");
+  if (!ValAndVReg)
     return None;
+  return ValAndVReg->Value;
+}
 
-  if (MI->getOperand(1).isImm())
-    return MI->getOperand(1).getImm();
+Optional<ValueAndVReg> llvm::getConstantVRegValWithLookThrough(
+    unsigned VReg, const MachineRegisterInfo &MRI, bool LookThroughInstrs) {
+  SmallVector<std::pair<unsigned, unsigned>, 4> SeenOpcodes;
+  MachineInstr *MI;
+  while ((MI = MRI.getVRegDef(VReg)) &&
+         MI->getOpcode() != TargetOpcode::G_CONSTANT && LookThroughInstrs) {
+    switch (MI->getOpcode()) {
+    case TargetOpcode::G_TRUNC:
+    case TargetOpcode::G_SEXT:
+    case TargetOpcode::G_ZEXT:
+      SeenOpcodes.push_back(std::make_pair(
+          MI->getOpcode(),
+          MRI.getType(MI->getOperand(0).getReg()).getSizeInBits()));
+      VReg = MI->getOperand(1).getReg();
+      break;
+    case TargetOpcode::COPY:
+      VReg = MI->getOperand(1).getReg();
+      if (TargetRegisterInfo::isPhysicalRegister(VReg))
+        return None;
+      break;
+    default:
+      return None;
+    }
+  }
+  if (!MI || MI->getOpcode() != TargetOpcode::G_CONSTANT ||
+      (!MI->getOperand(1).isImm() && !MI->getOperand(1).isCImm()))
+    return None;
 
-  if (MI->getOperand(1).isCImm() &&
-      MI->getOperand(1).getCImm()->getBitWidth() <= 64)
-    return MI->getOperand(1).getCImm()->getSExtValue();
+  const MachineOperand &CstVal = MI->getOperand(1);
+  unsigned BitWidth = MRI.getType(MI->getOperand(0).getReg()).getSizeInBits();
+  APInt Val = CstVal.isImm() ? APInt(BitWidth, CstVal.getImm())
+                             : CstVal.getCImm()->getValue();
+  assert(Val.getBitWidth() == BitWidth &&
+         "Value bitwidth doesn't match definition type");
+  while (!SeenOpcodes.empty()) {
+    std::pair<unsigned, unsigned> OpcodeAndSize = SeenOpcodes.pop_back_val();
+    switch (OpcodeAndSize.first) {
+    case TargetOpcode::G_TRUNC:
+      Val = Val.trunc(OpcodeAndSize.second);
+      break;
+    case TargetOpcode::G_SEXT:
+      Val = Val.sext(OpcodeAndSize.second);
+      break;
+    case TargetOpcode::G_ZEXT:
+      Val = Val.zext(OpcodeAndSize.second);
+      break;
+    }
+  }
 
+  if (Val.getBitWidth() > 64)
     return None;
+
+  return ValueAndVReg{Val.getSExtValue(), VReg};
 }
 
 const llvm::ConstantFP* llvm::getConstantFPVRegVal(unsigned VReg,
                                        const MachineRegisterInfo &MRI) {
   MachineInstr *MI = MRI.getVRegDef(VReg);
   if (TargetOpcode::G_FCONSTANT != MI->getOpcode())
     return nullptr;
   return MI->getOperand(1).getFPImm();

test/CodeGen/Mips/GlobalISel/llvm-ir/select.ll

 entry:
   %cond = select i1 %test, i32* %a, i32* %b
   ret i32* %cond
 }
 
 define i32 @select_with_negation(i32 %a, i32 %b, i32 %x, i32 %y) {
 ; MIPS32-LABEL: select_with_negation:
 ; MIPS32:       # %bb.0: # %entry
-; MIPS32-NEXT:    lui $1, 65535
-; MIPS32-NEXT:    ori $1, $1, 65535
 ; MIPS32-NEXT:    slt $4, $4, $5
-; MIPS32-NEXT:    xor $1, $4, $1
-; MIPS32-NEXT:    lui $4, 0
-; MIPS32-NEXT:    ori $4, $4, 1
-; MIPS32-NEXT:    and $1, $1, $4
-; MIPS32-NEXT:    movn $7, $6, $1
+; MIPS32-NEXT:    not $4, $4
+; MIPS32-NEXT:    lui $5, 0
+; MIPS32-NEXT:    ori $5, $5, 1
+; MIPS32-NEXT:    and $4, $4, $5
+; MIPS32-NEXT:    movn $7, $6, $4
 ; MIPS32-NEXT:    move $2, $7
 ; MIPS32-NEXT:    jr $ra
 ; MIPS32-NEXT:    nop
 entry:
   %cmp = icmp slt i32 %a, %b
   %lneg = xor i1 %cmp, true
   %cond = select i1 %lneg, i32 %x, i32 %y
   ret i32 %cond
 }

test/CodeGen/X86/GlobalISel/ashr-scalar.ll

; X64-NEXT:    retq
   %res = ashr i64 %arg1, 5
   ret i64 %res
 }
 
 define i64 @test_ashr_i64_imm1(i64 %arg1) {
 ; X64-LABEL: test_ashr_i64_imm1:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movq %rdi, %rax
-; X64-NEXT:    movq $1, %rcx
-; X64-NEXT:    sarq %cl, %rax
+; X64-NEXT:    sarq %rax
 ; X64-NEXT:    retq
   %res = ashr i64 %arg1, 1
   ret i64 %res
 }
 
 define i32 @test_ashr_i32(i32 %arg1, i32 %arg2) {
 ; X64-LABEL: test_ashr_i32:
 ; X64:       # %bb.0:
; X64-NEXT:    retq
   %res = ashr i32 %arg1, 5
   ret i32 %res
 }
 
 define i32 @test_ashr_i32_imm1(i32 %arg1) {
 ; X64-LABEL: test_ashr_i32_imm1:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    movl $1, %ecx
-; X64-NEXT:    sarl %cl, %eax
+; X64-NEXT:    sarl %eax
 ; X64-NEXT:    retq
   %res = ashr i32 %arg1, 1
   ret i32 %res
 }
 
 define i16 @test_ashr_i16(i32 %arg1, i32 %arg2) {
 ; X64-LABEL: test_ashr_i16:
 ; X64:       # %bb.0:
; X64-NEXT:    retq
   %res = ashr i16 %a, 5
   ret i16 %res
 }
 
 define i16 @test_ashr_i16_imm1(i32 %arg1) {
 ; X64-LABEL: test_ashr_i16_imm1:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    movw $1, %cx
-; X64-NEXT:    sarw %cl, %ax
+; X64-NEXT:    sarw %ax
 ; X64-NEXT:    # kill: def $ax killed $ax killed $eax
 ; X64-NEXT:    retq
   %a = trunc i32 %arg1 to i16
   %res = ashr i16 %a, 1
   ret i16 %res
 }
 
 define i8 @test_ashr_i8(i32 %arg1, i32 %arg2) {

test/CodeGen/X86/GlobalISel/lshr-scalar.ll

; X64-NEXT:    retq
   %res = lshr i64 %arg1, 5
   ret i64 %res
 }
 
 define i64 @test_lshr_i64_imm1(i64 %arg1) {
 ; X64-LABEL: test_lshr_i64_imm1:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movq %rdi, %rax
-; X64-NEXT:    movq $1, %rcx
-; X64-NEXT:    shrq %cl, %rax
+; X64-NEXT:    shrq %rax
 ; X64-NEXT:    retq
   %res = lshr i64 %arg1, 1
   ret i64 %res
 }
 
 define i32 @test_lshr_i32(i32 %arg1, i32 %arg2) {
 ; X64-LABEL: test_lshr_i32:
 ; X64:       # %bb.0:
; X64-NEXT:    retq
   %res = lshr i32 %arg1, 5
   ret i32 %res
 }
 
 define i32 @test_lshr_i32_imm1(i32 %arg1) {
 ; X64-LABEL: test_lshr_i32_imm1:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    movl $1, %ecx
-; X64-NEXT:    shrl %cl, %eax
+; X64-NEXT:    shrl %eax
 ; X64-NEXT:    retq
   %res = lshr i32 %arg1, 1
   ret i32 %res
 }
 
 define i16 @test_lshr_i16(i32 %arg1, i32 %arg2) {
 ; X64-LABEL: test_lshr_i16:
 ; X64:       # %bb.0:
; X64-NEXT:    retq
   %res = lshr i16 %a, 5
   ret i16 %res
 }
 
 define i16 @test_lshr_i16_imm1(i32 %arg1) {
 ; X64-LABEL: test_lshr_i16_imm1:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    movw $1, %cx
-; X64-NEXT:    shrw %cl, %ax
+; X64-NEXT:    shrw %ax
 ; X64-NEXT:    # kill: def $ax killed $ax killed $eax
 ; X64-NEXT:    retq
   %a = trunc i32 %arg1 to i16
   %res = lshr i16 %a, 1
   ret i16 %res
 }
 
 define i8 @test_lshr_i8(i32 %arg1, i32 %arg2) {

test/CodeGen/X86/GlobalISel/shl-scalar.ll

 ; X64-NEXT:    retq
   %res = shl i64 %arg1, 5
   ret i64 %res
 }
 
 define i64 @test_shl_i64_imm1(i64 %arg1) {
 ; X64-LABEL: test_shl_i64_imm1:
 ; X64:       # %bb.0:
-; X64-NEXT:    movq %rdi, %rax
-; X64-NEXT:    movq $1, %rcx
-; X64-NEXT:    shlq %cl, %rax
+; X64-NEXT:    leaq (%rdi,%rdi), %rax
 ; X64-NEXT:    retq
   %res = shl i64 %arg1, 1
   ret i64 %res
 }
 
 define i32 @test_shl_i32(i32 %arg1, i32 %arg2) {
 ; X64-LABEL: test_shl_i32:
 ; X64:       # %bb.0:

 ; X64-NEXT:    retq
   %res = shl i32 %arg1, 5
   ret i32 %res
 }
 
 define i32 @test_shl_i32_imm1(i32 %arg1) {
 ; X64-LABEL: test_shl_i32_imm1:
 ; X64:       # %bb.0:
-; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    movl $1, %ecx
-; X64-NEXT:    shll %cl, %eax
+; X64-NEXT:    # kill: def $edi killed $edi def $rdi
+; X64-NEXT:    leal (%rdi,%rdi), %eax
 ; X64-NEXT:    retq
   %res = shl i32 %arg1, 1
   ret i32 %res
 }
 
 define i16 @test_shl_i16(i32 %arg1, i32 %arg2) {
 ; X64-LABEL: test_shl_i16:
 ; X64:       # %bb.0:
; X64-NEXT:    retq
   %a = trunc i32 %arg1 to i16
   %res = shl i16 %a, 5
   ret i16 %res
 }
 
 define i16 @test_shl_i16_imm1(i32 %arg1) {
 ; X64-LABEL: test_shl_i16_imm1:
 ; X64:       # %bb.0:
-; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    movw $1, %cx
-; X64-NEXT:    shlw %cl, %ax
+; X64-NEXT:    # kill: def $edi killed $edi def $rdi
+; X64-NEXT:    leal (%rdi,%rdi), %eax
 ; X64-NEXT:    # kill: def $ax killed $ax killed $eax
 ; X64-NEXT:    retq
   %a = trunc i32 %arg1 to i16
   %res = shl i16 %a, 1
   ret i16 %res
 }
 
 define i8 @test_shl_i8(i32 %arg1, i32 %arg2) {