GlobalISel/MachineIRBuilder: Construct DstOp with VRegAttrs (#113581)

Allow construction of DstOp with VRegAttrs.
Also allow construction with register class or bank and LLT.
Intended to be used in lowering code for reg-bank-select where
new registers need to have both register bank and LLT.
Add support for new type of DstOp in CSEMIRBuilder.

Author: petar-avramovic

llvm/include/llvm/CodeGen/GlobalISel/CSEInfo.h

@@ -17,6 +17,7 @@
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
 #include "llvm/CodeGen/GlobalISel/GISelWorkList.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CodeGen.h"
 
@@ -177,6 +178,8 @@ class GISelInstProfileBuilder {
   const GISelInstProfileBuilder &addNodeIDOpcode(unsigned Opc) const;
   const GISelInstProfileBuilder &addNodeIDRegType(const LLT Ty) const;
   const GISelInstProfileBuilder &addNodeIDRegType(const Register) const;
+  const GISelInstProfileBuilder &
+      addNodeIDRegType(MachineRegisterInfo::VRegAttrs) const;
 
   const GISelInstProfileBuilder &
   addNodeIDRegType(const TargetRegisterClass *RC) const;

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

@@ -72,15 +72,20 @@ class DstOp {
     LLT LLTTy;
     Register Reg;
     const TargetRegisterClass *RC;
+    MachineRegisterInfo::VRegAttrs Attrs;
   };
 
 public:
-  enum class DstType { Ty_LLT, Ty_Reg, Ty_RC };
+  enum class DstType { Ty_LLT, Ty_Reg, Ty_RC, Ty_VRegAttrs };
   DstOp(unsigned R) : Reg(R), Ty(DstType::Ty_Reg) {}
   DstOp(Register R) : Reg(R), Ty(DstType::Ty_Reg) {}
   DstOp(const MachineOperand &Op) : Reg(Op.getReg()), Ty(DstType::Ty_Reg) {}
   DstOp(const LLT T) : LLTTy(T), Ty(DstType::Ty_LLT) {}
   DstOp(const TargetRegisterClass *TRC) : RC(TRC), Ty(DstType::Ty_RC) {}
+  DstOp(MachineRegisterInfo::VRegAttrs Attrs)
+      : Attrs(Attrs), Ty(DstType::Ty_VRegAttrs) {}
+  DstOp(RegClassOrRegBank RCOrRB, LLT Ty)
+      : Attrs({RCOrRB, Ty}), Ty(DstType::Ty_VRegAttrs) {}
 
   void addDefToMIB(MachineRegisterInfo &MRI, MachineInstrBuilder &MIB) const {
     switch (Ty) {
@@ -93,6 +98,9 @@ class DstOp {
     case DstType::Ty_RC:
       MIB.addDef(MRI.createVirtualRegister(RC));
       break;
+    case DstType::Ty_VRegAttrs:
+      MIB.addDef(MRI.createVirtualRegister(Attrs));
+      break;
     }
   }
 
@@ -104,6 +112,8 @@ class DstOp {
       return LLTTy;
     case DstType::Ty_Reg:
       return MRI.getType(Reg);
+    case DstType::Ty_VRegAttrs:
+      return Attrs.Ty;
     }
     llvm_unreachable("Unrecognised DstOp::DstType enum");
   }
@@ -114,12 +124,13 @@ class DstOp {
   }
 
   const TargetRegisterClass *getRegClass() const {
-    switch (Ty) {
-    case DstType::Ty_RC:
-      return RC;
-    default:
-      llvm_unreachable("Not a RC Operand");
-    }
+    assert(Ty == DstType::Ty_RC && "Not a RC Operand");
+    return RC;
+  }
+
+  MachineRegisterInfo::VRegAttrs getVRegAttrs() const {
+    assert(Ty == DstType::Ty_VRegAttrs && "Not a VRegAttrs Operand");
+    return Attrs;
   }
 
   DstType getDstOpKind() const { return Ty; }

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

@@ -754,7 +754,7 @@ class MachineRegisterInfo {
   /// Returns register class or bank and low level type of \p Reg. Always safe
   /// to use. Special values are returned when \p Reg does not have some of the
   /// attributes.
-  VRegAttrs getVRegAttrs(Register Reg) {
+  VRegAttrs getVRegAttrs(Register Reg) const {
     return {getRegClassOrRegBank(Reg), getType(Reg)};
   }
 

llvm/lib/CodeGen/GlobalISel/CSEInfo.cpp

@@ -356,6 +356,20 @@ GISelInstProfileBuilder::addNodeIDRegType(const RegisterBank *RB) const {
   return *this;
 }
 
+const GISelInstProfileBuilder &GISelInstProfileBuilder::addNodeIDRegType(
+    MachineRegisterInfo::VRegAttrs Attrs) const {
+  addNodeIDRegType(Attrs.Ty);
+
+  const RegClassOrRegBank &RCOrRB = Attrs.RCOrRB;
+  if (RCOrRB) {
+    if (const auto *RB = dyn_cast_if_present<const RegisterBank *>(RCOrRB))
+      addNodeIDRegType(RB);
+    else
+      addNodeIDRegType(cast<const TargetRegisterClass *>(RCOrRB));
+  }
+  return *this;
+}
+
 const GISelInstProfileBuilder &
 GISelInstProfileBuilder::addNodeIDImmediate(int64_t Imm) const {
   ID.AddInteger(Imm);
@@ -389,17 +403,7 @@ GISelInstProfileBuilder::addNodeIDFlag(unsigned Flag) const {
 
 const GISelInstProfileBuilder &
 GISelInstProfileBuilder::addNodeIDReg(Register Reg) const {
-  LLT Ty = MRI.getType(Reg);
-  if (Ty.isValid())
-    addNodeIDRegType(Ty);
-
-  if (const RegClassOrRegBank &RCOrRB = MRI.getRegClassOrRegBank(Reg)) {
-    if (const auto *RB = dyn_cast_if_present<const RegisterBank *>(RCOrRB))
-      addNodeIDRegType(RB);
-    else if (const auto *RC =
-                 dyn_cast_if_present<const TargetRegisterClass *>(RCOrRB))
-      addNodeIDRegType(RC);
-  }
+  addNodeIDRegType(MRI.getVRegAttrs(Reg));
   return *this;
 }
 

llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp

@@ -73,18 +73,24 @@ bool CSEMIRBuilder::canPerformCSEForOpc(unsigned Opc) const {
 void CSEMIRBuilder::profileDstOp(const DstOp &Op,
                                  GISelInstProfileBuilder &B) const {
   switch (Op.getDstOpKind()) {
-  case DstOp::DstType::Ty_RC:
+  case DstOp::DstType::Ty_RC: {
     B.addNodeIDRegType(Op.getRegClass());
     break;
+  }
   case DstOp::DstType::Ty_Reg: {
     // Regs can have LLT&(RB|RC). If those exist, profile them as well.
     B.addNodeIDReg(Op.getReg());
     break;
   }
-  default:
+  case DstOp::DstType::Ty_LLT: {
     B.addNodeIDRegType(Op.getLLTTy(*getMRI()));
     break;
   }
+  case DstOp::DstType::Ty_VRegAttrs: {
+    B.addNodeIDRegType(Op.getVRegAttrs());
+    break;
+  }
+  }
 }
 
 void CSEMIRBuilder::profileSrcOp(const SrcOp &Op,

llvm/unittests/Target/AMDGPU/CMakeLists.txt

@@ -11,13 +11,15 @@ set(LLVM_LINK_COMPONENTS
   CodeGen
   CodeGenTypes
   Core
+  GlobalISel
   MC
   Support
   TargetParser
   )
 
 add_llvm_target_unittest(AMDGPUTests
   AMDGPUUnitTests.cpp
+  CSETest.cpp
   DwarfRegMappings.cpp
   ExecMayBeModifiedBeforeAnyUse.cpp
   PALMetadata.cpp

llvm/unittests/Target/AMDGPU/CSETest.cpp

@@ -0,0 +1,74 @@
+//===- llvm/unittests/Target/AMDGPU/CSETest.cpp ---------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUTargetMachine.h"
+#include "AMDGPUUnitTests.h"
+#include "llvm/CodeGen/GlobalISel/CSEInfo.h"
+#include "llvm/CodeGen/GlobalISel/CSEMIRBuilder.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+TEST(AMDGPU, TestCSEForRegisterClassOrBankAndLLT) {
+  auto TM = createAMDGPUTargetMachine("amdgcn-amd-", "gfx1100", "");
+  if (!TM)
+    GTEST_SKIP();
+
+  GCNSubtarget ST(TM->getTargetTriple(), std::string(TM->getTargetCPU()),
+                  std::string(TM->getTargetFeatureString()), *TM);
+
+  LLVMContext Ctx;
+  Module Mod("Module", Ctx);
+  Mod.setDataLayout(TM->createDataLayout());
+
+  auto *Type = FunctionType::get(Type::getVoidTy(Ctx), false);
+  auto *F = Function::Create(Type, GlobalValue::ExternalLinkage, "Test", &Mod);
+
+  MachineModuleInfo MMI(TM.get());
+  auto MF =
+      std::make_unique<MachineFunction>(*F, *TM, ST, MMI.getContext(), 42);
+  auto *BB = MF->CreateMachineBasicBlock();
+  MF->push_back(BB);
+
+  MachineIRBuilder B(*MF);
+  B.setMBB(*BB);
+
+  LLT S32{LLT::scalar(32)};
+  Register R0 = B.buildCopy(S32, Register(AMDGPU::SGPR0)).getReg(0);
+  Register R1 = B.buildCopy(S32, Register(AMDGPU::SGPR1)).getReg(0);
+
+  GISelCSEInfo CSEInfo;
+  CSEInfo.setCSEConfig(std::make_unique<CSEConfigFull>());
+  CSEInfo.analyze(*MF);
+  B.setCSEInfo(&CSEInfo);
+  CSEMIRBuilder CSEB(B.getState());
+  CSEB.setInsertPt(B.getMBB(), B.getInsertPt());
+
+  const RegisterBankInfo &RBI = *MF->getSubtarget().getRegBankInfo();
+
+  const TargetRegisterClass *SgprRC = &AMDGPU::SReg_32RegClass;
+  const RegisterBank *SgprRB = &RBI.getRegBank(AMDGPU::SGPRRegBankID);
+  MachineRegisterInfo::VRegAttrs SgprRCS32 = {SgprRC, S32};
+  MachineRegisterInfo::VRegAttrs SgprRBS32 = {SgprRB, S32};
+
+  auto Add = CSEB.buildAdd(S32, R0, R1);
+  auto AddRC = CSEB.buildInstr(AMDGPU::G_ADD, {SgprRCS32}, {R0, R1});
+  auto AddRB = CSEB.buildInstr(AMDGPU::G_ADD, {{SgprRB, S32}}, {R0, R1});
+
+  EXPECT_NE(Add, AddRC);
+  EXPECT_NE(Add, AddRB);
+  EXPECT_NE(AddRC, AddRB);
+
+  auto Add_CSE = CSEB.buildAdd(S32, R0, R1);
+  auto AddRC_CSE = CSEB.buildInstr(AMDGPU::G_ADD, {{SgprRC, S32}}, {R0, R1});
+  auto AddRB_CSE = CSEB.buildInstr(AMDGPU::G_ADD, {SgprRBS32}, {R0, R1});
+
+  EXPECT_EQ(Add, Add_CSE);
+  EXPECT_EQ(AddRC, AddRC_CSE);
+  EXPECT_EQ(AddRB, AddRB_CSE);
+}