[mlir][nvvm] Introduce performance tuning directives

[guraypp]

Summary

PTX programming models provides some performance tuning directives; see https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#performance-tuning-directives

The downstream compiler namely ptxas leverages these information for better register allocation or to handle other resource management that improves the performance.

This revision introduce all the kernel based directives to MLIR's NVVM dialect. The list is below

maxnreg			-> 	max register per thread in CTA
maxntid			-> 	max threads per CTA
reqntid			-> 	exact number of threads per CTA
minnctapersm		-> 	min CTA per SM

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

53–59

Can we add a dialect attribute verifier for these to be of the right type?

[guraypp]

Add attribute verifiers

[vzakhari]

Hi Guray,

This change-set seems to break the build:

ld.lld: error: undefined symbol: mlir::extractFromI64ArrayAttr(mlir::Attribute)
>>> referenced by NVVMToLLVMIRTranslation.cpp:142 (/llvm-project/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp:142)
>>>               CMakeFiles/obj.MLIRNVVMToLLVMIRTranslation.dir/NVVMToLLVMIRTranslation.cpp.o:((anonymous namespace)::NVVMDialectLLVMIRTranslationInterface::amendOperation(mlir::Operation*, mlir::NamedAttribute, mlir::LLVM::ModuleTranslation&) const)
>>> referenced by NVVMToLLVMIRTranslation.cpp:152 (/llvm-project/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp:152)
>>>               CMakeFiles/obj.MLIRNVVMToLLVMIRTranslation.dir/NVVMToLLVMIRTranslation.cpp.o:((anonymous namespace)::NVVMDialectLLVMIRTranslationInterface::amendOperation(mlir::Operation*, mlir::NamedAttribute, mlir::LLVM::ModuleTranslation&) const)
collect2: error: ld returned 1 exit status

It looks like there is missing dependency on MLIRDialectUtils.

[vzakhari]

I tried to fix it in https://github.com/llvm/llvm-project/commit/b9255099382d0121e567979e7d5b4ca1328b4a41

[guraypp]

In D136931#3893571, @vzakhari wrote:

I tried to fix it in https://github.com/llvm/llvm-project/commit/b9255099382d0121e567979e7d5b4ca1328b4a41

Hi Slava, thanks for fixing it.
I am not sure how did I miss it, I remember the buildbot was fine.

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

def NVVM_Dialect : Dialect {
   let cppNamespace = "::mlir::NVVM";
   let dependentDialects = ["LLVM::LLVMDialect"];
   let hasOperationAttrVerify = 1;
 
   let extraClassDeclaration = [{
     /// Get the name of the attribute used to annotate external kernel
     /// functions.
     static StringRef getKernelFuncAttrName() { return "nvvm.kernel"; }
+    /// Get the name of the attribute used to annotate max threads required
+    /// per CTA for kernel functions.
+    static StringRef getMaxntidAttrName() { return "nvvm.maxntid"; }
+    /// Get the name of the metadata names for each dimension
+    static StringRef getMaxntidXName() { return "maxntidx"; }
+    static StringRef getMaxntidYName() { return "maxntidy"; }
+    static StringRef getMaxntidZName() { return "maxntidz"; }
+
+    /// Get the name of the attribute used to annotate exact threads required
+    /// per CTA for kernel functions.
+    static StringRef getReqntidAttrName() { return "nvvm.reqntid"; }
+    /// Get the name of the metadata names for each dimension
+    static StringRef getReqntidXName() { return "reqntidx"; }
+    static StringRef getReqntidYName() { return "reqntidy"; }
+    static StringRef getReqntidZName() { return "reqntidz"; }
+
+    /// Get the name of the attribute used to annotate min CTA required
+    /// per SM for kernel functions.
+    static StringRef getMinctasmAttrName() { return "nvvm.minctasm"; }
+
+    /// Get the name of the attribute used to annotate max number of 
+    /// registers that can be allocated per thread.
+    static StringRef getMaxnregAttrName() { return "nvvm.maxnreg"; }
   }];
 
   let useDefaultAttributePrinterParser = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // NVVM op definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp

 //===- NVVMToLLVMIRTranslation.cpp - Translate NVVM to LLVM IR ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a translation between the MLIR NVVM dialect and
 // LLVM IR.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
#include "mlir/Dialect/LLVMIR/NVVMConversions.inc"
 
     return failure();
   }
 
   /// Attaches module-level metadata for functions marked as kernels.
   LogicalResult
   amendOperation(Operation *op, NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const final {
-    if (attribute.getName() == NVVM::NVVMDialect::getKernelFuncAttrName()) {
-      auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
-      if (!func)
-        return failure();
-
-      llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
-      llvm::Function *llvmFunc =
-          moduleTranslation.lookupFunction(func.getName());
+    llvm::Function *llvmFunc = moduleTranslation.lookupFunction(func.getName());
+
+    auto generateMetadata = [&](int dim, StringRef name) {
       llvm::Metadata *llvmMetadata[] = {
           llvm::ValueAsMetadata::get(llvmFunc),
+          llvm::MDString::get(llvmContext, name),
+          llvm::ValueAsMetadata::get(llvm::ConstantInt::get(
+              llvm::Type::getInt32Ty(llvmContext), dim))};
+      llvm::MDNode *llvmMetadataNode =
+          llvm::MDNode::get(llvmContext, llvmMetadata);
+      moduleTranslation.getOrInsertNamedModuleMetadata("nvvm.annotations")
+          ->addOperand(llvmMetadataNode);
+    };
+    if (attribute.getName() == NVVM::NVVMDialect::getMaxntidAttrName()) {
+      if (!attribute.getValue().dyn_cast<ArrayAttr>())
+        return failure();
+      SmallVector<int64_t> values =
+          extractFromI64ArrayAttr(attribute.getValue());
+      if (!values.empty())
+        generateMetadata(values[0], NVVM::NVVMDialect::getMaxntidXName());
+      if (values.size() > 1)
+        generateMetadata(values[1], NVVM::NVVMDialect::getMaxntidYName());
+      if (values.size() > 2)
+        generateMetadata(values[2], NVVM::NVVMDialect::getMaxntidZName());
+    } else if (attribute.getName() == NVVM::NVVMDialect::getReqntidAttrName()) {
+      if (!attribute.getValue().dyn_cast<ArrayAttr>())
+        return failure();
+      SmallVector<int64_t> values =
+          extractFromI64ArrayAttr(attribute.getValue());
+      if (!values.empty())
+        generateMetadata(values[0], NVVM::NVVMDialect::getReqntidXName());
+      if (values.size() > 1)
+        generateMetadata(values[1], NVVM::NVVMDialect::getReqntidYName());
+      if (values.size() > 2)
+        generateMetadata(values[2], NVVM::NVVMDialect::getReqntidZName());
+    } else if (attribute.getName() ==
+               NVVM::NVVMDialect::getMinctasmAttrName()) {
+      auto value = attribute.getValue().dyn_cast<IntegerAttr>();
+      if (!value)
+        return success();
+      generateMetadata(value.getInt(), "minctasm");
+    } else if (attribute.getName() == NVVM::NVVMDialect::getMaxnregAttrName()) {
+      auto value = attribute.getValue().dyn_cast<IntegerAttr>();
+      if (!value)
+        return success();
+      generateMetadata(value.getInt(), "maxnreg");
+    } else if (attribute.getName() ==
+               NVVM::NVVMDialect::getKernelFuncAttrName()) {
+      llvm::Metadata *llvmMetadataKernel[] = {
+          llvm::ValueAsMetadata::get(llvmFunc),
           llvm::MDString::get(llvmContext, "kernel"),
           llvm::ValueAsMetadata::get(
               llvm::ConstantInt::get(llvm::Type::getInt32Ty(llvmContext), 1))};
       llvm::MDNode *llvmMetadataNode =
-          llvm::MDNode::get(llvmContext, llvmMetadata);
+          llvm::MDNode::get(llvmContext, llvmMetadataKernel);
       moduleTranslation.getOrInsertNamedModuleMetadata("nvvm.annotations")
           ->addOperand(llvmMetadataNode);
     }
     return success();
   }
 };
 } // namespace
 

