MLIR  21.0.0git
NVVMDialect.h
Go to the documentation of this file.
1 //===- NVVMDialect.h - MLIR NVVM IR dialect ---------------------*- 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 defines the NVVM IR dialect in MLIR, containing NVVM operations and
10 // NVVM specific extensions to the LLVM type system.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_DIALECT_LLVMIR_NVVMDIALECT_H_
15 #define MLIR_DIALECT_LLVMIR_NVVMDIALECT_H_
16 
22 #include "mlir/IR/Dialect.h"
23 #include "mlir/IR/OpDefinition.h"
27 #include "llvm/IR/IntrinsicsNVPTX.h"
28 
29 #include "mlir/Dialect/LLVMIR/NVVMOpsEnums.h.inc"
30 
31 namespace mlir {
32 namespace NVVM {
33 
34 // Shared memory has 128-bit alignment
35 constexpr int kSharedMemoryAlignmentBit = 128;
36 
37 /// NVVM memory space identifiers.
39  /// Generic memory space identifier.
41  /// Global memory space identifier.
43  /// Shared memory space identifier.
45  /// Constant memory space identifier.
47  /// Local memory space identifier.
49  /// Tensor memory space identifier.
50  /// Tensor memory is available only in arch-accelerated
51  /// variants from sm100 onwards.
53  /// Distributed shared memory space identifier.
54  /// Distributed shared memory is available only in sm90+.
56 };
57 
58 /// A pair type of LLVM's Intrinsic ID and args (which are llvm values).
59 /// This type is returned by the getIntrinsicIDAndArgs() methods.
60 using IDArgPair =
61  std::pair<llvm::Intrinsic::ID, llvm::SmallVector<llvm::Value *>>;
62 
63 /// Return the element type and number of elements associated with a wmma matrix
64 /// of given chracteristics. This matches the logic in IntrinsicsNVVM.td
65 /// WMMA_REGS structure.
66 std::pair<mlir::Type, unsigned> inferMMAType(mlir::NVVM::MMATypes type,
67  mlir::NVVM::MMAFrag frag, int nRow,
68  int nCol,
69  mlir::MLIRContext *context);
70 } // namespace NVVM
71 } // namespace mlir
72 
73 ///// Ops /////
74 #define GET_ATTRDEF_CLASSES
75 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.h.inc"
76 
77 #define GET_OP_CLASSES
78 #include "mlir/Dialect/LLVMIR/NVVMOps.h.inc"
79 
80 #include "mlir/Dialect/LLVMIR/NVVMOpsDialect.h.inc"
81 
82 #endif /* MLIR_DIALECT_LLVMIR_NVVMDIALECT_H_ */
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
std::pair< llvm::Intrinsic::ID, llvm::SmallVector< llvm::Value * > > IDArgPair
A pair type of LLVM's Intrinsic ID and args (which are llvm values).
Definition: NVVMDialect.h:61
constexpr int kSharedMemoryAlignmentBit
Definition: NVVMDialect.h:35
NVVMMemorySpace
NVVM memory space identifiers.
Definition: NVVMDialect.h:38
@ kGenericMemorySpace
Generic memory space identifier.
Definition: NVVMDialect.h:40
@ kGlobalMemorySpace
Global memory space identifier.
Definition: NVVMDialect.h:42
@ kConstantMemorySpace
Constant memory space identifier.
Definition: NVVMDialect.h:46
@ kSharedClusterMemorySpace
Distributed shared memory space identifier.
Definition: NVVMDialect.h:55
@ kLocalMemorySpace
Local memory space identifier.
Definition: NVVMDialect.h:48
@ kTensorMemorySpace
Tensor memory space identifier.
Definition: NVVMDialect.h:52
@ kSharedMemorySpace
Shared memory space identifier.
Definition: NVVMDialect.h:44
std::pair< mlir::Type, unsigned > inferMMAType(mlir::NVVM::MMATypes type, mlir::NVVM::MMAFrag frag, int nRow, int nCol, mlir::MLIRContext *context)
Return the element type and number of elements associated with a wmma matrix of given chracteristics.
Include the generated interface declarations.