xref: /llvm-project/mlir/include/mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.h (revision 25da11541c0740950e812a85d332a6d5cf4a5f87)
1 //===- BasicPtxBuilderInterface.td - PTX builder interface -*- 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 // Defines the interface to build PTX (Parallel Thread Execution) from NVVM Ops
10 // automatically. It is used by NVVM to LLVM pass.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef NVVM_DIALECT_NVVM_IR_BASICPTXBUILDERINTERFACE_H_
15 #define NVVM_DIALECT_NVVM_IR_BASICPTXBUILDERINTERFACE_H_
16 
17 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
18 #include "mlir/IR/BuiltinAttributes.h"
19 #include "mlir/IR/PatternMatch.h"
20 #include "mlir/IR/Value.h"
21 
22 namespace mlir {
23 namespace NVVM {
24 /// Register read/write modifier to build constraint string for PTX inline
25 /// https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html#parameters
26 enum class PTXRegisterMod {
27   /// Read register with no modifier
28   Read = 0,
29   /// Read register with '+' modifier
30   Write = 2,
31   /// Read register with '=' modifier.
32   /// Note that, this is not natively supported by LLVM, but it is possible to
33   /// set read and write for the same operand.
34   ReadWrite = 1,
35 };
36 } // namespace NVVM
37 } // namespace mlir
38 
39 /// Include the generated interface declarations.
40 #include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.h.inc"
41 
42 namespace mlir {
43 
44 namespace NVVM {
45 
46 /// A class to build PTX assembly automatically. It is used by
47 /// BasicPtxBuilderInterface.
48 class PtxBuilder {
49   // The interface op that is used to build the PTX.
50   BasicPtxBuilderInterface interfaceOp;
51   // Rewriter to create new operations.
52   PatternRewriter &rewriter;
53   // The operands for the PTX instruction
54   SmallVector<Value> ptxOperands;
55   // Register constraints (read, write, readwrite) and register data types
56   std::string registerConstraints;
57 
58   bool hasResult = false;
59 
60 public:
61   /// Single constructor that only initializes members.
PtxBuilder(Operation * op,PatternRewriter & rewriter)62   PtxBuilder(Operation *op, PatternRewriter &rewriter)
63       : interfaceOp(op), rewriter(rewriter) {}
64 
65   /// Add an operand with the read/write input type.
66   void insertValue(Value v, PTXRegisterMod itype = PTXRegisterMod::Read);
67 
68   /// Builds the inline assembly Op and returns it. The `insertValue` needs to
69   /// be called to pass operands before building the PTX.
70   LLVM::InlineAsmOp build();
71 
72   /// Shortcut to build the inline assembly Op and replace or erase the original
73   /// op with
74   void buildAndReplaceOp();
75 };
76 
77 } // namespace NVVM
78 } // namespace mlir
79 
80 #endif // NVVM_DIALECT_NVVM_IR_BASICPTXBUILDERINTERFACE_H_
81