[llvm] [AMDGPU] Set register bank for i1 arguments/return values (PR #96155)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Jun Wang (jwanggit86)

<details>
<summary>Changes</summary>

In planned work, the calling convention is to be updated such that i1 arguments and return values are assigned to SGPRs. For this change, we need to ensure the register banks are correctly assigned.

---
Full diff: https://github.com/llvm/llvm-project/pull/96155.diff


1 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp (+15) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp
index 0510a1d2eff88..ddf49f153d2c6 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp
@@ -3745,6 +3745,21 @@ AMDGPURegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
     if (!DstBank)
       DstBank = SrcBank;
 
+    // The calling convention is to be updated such that i1 function arguments
+    // or return values are assigned to SGPRs without promoting to i32. With
+    // this, for i1 function arguments, the call of getRegBank() above gives
+    // incorrect result. We set both src and dst banks to VCCRegBank.
+    if (!MI.getOperand(1).getReg().isVirtual() &&
+        MRI.getType(MI.getOperand(0).getReg()) == LLT::scalar(1)) {
+      DstBank = SrcBank = &AMDGPU::VCCRegBank;
+    }
+
+    // Similarly, for i1 return value, the dst reg is an SReg but we need to
+    // explicitly set the reg bank to VCCRegBank.
+    if (!MI.getOperand(0).getReg().isVirtual() &&
+        SrcBank == &AMDGPU::VCCRegBank)
+      DstBank = SrcBank;
+
     unsigned Size = getSizeInBits(MI.getOperand(0).getReg(), MRI, *TRI);
     if (MI.getOpcode() != AMDGPU::G_FREEZE &&
         cannotCopy(*DstBank, *SrcBank, TypeSize::getFixed(Size)))

``````````

</details>


https://github.com/llvm/llvm-project/pull/96155