[mlir][nvgpu] Add NVGPU dialect (architectural specific gpu dialect)

[ThomasRaoux]

Summary

This introduce a new dialect for vendor specific ptx operations. This
also adds the first operation ldmatrix as an example. More operations
will be added in follow up patches.
This new dialect is meant to be a bridge between GPU and Vector
dialectic and NVVM dialect.

This is based on the RFC proposed here:
https://discourse.llvm.org/t/rfc-add-nv-gpu-dialect-hw-specific-extension-of-gpu-dialect-for-nvidia-gpus/61466

[mravishankar]

Looks good to me. I am happy to stamp, but will weight for folks to weigh in.

[bondhugula]

In D123266#3434993, @mravishankar wrote:

Looks good to me. I am happy to stamp, but will weight for folks to weigh in.

I'd like to see more discussion on this -- posted some questions here: https://discourse.llvm.org/t/rfc-add-nv-gpu-dialect-hw-specific-extension-of-gpu-dialect-for-nvidia-gpus/61466/10?u=bondhugula

[ThomasRaoux]

Rename dialect to nvgpu

[ThomasRaoux]

In D123266#3435213, @bondhugula wrote:

In D123266#3434993, @mravishankar wrote:

Looks good to me. I am happy to stamp, but will weight for folks to weigh in.

I'd like to see more discussion on this -- posted some questions here: https://discourse.llvm.org/t/rfc-add-nv-gpu-dialect-hw-specific-extension-of-gpu-dialect-for-nvidia-gpus/61466/10?u=bondhugula

I renamed the dialect as suggested. If there are any fundamental points you think you should discussed please bring it up on discourse or feel free to comment on more details case on the review.

[herhut]

Looks good to me. We really need to figure out a way to group dialects :)

Please also wait for @bondhugula, who had concerns.

[ThomasRaoux]

In D123266#3448574, @herhut wrote:

Looks good to me. We really need to figure out a way to group dialects :)

Please also wait for @bondhugula, who had concerns.

Thanks @herhut. @bondhugula, do you still have any concerns?

[bondhugula]

In D123266#3448690, @ThomasRaoux wrote:

In D123266#3448574, @herhut wrote:

Looks good to me. We really need to figure out a way to group dialects :)

Please also wait for @bondhugula, who had concerns.

Thanks @herhut. @bondhugula, do you still have any concerns?

This looks fine to me. It was a matter of time before a dialect like this was created. We still have to be cautious about deciding what goes into the GPU dialect vs NVGPU dialect and the lowering paths for the ops that are added here.

[ThomasRaoux]

We still have to be cautious about deciding what goes into the GPU dialect vs NVGPU dialect and the lowering paths for the ops that are added here.

Yes, I definitely agree.

[nicolasvasilache]

mlir/include/mlir/Dialect/NVGPU/NVGPU.td

65

Is numTiles the same as the .num attribute in the PTX ISA doc ?

66

The PTX doc specifically mentions 16b elements, did you want to tighten the type here or allow more relaxed semantics with an implicit bitcast and make the verifier only check the final bitlength?

Hmm actually what about the fact that the shape seems to be prescribed to exactly 8x8xf16?
Do you want the op to model exactly that or relax it?

[christopherbate]

mlir/include/mlir/Dialect/NVGPU/NVGPU.td

65

Yes, but I thought numTiles was more descriptive, see my response to the other question below.

66

I originally authored this code, which was already merged in D123647; here it is code movement.

It is meant to be relaxed vs the stated requirement of 8x8xf16. You can restate the 8x8xf16 tile specification as 8x(4x32b) or 8x16Byte tiles and functionally it will work out. In fact, the NVVM intrinsic ldmatrix in the backend returns i32 values which then need to be bit casted into 2xf16 or 4xi8, etc. We have the tests covering all those cases in the NVVM dialect, but we do need to follow up here with a verifier for this operation.

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(EmitC)
 add_subdirectory(Func)
 add_subdirectory(GPU)
 add_subdirectory(Math)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)
 add_subdirectory(MLProgram)
+add_subdirectory(NVGPU)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Tensor)
 add_subdirectory(Tosa)
 add_subdirectory(Transform)
 add_subdirectory(Vector)
 add_subdirectory(X86Vector)

mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt

+add_mlir_dialect(NVGPU nvgpu)
+add_mlir_doc(NVGPU -gen-dialect-doc NVGPU Dialects/)
+
+set(LLVM_TARGET_DEFINITIONS NVGPU.td)

mlir/include/mlir/Dialect/NVGPU/NVGPU.td

