[mlir][gpu] Expose some utility functions from gpu-to-binary infra
#172205
+77
−67
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Member
|
@llvm/pr-subscribers-mlir-llvm @llvm/pr-subscribers-mlir Author: Ivan Butygin (Hardcode84) ChangesFor people who do not want to use a single monolithic pass. Full diff: https://github.com/llvm/llvm-project/pull/172205.diff 6 Files Affected:
diff --git a/mlir/include/mlir/Target/LLVM/ModuleToObject.h b/mlir/include/mlir/Target/LLVM/ModuleToObject.h
index 11fea6f0a4443..8c49e7712d6d3 100644
--- a/mlir/include/mlir/Target/LLVM/ModuleToObject.h
+++ b/mlir/include/mlir/Target/LLVM/ModuleToObject.h
@@ -23,6 +23,13 @@ class TargetMachine;
namespace mlir {
namespace LLVM {
+
+/// Translate LLVM module to textual ISA.
+FailureOr<std::string>
+translateModuleToISA(llvm::Module &llvmModule,
+ llvm::TargetMachine &targetMachine,
+ function_ref<InFlightDiagnostic()> emitError);
+
class ModuleTranslation;
/// Utility base class for transforming operations into binary objects, by
/// default it returns the serialized LLVM bitcode for the module. The
@@ -98,11 +105,6 @@ class ModuleToObject {
/// Optimize the module.
virtual LogicalResult optimizeModule(llvm::Module &module, int optL);
- /// Utility function for translating to ISA, returns `std::nullopt` on
- /// failure.
- static std::optional<std::string>
- translateToISA(llvm::Module &llvmModule, llvm::TargetMachine &targetMachine);
-
protected:
/// Module to transform to a binary object.
Operation &module;
diff --git a/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h b/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
index 8f5d4162984fa..9340d3d46a9be 100644
--- a/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
+++ b/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
@@ -43,6 +43,15 @@ enum class AMDGCNLibraries : uint32_t {
All = (LastLib << 1) - 1
};
+/// Assembles ISA to an object code.
+FailureOr<SmallVector<char, 0>>
+assembleIsa(StringRef isa, StringRef targetTriple, StringRef chip,
+ StringRef features, function_ref<InFlightDiagnostic()> emitError);
+
+FailureOr<SmallVector<char, 0>>
+linkObjectCode(ArrayRef<char> objectCode, StringRef toolkitPath,
+ function_ref<InFlightDiagnostic()> emitError);
+
/// Base class for all ROCDL serializations from GPU modules into binary
/// strings. By default this class serializes into LLVM bitcode.
class SerializeGPUModuleBase : public LLVM::ModuleToObject {
@@ -98,9 +107,6 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
moduleToObjectImpl(const gpu::TargetOptions &targetOptions,
llvm::Module &llvmModule);
- /// Returns the assembled ISA.
- std::optional<SmallVector<char, 0>> assembleIsa(StringRef isa);
-
/// ROCDL target attribute.
ROCDLTargetAttr target;
diff --git a/mlir/lib/Target/LLVM/ModuleToObject.cpp b/mlir/lib/Target/LLVM/ModuleToObject.cpp
index 4098ccc548dc1..eb99e69b626b6 100644
--- a/mlir/lib/Target/LLVM/ModuleToObject.cpp
+++ b/mlir/lib/Target/LLVM/ModuleToObject.cpp
@@ -205,9 +205,10 @@ LogicalResult ModuleToObject::optimizeModule(llvm::Module &module,
return success();
}
-std::optional<std::string>
-ModuleToObject::translateToISA(llvm::Module &llvmModule,
- llvm::TargetMachine &targetMachine) {
+FailureOr<std::string>
+mlir::LLVM::translateModuleToISA(llvm::Module &llvmModule,
+ llvm::TargetMachine &targetMachine,
+ function_ref<InFlightDiagnostic()> emitError) {
std::string targetISA;
llvm::raw_string_ostream stream(targetISA);
@@ -217,7 +218,7 @@ ModuleToObject::translateToISA(llvm::Module &llvmModule,
if (targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
llvm::CodeGenFileType::AssemblyFile))
- return std::nullopt;
+ return emitError() << "Target machine cannot emit assembly";
codegenPasses.run(llvmModule);
}
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 8760ea8588e2c..c22f8a23d2b2b 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -697,18 +697,20 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
return std::nullopt;
}
moduleToObjectTimer.startTimer();
- std::optional<std::string> serializedISA =
- translateToISA(llvmModule, **targetMachine);
+ FailureOr<std::string> serializedISA =
+ mlir::LLVM::translateModuleToISA(llvmModule, **targetMachine, [&]() {
+ return getOperation().emitError();
+ });
moduleToObjectTimer.stopTimer();
llvmToISATimeInMs = moduleToObjectTimer.getTotalTime().getWallTime() * 1000;
moduleToObjectTimer.clear();
- if (!serializedISA) {
+ if (failed(serializedISA)) {
getOperation().emitError() << "Failed translating the module to ISA.";
return std::nullopt;
}
if (isaCallback)
- isaCallback(serializedISA.value());
+ isaCallback(*serializedISA);
#define DEBUG_TYPE "serialize-to-isa"
LDBG() << "PTX for module: " << getOperation().getNameAttr() << "\n"
diff --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index f813f8db8fc94..a638f6e5eae7a 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -277,12 +277,10 @@ void SerializeGPUModuleBase::addControlVariables(
}
}
-std::optional<SmallVector<char, 0>>
-SerializeGPUModuleBase::assembleIsa(StringRef isa) {
- auto loc = getOperation().getLoc();
-
- StringRef targetTriple = this->triple;
-
+FailureOr<SmallVector<char, 0>>
+mlir::ROCDL::assembleIsa(StringRef isa, StringRef targetTriple, StringRef chip,
+ StringRef features,
+ function_ref<InFlightDiagnostic()> emitError) {
SmallVector<char, 0> result;
llvm::raw_svector_ostream os(result);
@@ -290,10 +288,8 @@ SerializeGPUModuleBase::assembleIsa(StringRef isa) {
std::string error;
const llvm::Target *target =
llvm::TargetRegistry::lookupTarget(triple, error);
- if (!target) {
- emitError(loc, Twine("failed to lookup target: ") + error);
- return std::nullopt;
- }
+ if (!target)
+ return emitError() << "failed to lookup target: " << error;
llvm::SourceMgr srcMgr;
srcMgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(isa), SMLoc());
@@ -330,50 +326,38 @@ SerializeGPUModuleBase::assembleIsa(StringRef isa) {
std::unique_ptr<llvm::MCTargetAsmParser> tap(
target->createMCAsmParser(*sti, *parser, *mcii, mcOptions));
- if (!tap) {
- emitError(loc, "assembler initialization error");
- return std::nullopt;
- }
+ if (!tap)
+ return emitError() << "assembler initialization error";
parser->setTargetParser(*tap);
parser->Run(false);
return std::move(result);
}
-std::optional<SmallVector<char, 0>>
-SerializeGPUModuleBase::compileToBinary(const std::string &serializedISA) {
- // Assemble the ISA.
- std::optional<SmallVector<char, 0>> isaBinary = assembleIsa(serializedISA);
-
- if (!isaBinary) {
- getOperation().emitError() << "failed during ISA assembling";
- return std::nullopt;
- }
-
+FailureOr<SmallVector<char, 0>>
+mlir::ROCDL::linkObjectCode(ArrayRef<char> objectCode, StringRef toolkitPath,
+ function_ref<InFlightDiagnostic()> emitError) {
// Save the ISA binary to a temp file.
int tempIsaBinaryFd = -1;
SmallString<128> tempIsaBinaryFilename;
if (llvm::sys::fs::createTemporaryFile("kernel%%", "o", tempIsaBinaryFd,
- tempIsaBinaryFilename)) {
- getOperation().emitError()
- << "failed to create a temporary file for dumping the ISA binary";
- return std::nullopt;
- }
+ tempIsaBinaryFilename))
+ return emitError()
+ << "failed to create a temporary file for dumping the ISA binary";
+
llvm::FileRemover cleanupIsaBinary(tempIsaBinaryFilename);
{
llvm::raw_fd_ostream tempIsaBinaryOs(tempIsaBinaryFd, true);
- tempIsaBinaryOs << StringRef(isaBinary->data(), isaBinary->size());
+ tempIsaBinaryOs << StringRef(objectCode.data(), objectCode.size());
tempIsaBinaryOs.flush();
}
// Create a temp file for HSA code object.
SmallString<128> tempHsacoFilename;
- if (llvm::sys::fs::createTemporaryFile("kernel", "hsaco",
- tempHsacoFilename)) {
- getOperation().emitError()
- << "failed to create a temporary file for the HSA code object";
- return std::nullopt;
- }
+ if (llvm::sys::fs::createTemporaryFile("kernel", "hsaco", tempHsacoFilename))
+ return emitError()
+ << "failed to create a temporary file for the HSA code object";
+
llvm::FileRemover cleanupHsaco(tempHsacoFilename);
llvm::SmallString<128> lldPath(toolkitPath);
@@ -381,25 +365,40 @@ SerializeGPUModuleBase::compileToBinary(const std::string &serializedISA) {
int lldResult = llvm::sys::ExecuteAndWait(
lldPath,
{"ld.lld", "-shared", tempIsaBinaryFilename, "-o", tempHsacoFilename});
- if (lldResult != 0) {
- getOperation().emitError() << "lld invocation failed";
- return std::nullopt;
- }
+ if (lldResult != 0)
+ return emitError() << "lld invocation failed";
// Load the HSA code object.
auto hsacoFile =
llvm::MemoryBuffer::getFile(tempHsacoFilename, /*IsText=*/false);
- if (!hsacoFile) {
- getOperation().emitError()
- << "failed to read the HSA code object from the temp file";
- return std::nullopt;
- }
+ if (!hsacoFile)
+ return emitError()
+ << "failed to read the HSA code object from the temp file";
StringRef buffer = (*hsacoFile)->getBuffer();
return SmallVector<char, 0>(buffer.begin(), buffer.end());
}
+std::optional<SmallVector<char, 0>>
+SerializeGPUModuleBase::compileToBinary(const std::string &serializedISA) {
+ auto errCallback = [&]() { return getOperation().emitError(); };
+ // Assemble the ISA.
+ FailureOr<SmallVector<char, 0>> isaBinary = ROCDL::assembleIsa(
+ serializedISA, this->triple, this->chip, this->features, errCallback);
+
+ if (failed(isaBinary))
+ return std::nullopt;
+
+ // Link the object code.
+ FailureOr<SmallVector<char, 0>> linkedCode =
+ ROCDL::linkObjectCode(*isaBinary, toolkitPath, errCallback);
+ if (failed(linkedCode))
+ return std::nullopt;
+
+ return linkedCode;
+}
+
std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
const gpu::TargetOptions &targetOptions, llvm::Module &llvmModule) {
// Return LLVM IR if the compilation target is offload.
@@ -422,9 +421,11 @@ std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
}
// Translate the Module to ISA.
- std::optional<std::string> serializedISA =
- translateToISA(llvmModule, **targetMachine);
- if (!serializedISA) {
+ FailureOr<std::string> serializedISA =
+ mlir::LLVM::translateModuleToISA(llvmModule, **targetMachine, [&]() {
+ return getOperation().emitError();
+ });
+ if (failed(serializedISA)) {
getOperation().emitError() << "failed translating the module to ISA";
return std::nullopt;
}
diff --git a/mlir/lib/Target/LLVM/XeVM/Target.cpp b/mlir/lib/Target/LLVM/XeVM/Target.cpp
index 1e6784a25d64a..d1bd4d78d28c0 100644
--- a/mlir/lib/Target/LLVM/XeVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/XeVM/Target.cpp
@@ -301,9 +301,11 @@ SPIRVSerializer::moduleToObject(llvm::Module &llvmModule) {
// Return SPIRV if the compilation target is `assembly`.
if (targetOptions.getCompilationTarget() ==
gpu::CompilationTarget::Assembly) {
- std::optional<std::string> serializedISA =
- translateToISA(llvmModule, **targetMachine);
- if (!serializedISA) {
+ FailureOr<std::string> serializedISA =
+ mlir::LLVM::translateModuleToISA(llvmModule, **targetMachine, [&]() {
+ return getGPUModuleOp().emitError();
+ });
+ if (failed(serializedISA)) {
getGPUModuleOp().emitError() << "Failed translating the module to ISA."
<< triple << ", can't compile with LLVM\n";
return std::nullopt;
|
fabianmcg
reviewed
Dec 14, 2025
Contributor
fabianmcg
left a comment
fabianmcg
left a comment
There was a problem hiding this comment.
In general LGTM, just look at the comments.
fabianmcg
approved these changes
Dec 14, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
For people who do not want to use a single monolithic pass.