[CIR][CUDA] Add NVPTX ABI info skeleton

llvm · Feb 1, 2025 · 6bbe08c · 6bbe08c
1 parent c4e5842
commit 6bbe08c
Showing 1 changed file with 50 additions and 0 deletions.
diff --git a/clang/lib/CIR/CodeGen/TargetInfo.cpp b/clang/lib/CIR/CodeGen/TargetInfo.cpp
@@ -305,6 +305,33 @@ class SPIRVTargetCIRGenInfo : public CommonSPIRTargetCIRGenInfo {
 
 } // namespace
 
+//===----------------------------------------------------------------------===//
+// NVPTX ABI Implementation
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+class NVPTXABIInfo : public ABIInfo {
+  mlir::MLIRContext &MlCtx;
+
+public:
+  NVPTXABIInfo(CIRGenTypes &CGT, mlir::MLIRContext &MlCtx)
+      : ABIInfo(CGT), MlCtx(MlCtx) {}
+
+  cir::ABIArgInfo classifyReturnType(QualType RetTy) const;
+  cir::ABIArgInfo classifyArgumentType(QualType Ty) const;
+
+  void computeInfo(CIRGenFunctionInfo &FnInfo) const override;
+};
+
+class NVPTXTargetCIRGenInfo : public TargetCIRGenInfo {
+public:
+  NVPTXTargetCIRGenInfo(CIRGenTypes &CGT, mlir::MLIRContext &MlCtx)
+      : TargetCIRGenInfo(std::make_unique<NVPTXABIInfo>(CGT, MlCtx)) {}
+};
+
+} // namespace
+
 // TODO(cir): remove the attribute once this gets used.
 LLVM_ATTRIBUTE_UNUSED
 static bool classifyReturnType(const CIRGenCXXABI &CXXABI,
@@ -443,6 +470,25 @@ cir::ABIArgInfo X86_64ABIInfo::classifyArgumentType(QualType Ty,
   return cir::ABIArgInfo::getDirect(ResType);
 }
 
+// Skeleton only. Implement when used in TargetLower stage.
+cir::ABIArgInfo NVPTXABIInfo::classifyReturnType(QualType RetTy) const {
+  llvm_unreachable("not yet implemented");
+}
+
+cir::ABIArgInfo NVPTXABIInfo::classifyArgumentType(QualType Ty) const {
+  llvm_unreachable("not yet implemented");
+}
+
+void NVPTXABIInfo::computeInfo(CIRGenFunctionInfo &FnInfo) const {
+  if (!getCXXABI().classifyReturnType(FnInfo))
+    FnInfo.getReturnInfo() = classifyReturnType(FnInfo.getReturnType());
+
+  for (auto &&[ArgumentsCount, I] : llvm::enumerate(FnInfo.arguments()))
+    I.info = ArgumentsCount < FnInfo.getNumRequiredArgs()
+                 ? classifyArgumentType(I.type)
+                 : cir::ABIArgInfo::getDirect();
+}
+
 ABIInfo::~ABIInfo() {}
 
 bool ABIInfo::isPromotableIntegerTypeForABI(QualType Ty) const {
@@ -634,5 +680,9 @@ const TargetCIRGenInfo &CIRGenModule::getTargetCIRGenInfo() {
   case llvm::Triple::spirv64: {
     return SetCIRGenInfo(new SPIRVTargetCIRGenInfo(genTypes));
   }
+
+  case llvm::Triple::nvptx64: {
+    return SetCIRGenInfo(new NVPTXTargetCIRGenInfo(genTypes, getMLIRContext()));
+  }
   }
 }