+//===-- NVGPU.td - NVGPU dialect operation definitions *- tablegen -*------===//
+//
+// 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 defines the basic operations for the NVGPU dialect.
+//
+// This NVGPU provides a bridge between the target agnostic GPU and Vector
+// dialects and lower level NVVM dialect. This allow representing PTX specific
+// operations while using MLIR high level concepts like memref and 2-D vector.
+//
+// Ops semantic are going to be based on vendor specific PTX defintion:
+// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef NVGPU
+#define NVGPU
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
+
+def NVGPU_Dialect : Dialect {
+  let name = "nvgpu";
+  let cppNamespace = "::mlir::nvgpu";
+  let description = [{
+    This `NVGPU` dialect provides a bridge between the target agnostic GPU and
+    Vector dialects and the lower level LLVM IR based NVVM dialect. This allow
+    representing PTX specific operations while using MLIR high level concepts
+    like memref and 2-D vector.
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// NVGPU Op definitions
+//===----------------------------------------------------------------------===//
+
+class NVGPU_Op<string mnemonic, list<Trait> traits = []> :
+  Op<NVGPU_Dialect, mnemonic, traits> {}
+
+def NVGPU_LdMatrixOp : NVGPU_Op<"ldmatrix",
+                                [MemoryEffects<[MemRead]>]> {
+  let description = [{
+  The `nvgpu.ldmatrix` op represents loading a matrix fragment from
+  memory. The load source and result type must be compatible with lowering
+  to the `nvvm.ldmatrix` instruction. This op is meant to represent
+  the distributed version of a `vector.transfer_read` as an intermediate
+  step between lowering from `vector.transfer_read` to `nvvm.ldmatrix`.
+
+  This operation is meant to follow the semantic of described here:
+  https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix
+
+  Example:
+  ```mlir
+  %0 = nvgpu.ldmatrix %sm[%c0, %c0] {numTiles = 4 : i32, transpose = false} :
+    memref<?x?xf16, 3> -> vector<4x2xf16>
+  ```
+  }];
+
+  let arguments = (ins Arg<AnyMemRef, "", [MemRead]>:$srcMemref,
+                           Variadic<Index>:$indices, BoolAttr:$transpose,
+                           I32Attr:$numTiles);
+  let results = (outs AnyVector:$res);
+  let assemblyFormat = [{
+    $srcMemref`[` $indices `]` attr-dict `:` type($srcMemref) `->` type($res)
+  }];
+}
+
+#endif // NVGPU

mlir/include/mlir/Dialect/NVGPU/NVGPUDialect.h

+//===- NVGPUDialect.h - MLIR Dialect for NVGPU ------------------*- C++ -*-===//
+//
+// 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 declares the Target dialect for NVGPU in MLIR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_
+#define MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "mlir/Dialect/NVGPU/NVGPUDialect.h.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/NVGPU/NVGPU.h.inc"
+
+#endif // MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/NVGPU/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
registry.insert<acc::OpenACCDialect,
                   emitc::EmitCDialect,
                   func::FuncDialect,
                   gpu::GPUDialect,
                   LLVM::LLVMDialect,
                   linalg::LinalgDialect,
                   math::MathDialect,
                   memref::MemRefDialect,
                   ml_program::MLProgramDialect,
+                  nvgpu::NVGPUDialect,
                   scf::SCFDialect,
                   omp::OpenMPDialect,
                   pdl::PDLDialect,
                   pdl_interp::PDLInterpDialect,
                   quant::QuantizationDialect,
                   spirv::SPIRVDialect,
                   arm_sve::ArmSVEDialect,
                   vector::VectorDialect,

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(EmitC)
 add_subdirectory(Func)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(Math)
 add_subdirectory(MemRef)
 add_subdirectory(MLProgram)
+add_subdirectory(NVGPU)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SparseTensor)

mlir/lib/Dialect/NVGPU/CMakeLists.txt

+add_subdirectory(IR)

mlir/lib/Dialect/NVGPU/IR/CMakeLists.txt

+add_mlir_dialect_library(MLIRNVGPU
+  NVGPUDialect.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/NVGPU
+
+  DEPENDS
+  MLIRNVGPUIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRSideEffectInterfaces
+  )

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

+//===- NVGPUDialect.cpp - MLIR NVGPU ops implementation -------------------===//
+//
+// 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 the NVGPU dialect and its operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/NVGPU/NVGPUDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/TypeUtilities.h"
+
+using namespace mlir;
+
+#include "mlir/Dialect/NVGPU/NVGPUDialect.cpp.inc"
+
+void nvgpu::NVGPUDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/NVGPU/NVGPU.cpp.inc"
+      >();
+}
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/NVGPU/NVGPU.cpp.inc"

mlir/test/Dialect/NVGPU/roundtrip.mlir

+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// CHECK-LABEL: func @ldmatrix(
+func @ldmatrix(%arg0: memref<?x?xf16, 3>, %x: index, %y: index) {
+//      CHECK: nvgpu.ldmatrix %{{.*}}[%{{.*}}, %{{.*}}]
+// CHECK-SAME: {numTiles = 4 : i32, transpose = false} : memref<?x?xf16, 3> -> vector<4x2xf16>
+  %l = nvgpu.ldmatrix %arg0[%x, %y] {numTiles = 4 : i32, transpose = false} :
+    memref<?x?xf16, 3> -> vector<4x2xf16>
+  return
+}

mlir/test/mlir-opt/commandline.mlir

 // CHECK-NEXT: emitc
 // CHECK-NEXT: func
 // CHECK-NEXT: gpu
 // CHECK-NEXT: linalg
 // CHECK-NEXT: llvm
 // CHECK-NEXT: math
 // CHECK-NEXT: memref
 // CHECK-NEXT: ml_program
+// CHECK-NEXT: nvgpu
 // CHECK-NEXT: nvvm
 // CHECK-NEXT: omp
 // CHECK-NEXT: pdl
 // CHECK-NEXT: pdl_interp
 // CHECK-NEXT: quant
 // CHECK-NEXT: rocdl
 // CHECK-NEXT: scf
 // CHECK-NEXT: shape
 // CHECK-NEXT: sparse_tensor
 // CHECK-NEXT: spv
 // CHECK-NEXT: tensor
 // CHECK-NEXT: test
 // CHECK-NEXT: tosa
 // CHECK-NEXT: transform
 // CHECK-NEXT: vector
 // CHECK-NEXT: x86vector

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":SparseTensor",
         ":SparseTensorTransforms",
         ":TensorTransforms",
         ":VectorToLLVM",
         ":VectorTransforms",
     ],
 )
 
