Reland [mlir][gpu] Use SmallString, FailureOr and StringRef in module-to-binary infra (NFC)
#172390
+154
−177
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
…-to-binary` infra (NFC) (llvm#172284) Instead of `std::string`, `std::optional` and `const std::string&`.
Member
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-llvm Author: Ivan Butygin (Hardcode84) ChangesReland #172284
Patch is 29.18 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/172390.diff 7 Files Affected:
diff --git a/mlir/include/mlir/Target/LLVM/ModuleToObject.h b/mlir/include/mlir/Target/LLVM/ModuleToObject.h
index ea0b4b4a6ab7f..7d232dcf9ab90 100644
--- a/mlir/include/mlir/Target/LLVM/ModuleToObject.h
+++ b/mlir/include/mlir/Target/LLVM/ModuleToObject.h
@@ -45,8 +45,7 @@ class ModuleToObject {
virtual std::optional<SmallVector<char, 0>> run();
/// Translate LLVM module to textual ISA.
- /// TODO: switch to SmallString
- static FailureOr<std::string>
+ static FailureOr<SmallString<0>>
translateModuleToISA(llvm::Module &llvmModule,
llvm::TargetMachine &targetMachine,
function_ref<InFlightDiagnostic()> emitError);
@@ -76,13 +75,13 @@ class ModuleToObject {
/// Serializes the LLVM IR bitcode to an object file, by default it serializes
/// to LLVM bitcode.
- virtual std::optional<SmallVector<char, 0>>
+ virtual FailureOr<SmallVector<char, 0>>
moduleToObject(llvm::Module &llvmModule);
protected:
/// Create the target machine based on the target triple and chip.
/// This can fail if the target is not available.
- std::optional<llvm::TargetMachine *> getOrCreateTargetMachine();
+ FailureOr<llvm::TargetMachine *> getOrCreateTargetMachine();
/// Loads a bitcode file from path.
std::unique_ptr<llvm::Module> loadBitcodeFile(llvm::LLVMContext &context,
diff --git a/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h b/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
index 9340d3d46a9be..302ff6e7b8575 100644
--- a/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
+++ b/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
@@ -99,11 +99,11 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
StringRef abiVer);
/// Compiles assembly to a binary.
- virtual std::optional<SmallVector<char, 0>>
- compileToBinary(const std::string &serializedISA);
+ virtual FailureOr<SmallVector<char, 0>>
+ compileToBinary(StringRef serializedISA);
/// Default implementation of `ModuleToObject::moduleToObject`.
- std::optional<SmallVector<char, 0>>
+ FailureOr<SmallVector<char, 0>>
moduleToObjectImpl(const gpu::TargetOptions &targetOptions,
llvm::Module &llvmModule);
diff --git a/mlir/include/mlir/Target/LLVM/XeVM/Utils.h b/mlir/include/mlir/Target/LLVM/XeVM/Utils.h
index 5d523f1ec457f..455c8303a9aa8 100644
--- a/mlir/include/mlir/Target/LLVM/XeVM/Utils.h
+++ b/mlir/include/mlir/Target/LLVM/XeVM/Utils.h
@@ -40,8 +40,8 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
gpu::GPUModuleOp getGPUModuleOp();
/// Compiles to native code using `ocloc`.
- std::optional<SmallVector<char, 0>> compileToBinary(const std::string &asmStr,
- StringRef inputFormat);
+ FailureOr<SmallVector<char, 0>> compileToBinary(StringRef asmStr,
+ StringRef inputFormat);
protected:
/// XeVM Target attribute.
diff --git a/mlir/lib/Target/LLVM/ModuleToObject.cpp b/mlir/lib/Target/LLVM/ModuleToObject.cpp
index e20fd544ed8f7..8eaf8d9cadcd2 100644
--- a/mlir/lib/Target/LLVM/ModuleToObject.cpp
+++ b/mlir/lib/Target/LLVM/ModuleToObject.cpp
@@ -50,8 +50,7 @@ ModuleToObject::~ModuleToObject() = default;
Operation &ModuleToObject::getOperation() { return module; }
-std::optional<llvm::TargetMachine *>
-ModuleToObject::getOrCreateTargetMachine() {
+FailureOr<llvm::TargetMachine *> ModuleToObject::getOrCreateTargetMachine() {
if (targetMachine)
return targetMachine.get();
// Load the target.
@@ -59,17 +58,17 @@ ModuleToObject::getOrCreateTargetMachine() {
llvm::Triple parsedTriple(triple);
const llvm::Target *target =
llvm::TargetRegistry::lookupTarget(parsedTriple, error);
- if (!target) {
- getOperation().emitError()
- << "Failed to lookup target for triple '" << triple << "' " << error;
- return std::nullopt;
- }
+ if (!target)
+ return getOperation().emitError()
+ << "Failed to lookup target for triple '" << triple << "' " << error;
// Create the target machine using the target.
targetMachine.reset(
target->createTargetMachine(parsedTriple, chip, features, {}, {}));
if (!targetMachine)
- return std::nullopt;
+ return getOperation().emitError()
+ << "Failed to create target machine for triple '" << triple << "'";
+
return targetMachine.get();
}
@@ -183,9 +182,8 @@ LogicalResult ModuleToObject::optimizeModule(llvm::Module &module,
return getOperation().emitError()
<< "Invalid optimization level: " << optLevel << ".";
- std::optional<llvm::TargetMachine *> targetMachine =
- getOrCreateTargetMachine();
- if (!targetMachine)
+ FailureOr<llvm::TargetMachine *> targetMachine = getOrCreateTargetMachine();
+ if (failed(targetMachine))
return getOperation().emitError()
<< "Target Machine unavailable for triple " << triple
<< ", can't optimize with LLVM\n";
@@ -205,11 +203,11 @@ LogicalResult ModuleToObject::optimizeModule(llvm::Module &module,
return success();
}
-FailureOr<std::string> ModuleToObject::translateModuleToISA(
+FailureOr<SmallString<0>> ModuleToObject::translateModuleToISA(
llvm::Module &llvmModule, llvm::TargetMachine &targetMachine,
function_ref<InFlightDiagnostic()> emitError) {
- std::string targetISA;
- llvm::raw_string_ostream stream(targetISA);
+ SmallString<0> targetISA;
+ llvm::raw_svector_ostream stream(targetISA);
{ // Drop pstream after this to prevent the ISA from being stuck buffering
llvm::buffer_ostream pstream(stream);
@@ -226,16 +224,16 @@ FailureOr<std::string> ModuleToObject::translateModuleToISA(
void ModuleToObject::setDataLayoutAndTriple(llvm::Module &module) {
// Create the target machine.
- std::optional<llvm::TargetMachine *> targetMachine =
- getOrCreateTargetMachine();
- if (targetMachine) {
- // Set the data layout and target triple of the module.
- module.setDataLayout((*targetMachine)->createDataLayout());
- module.setTargetTriple((*targetMachine)->getTargetTriple());
- }
+ FailureOr<llvm::TargetMachine *> targetMachine = getOrCreateTargetMachine();
+ if (failed(targetMachine))
+ return;
+
+ // Set the data layout and target triple of the module.
+ module.setDataLayout((*targetMachine)->createDataLayout());
+ module.setTargetTriple((*targetMachine)->getTargetTriple());
}
-std::optional<SmallVector<char, 0>>
+FailureOr<SmallVector<char, 0>>
ModuleToObject::moduleToObject(llvm::Module &llvmModule) {
SmallVector<char, 0> binaryData;
// Write the LLVM module bitcode to a buffer.
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index cae5adfc086df..43063d9eb5bee 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -212,16 +212,14 @@ class NVPTXSerializer : public SerializeGPUModuleBase {
gpu::GPUModuleOp getOperation();
/// Compiles PTX to cubin using `ptxas`.
- std::optional<SmallVector<char, 0>>
- compileToBinary(const std::string &ptxCode);
+ FailureOr<SmallVector<char, 0>> compileToBinary(StringRef ptxCode);
/// Compiles PTX to cubin using the `nvptxcompiler` library.
- std::optional<SmallVector<char, 0>>
- compileToBinaryNVPTX(const std::string &ptxCode);
+ FailureOr<SmallVector<char, 0>> compileToBinaryNVPTX(StringRef ptxCode);
/// Serializes the LLVM module to an object format, depending on the
/// compilation target selected in target options.
- std::optional<SmallVector<char, 0>>
+ FailureOr<SmallVector<char, 0>>
moduleToObject(llvm::Module &llvmModule) override;
/// Get LLVMIR->ISA performance result.
@@ -346,8 +344,8 @@ static void setOptionalCommandlineArguments(NVVMTargetAttr target,
// TODO: clean this method & have a generic tool driver or never emit binaries
// with this mechanism and let another stage take care of it.
-std::optional<SmallVector<char, 0>>
-NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
+FailureOr<SmallVector<char, 0>>
+NVPTXSerializer::compileToBinary(StringRef ptxCode) {
// Determine if the serializer should create a fatbinary with the PTX embeded
// or a simple CUBIN binary.
const bool createFatbin =
@@ -356,12 +354,12 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
// Find the `ptxas` & `fatbinary` tools.
std::optional<std::string> ptxasCompiler = findTool("ptxas");
if (!ptxasCompiler)
- return std::nullopt;
+ return failure();
std::optional<std::string> fatbinaryTool;
if (createFatbin) {
fatbinaryTool = findTool("fatbinary");
if (!fatbinaryTool)
- return std::nullopt;
+ return failure();
}
Location loc = getOperation().getLoc();
@@ -373,13 +371,13 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
// Create temp files:
std::optional<TmpFile> ptxFile = createTemp(basename, "ptx");
if (!ptxFile)
- return std::nullopt;
+ return failure();
std::optional<TmpFile> logFile = createTemp(basename, "log");
if (!logFile)
- return std::nullopt;
+ return failure();
std::optional<TmpFile> binaryFile = createTemp(basename, "bin");
if (!binaryFile)
- return std::nullopt;
+ return failure();
TmpFile cubinFile;
if (createFatbin) {
std::string cubinFilename = (ptxFile->first + ".cubin").str();
@@ -392,17 +390,15 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
// Dump the PTX to a temp file.
{
llvm::raw_fd_ostream ptxStream(ptxFile->first, ec);
- if (ec) {
- emitError(loc) << "Couldn't open the file: `" << ptxFile->first
- << "`, error message: " << ec.message();
- return std::nullopt;
- }
+ if (ec)
+ return emitError(loc) << "Couldn't open the file: `" << ptxFile->first
+ << "`, error message: " << ec.message();
+
ptxStream << ptxCode;
- if (ptxStream.has_error()) {
- emitError(loc) << "An error occurred while writing the PTX to: `"
- << ptxFile->first << "`.";
- return std::nullopt;
- }
+ if (ptxStream.has_error())
+ return emitError(loc) << "An error occurred while writing the PTX to: `"
+ << ptxFile->first << "`.";
+
ptxStream.flush();
}
@@ -466,20 +462,18 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
// Helper function for printing tool error logs.
std::string message;
auto emitLogError =
- [&](StringRef toolName) -> std::optional<SmallVector<char, 0>> {
+ [&](StringRef toolName) -> FailureOr<SmallVector<char, 0>> {
if (message.empty()) {
llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> toolStderr =
llvm::MemoryBuffer::getFile(logFile->first);
if (toolStderr)
- emitError(loc) << toolName << " invocation failed. Log:\n"
- << toolStderr->get()->getBuffer();
+ return emitError(loc) << toolName << " invocation failed. Log:\n"
+ << toolStderr->get()->getBuffer();
else
- emitError(loc) << toolName << " invocation failed.";
- return std::nullopt;
+ return emitError(loc) << toolName << " invocation failed.";
}
- emitError(loc) << toolName
- << " invocation failed, error message: " << message;
- return std::nullopt;
+ return emitError(loc) << toolName
+ << " invocation failed, error message: " << message;
};
// Invoke PTXAS.
@@ -536,11 +530,11 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
// Read the fatbin.
llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> binaryBuffer =
llvm::MemoryBuffer::getFile(binaryFile->first);
- if (!binaryBuffer) {
- emitError(loc) << "Couldn't open the file: `" << binaryFile->first
- << "`, error message: " << binaryBuffer.getError().message();
- return std::nullopt;
- }
+ if (!binaryBuffer)
+ return emitError(loc) << "Couldn't open the file: `" << binaryFile->first
+ << "`, error message: "
+ << binaryBuffer.getError().message();
+
StringRef fatbin = (*binaryBuffer)->getBuffer();
return SmallVector<char, 0>(fatbin.begin(), fatbin.end());
}
@@ -569,8 +563,8 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
} \
} while (false)
-std::optional<SmallVector<char, 0>>
-NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) {
+FailureOr<SmallVector<char, 0>>
+NVPTXSerializer::compileToBinaryNVPTX(StringRef ptxCode) {
Location loc = getOperation().getLoc();
nvPTXCompilerHandle compiler = nullptr;
nvPTXCompileResult status;
@@ -601,13 +595,12 @@ NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) {
SmallVector<char> log(logSize + 1, 0);
RETURN_ON_NVPTXCOMPILER_ERROR(
nvPTXCompilerGetErrorLog(compiler, log.data()));
- emitError(loc) << "NVPTX compiler invocation failed, error log: "
- << log.data();
+ return emitError(loc)
+ << "NVPTX compiler invocation failed, error log: " << log.data();
} else {
- emitError(loc) << "NVPTX compiler invocation failed with error code: "
- << status;
+ return emitError(loc)
+ << "NVPTX compiler invocation failed with error code: " << status;
}
- return std::nullopt;
}
// Retrieve the binary.
@@ -666,7 +659,7 @@ NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) {
}
#endif // MLIR_ENABLE_NVPTXCOMPILER
-std::optional<SmallVector<char, 0>>
+FailureOr<SmallVector<char, 0>>
NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
llvm::Timer moduleToObjectTimer(
"moduleToObjectTimer",
@@ -683,30 +676,27 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
return SerializeGPUModuleBase::moduleToObject(llvmModule);
#if !LLVM_HAS_NVPTX_TARGET
- getOperation()->emitError(
+ return getOperation()->emitError(
"The `NVPTX` target was not built. Please enable it when building LLVM.");
- return std::nullopt;
#endif // LLVM_HAS_NVPTX_TARGET
// Emit PTX code.
- std::optional<llvm::TargetMachine *> targetMachine =
- getOrCreateTargetMachine();
- if (!targetMachine) {
- getOperation().emitError() << "Target Machine unavailable for triple "
- << triple << ", can't optimize with LLVM\n";
- return std::nullopt;
- }
+ FailureOr<llvm::TargetMachine *> targetMachine = getOrCreateTargetMachine();
+ if (failed(targetMachine))
+ return getOperation().emitError()
+ << "Target Machine unavailable for triple " << triple
+ << ", can't optimize with LLVM\n";
+
moduleToObjectTimer.startTimer();
- FailureOr<std::string> serializedISA =
+ FailureOr<SmallString<0>> serializedISA =
translateModuleToISA(llvmModule, **targetMachine,
[&]() { return getOperation().emitError(); });
moduleToObjectTimer.stopTimer();
llvmToISATimeInMs = moduleToObjectTimer.getTotalTime().getWallTime() * 1000;
moduleToObjectTimer.clear();
- if (failed(serializedISA)) {
- getOperation().emitError() << "Failed translating the module to ISA.";
- return std::nullopt;
- }
+ if (failed(serializedISA))
+ return getOperation().emitError()
+ << "Failed translating the module to ISA.";
if (isaCallback)
isaCallback(*serializedISA);
@@ -720,7 +710,7 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
if (targetOptions.getCompilationTarget() == gpu::CompilationTarget::Assembly)
return SmallVector<char, 0>(serializedISA->begin(), serializedISA->end());
- std::optional<SmallVector<char, 0>> result;
+ FailureOr<SmallVector<char, 0>> result;
moduleToObjectTimer.startTimer();
// Compile to binary.
#if MLIR_ENABLE_NVPTXCOMPILER
diff --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index a1970dac8785a..1af7eabfb4b10 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -292,7 +292,8 @@ mlir::ROCDL::assembleIsa(StringRef isa, StringRef targetTriple, StringRef chip,
return emitError() << "failed to lookup target: " << error;
llvm::SourceMgr srcMgr;
- srcMgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(isa), SMLoc());
+ // Copy buffer to ensure it's null terminated.
+ srcMgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBufferCopy(isa), SMLoc());
const llvm::MCTargetOptions mcOptions;
std::unique_ptr<llvm::MCRegisterInfo> mri(target->createMCRegInfo(triple));
@@ -380,26 +381,26 @@ mlir::ROCDL::linkObjectCode(ArrayRef<char> objectCode, StringRef toolkitPath,
return SmallVector<char, 0>(buffer.begin(), buffer.end());
}
-std::optional<SmallVector<char, 0>>
-SerializeGPUModuleBase::compileToBinary(const std::string &serializedISA) {
+FailureOr<SmallVector<char, 0>>
+SerializeGPUModuleBase::compileToBinary(StringRef 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;
+ return failure();
// Link the object code.
FailureOr<SmallVector<char, 0>> linkedCode =
ROCDL::linkObjectCode(*isaBinary, toolkitPath, errCallback);
if (failed(linkedCode))
- return std::nullopt;
+ return failure();
return linkedCode;
}
-std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
+FailureOr<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
const gpu::TargetOptions &targetOptions, llvm::Module &llvmModule) {
// Return LLVM IR if the compilation target is offload.
#define DEBUG_TYPE "serialize-to-llvm"
@@ -412,22 +413,19 @@ std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
if (targetOptions.getCompilationTarget() == gpu::CompilationTarget::Offload)
return SerializeGPUModuleBase::moduleToObject(llvmModule);
- std::optional<llvm::TargetMachine *> targetMachine =
- getOrCreateTargetMachine();
- if (!targetMachine) {
- getOperation().emitError() << "target Machine unavailable for triple "
- << triple << ", can't compile with LLVM";
- return std::nullopt;
- }
+ FailureOr<llvm::TargetMachine *> targetMachine = getOrCreateTargetMachine();
+ if (failed(targetMachine))
+ return getOperation().emitError()
+ << "target Machine unavailable for triple " << triple
+ << ", can't compile with LLVM";
// Translate the Module to ISA.
- FailureOr<std::string> serializedISA =
+ FailureOr<SmallString<0>> serializedISA =
translateModuleToISA(llvmModule, **targetMachine,
[&]() { return getOperation().emitError(); });
- if (failed(serializedISA)) {
- getOperation().emitError() << "failed translating the module to ISA";
- return std::nullopt;
- }
+ if (failed(serializedISA))
+ return getOperation().emitError() << "failed translating the module to ISA";
+
#define DEBUG_TYPE "serialize-to-isa"
LLVM_DEBUG({
llvm::dbgs() << "ISA for module: "
@@ -440,10 +438,9 @@ std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
return SmallVector<char, 0>(serializedISA->begin(), serializedISA->end());
// Compiling to binary requires a valid ROCm path, fail if it's not found.
- if (getToolkitPath().empty()) {
- getOperation().emitError() << "invalid ROCm path, please set a valid path";
- return std::nullopt;
- }
+ if (getToolkitPath().empty())
+ return getOperation().emitError()
+ << "invalid ROCm path, please set a valid path";
// Compile to binary.
return compileToBinary(*serializedISA);
@@ -456,7 +453,7 @@ class AMDGPUSerializer : public SerializeGPUModuleBase {
AMDGPUSerializer(Operation &module, ROCDLTargetAttr target,
const gpu::TargetOptions &targetOptions);
- std::optional<SmallVector<char, 0>>
+ FailureOr<SmallVector<char, 0>>
moduleToObject(llvm::Module &llvmModule) override;
private:
@@ -470,7 +467,7 @@ ...
[truncated]
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
Reland #172284
MCAsmParserexpects buffer to be null terminated, had to usegetMemBufferCopywhich is unfortunate.