[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.

Author: Jun Wang

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