[llvm] AMDGPU/GlobalISel: RegBankLegalize rules for gfx90a/gfx942 MFMAs (PR #194076)

[Petar Avramovic via llvm-commits]

================
@@ -2031,21 +2032,64 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
       .Uni(S16, {{UniInVgprS16}, {IntrId, Vgpr32}})
       .Div(S16, {{Vgpr16}, {IntrId, Vgpr32}});
 
-  // TODO: Add handling for GFX90A+ which should use VGPRs instead of AGPRs.
+  // Predicates to check if we want to use the VGPR form or the AGPR form
+  // for an MFMA / SMFMAC intrinsic on gfx90a+.
+  Predicate needAGPR([](const MachineInstr &MI) -> bool {
+    const MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
+    unsigned DstSize = MRI.getType(MI.getOperand(0).getReg()).getSizeInBits();
+    unsigned MinNumRegsRequired = DstSize / 32;
+    const SIMachineFunctionInfo *Info =
+        MI.getMF()->getInfo<SIMachineFunctionInfo>();
+    return Info->selectAGPRFormMFMA(MinNumRegsRequired);
+  });
+  Predicate needVGPR = !needAGPR;
+
   bool HasGFX90AInsts = ST->hasGFX90AInsts();
-  addRulesForIOpcs({amdgcn_mfma_f32_32x32x1f32,  amdgcn_mfma_f32_16x16x1f32,
-                    amdgcn_mfma_f32_4x4x1f32,    amdgcn_mfma_f32_32x32x2f32,
-                    amdgcn_mfma_f32_16x16x4f32,  amdgcn_mfma_f32_32x32x4f16,
-                    amdgcn_mfma_f32_16x16x4f16,  amdgcn_mfma_f32_4x4x4f16,
-                    amdgcn_mfma_f32_32x32x8f16,  amdgcn_mfma_f32_16x16x16f16,
-                    amdgcn_mfma_i32_32x32x4i8,   amdgcn_mfma_i32_16x16x4i8,
-                    amdgcn_mfma_i32_4x4x4i8,     amdgcn_mfma_i32_32x32x8i8,
-                    amdgcn_mfma_i32_16x16x16i8,  amdgcn_mfma_f32_32x32x2bf16,
-                    amdgcn_mfma_f32_16x16x2bf16, amdgcn_mfma_f32_4x4x2bf16,
-                    amdgcn_mfma_f32_32x32x4bf16, amdgcn_mfma_f32_16x16x8bf16})
+  addRulesForIOpcs(
+      {amdgcn_mfma_f32_32x32x1f32,       amdgcn_mfma_f32_16x16x1f32,
+       amdgcn_mfma_f32_4x4x1f32,         amdgcn_mfma_f32_32x32x2f32,
+       amdgcn_mfma_f32_16x16x4f32,       amdgcn_mfma_f32_32x32x4f16,
+       amdgcn_mfma_f32_16x16x4f16,       amdgcn_mfma_f32_4x4x4f16,
+       amdgcn_mfma_f32_32x32x8f16,       amdgcn_mfma_f32_16x16x16f16,
+       amdgcn_mfma_i32_32x32x4i8,        amdgcn_mfma_i32_16x16x4i8,
+       amdgcn_mfma_i32_4x4x4i8,          amdgcn_mfma_i32_32x32x8i8,
+       amdgcn_mfma_i32_16x16x16i8,       amdgcn_mfma_f32_32x32x2bf16,
+       amdgcn_mfma_f32_16x16x2bf16,      amdgcn_mfma_f32_4x4x2bf16,
+       amdgcn_mfma_f32_32x32x4bf16,      amdgcn_mfma_f32_16x16x8bf16,
+       amdgcn_mfma_f32_32x32x4bf16_1k,   amdgcn_mfma_f32_16x16x4bf16_1k,
+       amdgcn_mfma_f32_4x4x4bf16_1k,     amdgcn_mfma_f32_32x32x8bf16_1k,
+       amdgcn_mfma_f32_16x16x16bf16_1k,  amdgcn_mfma_f64_16x16x4f64,
+       amdgcn_mfma_f64_4x4x4f64,         amdgcn_mfma_i32_16x16x32_i8,
+       amdgcn_mfma_i32_32x32x16_i8,      amdgcn_mfma_f32_16x16x8_xf32,
+       amdgcn_mfma_f32_32x32x4_xf32,     amdgcn_mfma_f32_16x16x32_bf8_bf8,
+       amdgcn_mfma_f32_16x16x32_bf8_fp8, amdgcn_mfma_f32_16x16x32_fp8_bf8,
+       amdgcn_mfma_f32_16x16x32_fp8_fp8, amdgcn_mfma_f32_32x32x16_bf8_bf8,
+       amdgcn_mfma_f32_32x32x16_bf8_fp8, amdgcn_mfma_f32_32x32x16_fp8_bf8,
+       amdgcn_mfma_f32_32x32x16_fp8_fp8})
----------------
petar-avramovic wrote:

I would prefer to split these opcodes into another group since they are not available when !HasGFX90AInsts

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