[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/Dialect/LLVMIR/IR/NVVMDialect.cpp

 //
 // The NVVM dialect only contains GPU specific additions on top of the general
 // LLVM dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/OperationSupport.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Type.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 using namespace NVVM;
 
 #include "mlir/Dialect/LLVMIR/NVVMOpsDialect.cpp.inc"
 #include "mlir/Dialect/LLVMIR/NVVMOpsEnums.cpp.inc"
 
#include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
 
   // Support unknown operations because not all NVVM operations are
   // registered.
   allowUnknownOperations();
 }
 
 LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
                                                     NamedAttribute attr) {
+  StringAttr attrName = attr.getName();
   // Kernel function attribute should be attached to functions.
-  if (attr.getName() == NVVMDialect::getKernelFuncAttrName()) {
+  if (attrName == NVVMDialect::getKernelFuncAttrName()) {
     if (!isa<LLVM::LLVMFuncOp>(op)) {
       return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
                              << "' attribute attached to unexpected op";
     }
   }
+  // If maxntid and reqntid exist, it must be an array with max 3 dim
+  if (attrName == NVVMDialect::getMaxntidAttrName() ||
+      attrName == NVVMDialect::getReqntidAttrName()) {
+    auto values = attr.getValue().dyn_cast<ArrayAttr>();
+    if (!values || values.empty() || values.size() > 3)
+      return op->emitError()
+             << "'" << attrName
+             << "' attribute must be integer array with maximum 3 index";
+    for (auto val : attr.getValue().cast<ArrayAttr>()) {
+      if (!val.dyn_cast<IntegerAttr>())
+        return op->emitError()
+               << "'" << attrName
+               << "' attribute must be integer array with maximum 3 index";
+    }
+  }
+  // If minctasm and maxnreg exist, it must be an array with max 3 dim
+  if (attrName == NVVMDialect::getMinctasmAttrName() ||
+      attrName == NVVMDialect::getMaxnregAttrName()) {
+    if (!attr.getValue().dyn_cast<IntegerAttr>())
+      return op->emitError()
+             << "'" << attrName << "' attribute must be integer constant";
+  }
+
   return success();
 }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"

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/Support/LogicalResult.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 using mlir::LLVM::detail::createIntrinsicCall;
#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());
+      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());
+      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>();
+      generateMetadata(value.getInt(), "minctasm");
+    } else if (attribute.getName() == NVVM::NVVMDialect::getMaxnregAttrName()) {
+      auto value = attribute.getValue().dyn_cast<IntegerAttr>();
+      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 --verify-diagnostics | 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}
+
+// -----
+// expected-error @below {{'"nvvm.minctasm"' attribute must be integer constant}}  
+llvm.func @kernel_func() attributes {nvvm.kernel, 
+nvvm.minctasm = "foo"} {
+  llvm.return
+}
+
+
+// -----
+// expected-error @below {{'"nvvm.maxnreg"' attribute must be integer constant}}  
+llvm.func @kernel_func() attributes {nvvm.kernel, 
+nvvm.maxnreg = "boo"} {
+  llvm.return
+}
+// -----
+// expected-error @below {{'"nvvm.reqntid"' attribute must be integer array with maximum 3 index}}  
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.reqntid = [3,4,5,6]} {
+  llvm.return
+}
+
+// -----
+// expected-error @below {{'"nvvm.maxntid"' attribute must be integer array with maximum 3 index}}  
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.maxntid = [3,4,5,6]} {
+  llvm.return
+}
+