+##---------------------------------------------------------------------------##
+# NVGPU dialect.
+##---------------------------------------------------------------------------##
+
+td_library(
+    name = "NVGPUTdFiles",
+    srcs = ["include/mlir/Dialect/NVGPU/NVGPU.td"],
+    includes = ["include"],
+    deps = [
+        ":SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "NVGPUIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            [
+                "-gen-dialect-decls",
+                "-dialect=nvgpu",
+            ],
+            "include/mlir/Dialect/NVGPU/NVGPUDialect.h.inc",
+        ),
+        (
+            [
+                "-gen-dialect-defs",
+                "-dialect=nvgpu",
+            ],
+            "include/mlir/Dialect/NVGPU/NVGPUDialect.cpp.inc",
+        ),
+        (
+            ["-gen-op-decls"],
+            "include/mlir/Dialect/NVGPU/NVGPU.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/mlir/Dialect/NVGPU/NVGPU.cpp.inc",
+        ),
+        (
+            ["-gen-op-doc"],
+            "g3doc/Dialects/NVGPU/NVGPU.md",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/NVGPU/NVGPU.td",
+    deps = [":NVGPUTdFiles"],
+)
+
+cc_library(
+    name = "NVGPU",
+    srcs = ["lib/Dialect/NVGPU/IR/NVGPUDialect.cpp"],
+    hdrs = ["include/mlir/Dialect/NVGPU/NVGPUDialect.h"],
+    includes = ["include"],
+    deps = [
+        ":IR",
+        ":NVGPUIncGen",
+        ":SideEffectInterfaces",
+        "//llvm:Core",
+        "//llvm:Support",
+    ],
+)
+
 td_library(
     name = "FuncTdFiles",
     srcs = [
         "include/mlir/Dialect/Func/IR/FuncOps.td",
     ],
     includes = ["include"],
     deps = [
         ":CallInterfacesTdFiles",
deps = [
         ":MathToLLVM",
         ":MathToLibm",
         ":MathToSPIRV",
         ":MathTransforms",
         ":MemRefDialect",
         ":MemRefToLLVM",
         ":MemRefToSPIRV",
         ":MemRefTransforms",
+        ":NVGPU",
         ":NVVMDialect",
         ":OpenACCDialect",
         ":OpenMPDialect",
         ":OpenMPToLLVM",
         ":PDLDialect",
         ":PDLInterpDialect",
         ":PDLToPDLInterp",
         ":QuantOps",