Skip to content

Commit 4c56494

Browse files
ThomasRaouxThomasRaoux
committed
[mlir][nvgpu] Add NVGPU dialect (architectural specific gpu dialect)
This introduce a new dialect for vendro 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 dialectis 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/8 Differential Revision: https://reviews.llvm.org/D123266
1 parent f14ebe9 commit 4c56494

File tree

12 files changed

+225
-0
lines changed

12 files changed

+225
-0
lines changed

‎mlir/include/mlir/Dialect/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ add_subdirectory(Linalg)
1616
add_subdirectory(LLVMIR)
1717
add_subdirectory(MemRef)
1818
add_subdirectory(MLProgram)
19+
add_subdirectory(NVGPU)
1920
add_subdirectory(OpenACC)
2021
add_subdirectory(OpenMP)
2122
add_subdirectory(PDL)

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
add_mlir_dialect(NVGPU nvgpu)
2+
add_mlir_doc(NVGPU -gen-dialect-doc NVGPU Dialects/)
3+
4+
set(LLVM_TARGET_DEFINITIONS NVGPU.td)

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
//===-- NVGPU.td - NVGPU dialect operation definitions *- tablegen -*------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file defines the basic operations for the NVGPU dialect.
10+
//
11+
// This NVGPU provides a bridge between the target agnostic GPU and Vector
12+
// dialects and lower level NVVM dialect. This allow representing PTX specific
13+
// operations while using MLIR high level concepts like memref and 2-D vector.
14+
//
15+
// Ops semantic are going to be based on vendor specific PTX defintion:
16+
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html
17+
//
18+
//===----------------------------------------------------------------------===//
19+
20+
#ifndef NVGPU
21+
#define NVGPU
22+
23+
include "mlir/Interfaces/SideEffectInterfaces.td"
24+
include "mlir/IR/OpBase.td"
25+
26+
def NVGPU_Dialect : Dialect {
27+
let name = "nvgpu";
28+
let cppNamespace = "::mlir::nvgpu";
29+
let description = [{
30+
This `NVGPU` dialect provides a bridge between the target agnostic GPU and
31+
Vector dialects and the lower level LLVM IR based NVVM dialect. This allow
32+
representing PTX specific operations while using MLIR high level concepts
33+
like memref and 2-D vector.
34+
}];
35+
}
36+
37+
//===----------------------------------------------------------------------===//
38+
// NVGPU Op definitions
39+
//===----------------------------------------------------------------------===//
40+
41+
class NVGPU_Op<string mnemonic, list<Trait> traits = []> :
42+
Op<NVGPU_Dialect, mnemonic, traits> {}
43+
44+
def NVGPU_LdMatrixOp : NVGPU_Op<"ldmatrix",
45+
[MemoryEffects<[MemRead]>]> {
46+
let description = [{
47+
The `nvgpu.ldmatrix` op represents loading a matrix fragment from
48+
memory. The load source and result type must be compatible with lowering
49+
to the `nvvm.ldmatrix` instruction. This op is meant to represent
50+
the distributed version of a `vector.transfer_read` as an intermediate
51+
step between lowering from `vector.transfer_read` to `nvvm.ldmatrix`.
52+
53+
This operation is meant to follow the semantic of described here:
54+
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix
55+
56+
Example:
57+
```mlir
58+
%0 = nvgpu.ldmatrix %sm[%c0, %c0] {numTiles = 4 : i32, transpose = false} :
59+
memref<?x?xf16, 3> -> vector<4x2xf16>
60+
```
61+
}];
62+
63+
let arguments = (ins Arg<AnyMemRef, "", [MemRead]>:$srcMemref,
64+
Variadic<Index>:$indices, BoolAttr:$transpose,
65+
I32Attr:$numTiles);
66+
let results = (outs AnyVector:$res);
67+
let assemblyFormat = [{
68+
$srcMemref`[` $indices `]` attr-dict `:` type($srcMemref) `->` type($res)
69+
}];
70+
}
71+
72+
#endif // NVGPU

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
//===- NVGPUDialect.h - MLIR Dialect for NVGPU ------------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file declares the Target dialect for NVGPU in MLIR.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_
14+
#define MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_
15+
16+
#include "mlir/IR/BuiltinTypes.h"
17+
#include "mlir/IR/Dialect.h"
18+
#include "mlir/IR/OpDefinition.h"
19+
#include "mlir/Interfaces/SideEffectInterfaces.h"
20+
21+
#include "mlir/Dialect/NVGPU/NVGPUDialect.h.inc"
22+
23+
#define GET_OP_CLASSES
24+
#include "mlir/Dialect/NVGPU/NVGPU.h.inc"
25+
26+
#endif // MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_

