From 2a5050aa5ef56b01cf4a8f73e0d0eddd6d9ce2a9 Mon Sep 17 00:00:00 2001 From: Zichen Lu Date: Tue, 11 Feb 2025 19:58:58 +0800 Subject: [PATCH] [mlir][target][nvvm] Perf by stage and store into properties (#126178) Implement the feature about perf by stage(llvm-ir -> isa, isa->binary). The results will be stored into the properties, then users can use them after using GpuModuleToBinary Pass. --- mlir/lib/Target/LLVM/NVVM/Target.cpp | 70 +++++++++++++++++-- .../Dialect/GPU/module-to-binary-nvvm.mlir | 2 +- .../dialects/gpu/module-to-binary-nvvm.py | 2 +- 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp index b7d60ed59db0..e240a7ae4917 100644 --- a/mlir/lib/Target/LLVM/NVVM/Target.cpp +++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp @@ -35,10 +35,12 @@ #include "llvm/Support/Process.h" #include "llvm/Support/Program.h" #include "llvm/Support/TargetSelect.h" +#include "llvm/Support/Timer.h" #include "llvm/Support/raw_ostream.h" #include #include +#include using namespace mlir; using namespace mlir::NVVM; @@ -220,6 +222,16 @@ public: std::optional> moduleToObject(llvm::Module &llvmModule) override; + /// Get LLVMIR->ISA performance result. + /// Return nullopt if moduleToObject has not been called or the target format + /// is LLVMIR. + std::optional getLLVMIRToISATimeInMs(); + + /// Get ISA->Binary performance result. + /// Return nullopt if moduleToObject has not been called or the target format + /// is LLVMIR or ISA. + std::optional getISAToBinaryTimeInMs(); + private: using TmpFile = std::pair, llvm::FileRemover>; @@ -235,13 +247,20 @@ private: /// Target options. gpu::TargetOptions targetOptions; + + /// LLVMIR->ISA perf result. + std::optional llvmToISATimeInMs; + + /// ISA->Binary perf result. + std::optional isaToBinaryTimeInMs; }; } // namespace NVPTXSerializer::NVPTXSerializer(Operation &module, NVVMTargetAttr target, const gpu::TargetOptions &targetOptions) : SerializeGPUModuleBase(module, target, targetOptions), - targetOptions(targetOptions) {} + targetOptions(targetOptions), llvmToISATimeInMs(std::nullopt), + isaToBinaryTimeInMs(std::nullopt) {} std::optional NVPTXSerializer::createTemp(StringRef name, StringRef suffix) { @@ -256,6 +275,14 @@ NVPTXSerializer::createTemp(StringRef name, StringRef suffix) { return TmpFile(filename, llvm::FileRemover(filename.c_str())); } +std::optional NVPTXSerializer::getLLVMIRToISATimeInMs() { + return llvmToISATimeInMs; +} + +std::optional NVPTXSerializer::getISAToBinaryTimeInMs() { + return isaToBinaryTimeInMs; +} + gpu::GPUModuleOp NVPTXSerializer::getOperation() { return dyn_cast(&SerializeGPUModuleBase::getOperation()); } @@ -618,6 +645,10 @@ NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) { std::optional> NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) { + llvm::Timer moduleToObjectTimer( + "moduleToObjectTimer", + "Timer for perf llvm-ir -> isa and isa -> binary."); + moduleToObjectTimer.startTimer(); // Return LLVM IR if the compilation target is `offload`. #define DEBUG_TYPE "serialize-to-llvm" LLVM_DEBUG({ @@ -650,6 +681,11 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) { getOperation().emitError() << "Failed translating the module to ISA."; return std::nullopt; } + + moduleToObjectTimer.stopTimer(); + llvmToISATimeInMs = moduleToObjectTimer.getTotalTime().getWallTime() * 1000; + moduleToObjectTimer.clear(); + moduleToObjectTimer.startTimer(); if (isaCallback) isaCallback(serializedISA.value()); @@ -669,17 +705,24 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) { return SmallVector(bin.begin(), bin.end()); } + std::optional> result; // Compile to binary. #if MLIR_ENABLE_NVPTXCOMPILER - return compileToBinaryNVPTX(*serializedISA); + result = compileToBinaryNVPTX(*serializedISA); #else - return compileToBinary(*serializedISA); + result = compileToBinary(*serializedISA); #endif // MLIR_ENABLE_NVPTXCOMPILER + + moduleToObjectTimer.stopTimer(); + isaToBinaryTimeInMs = moduleToObjectTimer.getTotalTime().getWallTime() * 1000; + moduleToObjectTimer.clear(); + return result; } std::optional> NVVMTargetAttrImpl::serializeToObject(Attribute attribute, Operation *module, const gpu::TargetOptions &options) const { + Builder builder(attribute.getContext()); assert(module && "The module must be non null."); if (!module) return std::nullopt; @@ -689,7 +732,16 @@ NVVMTargetAttrImpl::serializeToObject(Attribute attribute, Operation *module, } NVPTXSerializer serializer(*module, cast(attribute), options); serializer.init(); - return serializer.run(); + std::optional> result = serializer.run(); + auto llvmToISATimeInMs = serializer.getLLVMIRToISATimeInMs(); + if (llvmToISATimeInMs.has_value()) + module->setAttr("LLVMIRToISATimeInMs", + builder.getI64IntegerAttr(*llvmToISATimeInMs)); + auto isaToBinaryTimeInMs = serializer.getISAToBinaryTimeInMs(); + if (isaToBinaryTimeInMs.has_value()) + module->setAttr("ISAToBinaryTimeInMs", + builder.getI64IntegerAttr(*isaToBinaryTimeInMs)); + return result; } Attribute @@ -700,7 +752,7 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module, gpu::CompilationTarget format = options.getCompilationTarget(); DictionaryAttr objectProps; Builder builder(attribute.getContext()); - SmallVector properties; + SmallVector properties; if (format == gpu::CompilationTarget::Assembly) properties.push_back( builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO()))); @@ -709,6 +761,14 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module, properties.push_back(builder.getNamedAttr(gpu::elfSectionName, builder.getStringAttr(section))); + for (const auto *perfName : {"LLVMIRToISATimeInMs", "ISAToBinaryTimeInMs"}) { + if (module->hasAttr(perfName)) { + IntegerAttr attr = llvm::dyn_cast(module->getAttr(perfName)); + properties.push_back(builder.getNamedAttr( + perfName, builder.getI64IntegerAttr(attr.getInt()))); + } + } + if (!properties.empty()) objectProps = builder.getDictionaryAttr(properties); diff --git a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir index e6284ccf94b5..dee680208bdf 100644 --- a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir +++ b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir @@ -16,7 +16,7 @@ module attributes {gpu.container_module} { } // CHECK-LABEL:gpu.binary @kernel_module2 - // CHECK-ISA:[#gpu.object<#nvvm.target, properties = {O = 2 : i32}, assembly = "{{.*}}">, #gpu.object<#nvvm.target, properties = {O = 2 : i32}, assembly = "{{.*}}">] + // CHECK-ISA:[#gpu.object<#nvvm.target, properties = {LLVMIRToISATimeInMs = {{[0-9]+}} : i64, O = 2 : i32}, assembly = "{{.*}}">, #gpu.object<#nvvm.target, properties = {LLVMIRToISATimeInMs = {{[0-9]+}} : i64, O = 2 : i32}, assembly = "{{.*}}">] gpu.module @kernel_module2 [#nvvm.target, #nvvm.target] { llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i64, %arg4: i64, diff --git a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py index e0225911b214..96e3f6efcc5d 100644 --- a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py +++ b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py @@ -82,5 +82,5 @@ module attributes {gpu.container_module} { print(gpu.CompilationTarget(o.format)) # CHECK: b'//\n// Generated by LLVM NVPTX Back-End{{.*}}' print(o.object) - # CHECK: {O = 2 : i32} + # CHECK: {LLVMIRToISATimeInMs = {{[0-9]+}} : i64, O = 2 : i32} print(o.properties)