[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.

Author: MikaOvO

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 <cstdint>
 #include <cstdlib>
+#include <optional>
 
 using namespace mlir;
 using namespace mlir::NVVM;
@@ -220,6 +222,16 @@ class NVPTXSerializer : public SerializeGPUModuleBase {
   std::optional<SmallVector<char, 0>>
   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<int64_t> getLLVMIRToISATimeInMs();
+
+  /// Get ISA->Binary performance result.
+  /// Return nullopt if moduleToObject has not been called or the target format
+  /// is LLVMIR or ISA.
+  std::optional<int64_t> getISAToBinaryTimeInMs();
+
 private:
   using TmpFile = std::pair<llvm::SmallString<128>, llvm::FileRemover>;
 
@@ -235,13 +247,20 @@ class NVPTXSerializer : public SerializeGPUModuleBase {
 
   /// Target options.
   gpu::TargetOptions targetOptions;
+
+  /// LLVMIR->ISA perf result.
+  std::optional<int64_t> llvmToISATimeInMs;
+
+  /// ISA->Binary perf result.
+  std::optional<int64_t> 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::TmpFile>
 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<int64_t> NVPTXSerializer::getLLVMIRToISATimeInMs() {
+  return llvmToISATimeInMs;
+}
+
+std::optional<int64_t> NVPTXSerializer::getISAToBinaryTimeInMs() {
+  return isaToBinaryTimeInMs;
+}
+
 gpu::GPUModuleOp NVPTXSerializer::getOperation() {
   return dyn_cast<gpu::GPUModuleOp>(&SerializeGPUModuleBase::getOperation());
 }
@@ -618,6 +645,10 @@ NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) {
 
 std::optional<SmallVector<char, 0>>
 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<char, 0>(bin.begin(), bin.end());
   }
 
+  std::optional<SmallVector<char, 0>> 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<SmallVector<char, 0>>
 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<NVVMTargetAttr>(attribute), options);
   serializer.init();
-  return serializer.run();
+  std::optional<SmallVector<char, 0>> 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<NamedAttribute, 2> properties;
+  SmallVector<NamedAttribute, 4> 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<IntegerAttr>(module->getAttr(perfName));
+      properties.push_back(builder.getNamedAttr(
+          perfName, builder.getI64IntegerAttr(attr.getInt())));
+    }
+  }
+
   if (!properties.empty())
     objectProps = builder.getDictionaryAttr(properties);
 

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<flags = {fast}>, properties = {O = 2 : i32}, assembly = "{{.*}}">, #gpu.object<#nvvm.target, properties = {O = 2 : i32}, assembly = "{{.*}}">]
+  // CHECK-ISA:[#gpu.object<#nvvm.target<flags = {fast}>, 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<flags = {fast}>, #nvvm.target] {
     llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
         %arg2: !llvm.ptr, %arg3: i64, %arg4: i64,

mlir/test/python/dialects/gpu/module-to-binary-nvvm.py

@@ -82,5 +82,5 @@ def testGPUToASMBin():
     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)