‎mlir/include/mlir/InitAllDialects.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
3737
#include "mlir/Dialect/Math/IR/Math.h"
3838
#include "mlir/Dialect/MemRef/IR/MemRef.h"
39+
#include "mlir/Dialect/NVGPU/NVGPUDialect.h"
3940
#include "mlir/Dialect/OpenACC/OpenACC.h"
4041
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
4142
#include "mlir/Dialect/PDL/IR/PDL.h"
@@ -80,6 +81,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
8081
math::MathDialect,
8182
memref::MemRefDialect,
8283
ml_program::MLProgramDialect,
84+
nvgpu::NVGPUDialect,
8385
scf::SCFDialect,
8486
omp::OpenMPDialect,
8587
pdl::PDLDialect,

‎mlir/lib/Dialect/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ add_subdirectory(LLVMIR)
1616
add_subdirectory(Math)
1717
add_subdirectory(MemRef)
1818
add_subdirectory(MLProgram)
19+
add_subdirectory(NVGPU)
1920
add_subdirectory(OpenACC)
2021
add_subdirectory(OpenMP)
2122
add_subdirectory(PDL)

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
add_subdirectory(IR)

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
add_mlir_dialect_library(MLIRNVGPU
2+
NVGPUDialect.cpp
3+
4+
ADDITIONAL_HEADER_DIRS
5+
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/NVGPU
6+
7+
DEPENDS
8+
MLIRNVGPUIncGen
9+
10+
LINK_LIBS PUBLIC
11+
MLIRIR
12+
MLIRSideEffectInterfaces
13+
)

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
//===- NVGPUDialect.cpp - MLIR NVGPU ops implementation -------------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file implements the NVGPU dialect and its operations.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#include "mlir/Dialect/NVGPU/NVGPUDialect.h"
14+
#include "mlir/IR/Builders.h"
15+
#include "mlir/IR/OpImplementation.h"
16+
#include "mlir/IR/TypeUtilities.h"
17+
18+
using namespace mlir;
19+
20+
#include "mlir/Dialect/NVGPU/NVGPUDialect.cpp.inc"
21+
22+
void nvgpu::NVGPUDialect::initialize() {
23+
addOperations<
24+
#define GET_OP_LIST
25+
#include "mlir/Dialect/NVGPU/NVGPU.cpp.inc"
26+
>();
27+
}
28+
29+
#define GET_OP_CLASSES
30+
#include "mlir/Dialect/NVGPU/NVGPU.cpp.inc"

‎mlir/test/Dialect/NVGPU/roundtrip.mlir‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
// RUN: mlir-opt %s | mlir-opt | FileCheck %s
2+
3+
// CHECK-LABEL: func @ldmatrix(
4+
func @ldmatrix(%arg0: memref<?x?xf16, 3>, %x: index, %y: index) {
5+
// CHECK: nvgpu.ldmatrix %{{.*}}[%{{.*}}, %{{.*}}]
6+
// CHECK-SAME: {numTiles = 4 : i32, transpose = false} : memref<?x?xf16, 3> -> vector<4x2xf16>
7+
%l = nvgpu.ldmatrix %arg0[%x, %y] {numTiles = 4 : i32, transpose = false} :
8+
memref<?x?xf16, 3> -> vector<4x2xf16>
9+
return
10+
}

0 commit comments

Comments
 (0)