mlir/test/Target/LLVMIR/nvvmir.mlir

-// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+// RUN: mlir-translate -mlir-to-llvmir %s  -split-input-file | FileCheck %s
 
 // CHECK-LABEL: @nvvm_special_regs
 llvm.func @nvvm_special_regs() -> i32 {
   // CHECK: %1 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
   %1 = nvvm.read.ptx.sreg.tid.x : i32
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.tid.y()
   %2 = nvvm.read.ptx.sreg.tid.y : i32
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.tid.z()

 // NVVM annotations after conversion.
 llvm.func @kernel_func() attributes {nvvm.kernel} {
   llvm.return
 }
 
 // CHECK:     !nvvm.annotations =
 // CHECK-NOT: {ptr @nvvm_special_regs, !"kernel", i32 1}
 // CHECK:     {ptr @kernel_func, !"kernel", i32 1}
+
+// -----
+
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.maxntid = [1,23,32]} {
+  llvm.return
+}
+
+// CHECK:     !nvvm.annotations =
+// CHECK-NOT: {ptr @nvvm_special_regs, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"maxntidx", i32 1}
+// CHECK:     {ptr @kernel_func, !"maxntidy", i32 23}
+// CHECK:     {ptr @kernel_func, !"maxntidz", i32 32}
+// -----
+
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.reqntid = [1,23,32]} {
+  llvm.return
+}
+
+// CHECK:     !nvvm.annotations =
+// CHECK-NOT: {ptr @nvvm_special_regs, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"reqntidx", i32 1}
+// CHECK:     {ptr @kernel_func, !"reqntidy", i32 23}
+// CHECK:     {ptr @kernel_func, !"reqntidz", i32 32}
+// -----
+
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.minctasm = 16} {
+  llvm.return
+}
+
+// CHECK:     !nvvm.annotations =
+// CHECK-NOT: {ptr @nvvm_special_regs, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"minctasm", i32 16}
+// -----
+
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.maxnreg = 16} {
+  llvm.return
+}
+
+// CHECK:     !nvvm.annotations =
+// CHECK-NOT: {ptr @nvvm_special_regs, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"maxnreg", i32 16}
+// -----
+
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.maxntid = [1,23,32], 
+                                     nvvm.minctasm = 16, nvvm.maxnreg = 32} {
+  llvm.return
+}
+
+// CHECK:     !nvvm.annotations =
+// CHECK-NOT: {ptr @nvvm_special_regs, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"kernel", i32 1}
+// CHECK:     {ptr @kernel_func, !"maxnreg", i32 32}
+// CHECK:     {ptr @kernel_func, !"maxntidx", i32 1}
+// CHECK:     {ptr @kernel_func, !"maxntidy", i32 23}
+// CHECK:     {ptr @kernel_func, !"maxntidz", i32 32}
+// CHECK:     {ptr @kernel_func, !"minctasm", i32 16}
+