[AMDGPU] Provide control over AGPR/VGPR MFMA form

Change-Id: Ife390264aef869c61a25d032f47301fab5554b4f

Author: jrbyrnes

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

@@ -29,6 +29,24 @@ enum { MAX_LANES = 64 };
 
 using namespace llvm;
 
+namespace {
+enum MFMARegClass {
+  Unspecified,
+  VGPR,
+  AGPR,
+};
+}
+
+cl::opt<MFMARegClass>
+    MFMAForm("amdgpu-mfma-form", cl::Hidden,
+             cl::desc("Register class to use for Opc and Dest of MFMA. If "
+                      "unspecified, default to compiler heuristics"),
+             cl::init(MFMARegClass::Unspecified),
+             cl::values(clEnumValN(MFMARegClass::VGPR, "vgpr",
+                                   "Use the VGPR MFMA form."),
+                        clEnumValN(MFMARegClass::AGPR, "agpr",
+                                   "Use the VGPR MFMA form.")));
+
 const GCNTargetMachine &getTM(const GCNSubtarget *STI) {
   const SITargetLowering *TLI = STI->getTargetLowering();
   return static_cast<const GCNTargetMachine &>(TLI->getTargetMachine());
@@ -70,11 +88,14 @@ SIMachineFunctionInfo::SIMachineFunctionInfo(const Function &F,
   }
 
   MayNeedAGPRs = ST.hasMAIInsts();
-  if (ST.hasGFX90AInsts() &&
+  if (MFMAForm == MFMARegClass::Unspecified && ST.hasGFX90AInsts() &&
       ST.getMaxNumVGPRs(F) <= AMDGPU::VGPR_32RegClass.getNumRegs() &&
       !mayUseAGPRs(F))
     MayNeedAGPRs = false; // We will select all MAI with VGPR operands.
 
+  else if (MFMAForm != MFMARegClass::Unspecified)
+    MayNeedAGPRs = MFMAForm == MFMARegClass::AGPR;
+
   if (AMDGPU::isChainCC(CC)) {
     // Chain functions don't receive an SP from their caller, but are free to
     // set one up. For now, we can use s32 to match what amdgpu_gfx functions