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

[Jun Wang via llvm-commits]

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

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.

>From 860ddbec1125a5ae07adb4ae6a2e65afbe77254b Mon Sep 17 00:00:00 2001
From: Jun Wang <jun.wang7 at amd.com>
Date: Thu, 20 Jun 2024 04:48:14 -0500
Subject: [PATCH] [AMDGPU] Set register bank for i1 arguments/return values for
 planned calling convention update

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.
---
 llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

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)))