1 | //===- IndexIntrinsicsOpLowering.h - GPU IndexOps Lowering class *- 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 | #ifndef MLIR_CONVERSION_GPUCOMMON_INDEXINTRINSICSOPLOWERING_H_ |
9 | #define MLIR_CONVERSION_GPUCOMMON_INDEXINTRINSICSOPLOWERING_H_ |
10 | |
11 | #include "mlir/Conversion/LLVMCommon/Pattern.h" |
12 | #include "mlir/Dialect/GPU/IR/GPUDialect.h" |
13 | #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
14 | #include "mlir/IR/BuiltinAttributes.h" |
15 | |
16 | namespace mlir { |
17 | |
18 | // Rewriting that replaces Op with XOp, YOp, or ZOp depending on the dimension |
19 | // that Op operates on. Op is assumed to return an `index` value and |
20 | // XOp, YOp and ZOp are assumed to return an `llvm.i32` value. Depending on |
21 | // `indexBitwidth`, sign-extend or truncate the resulting value to match the |
22 | // bitwidth expected by the consumers of the value. |
23 | template <typename Op, typename XOp, typename YOp, typename ZOp> |
24 | struct GPUIndexIntrinsicOpLowering : public ConvertOpToLLVMPattern<Op> { |
25 | private: |
26 | unsigned indexBitwidth; |
27 | StringRef boundsAttrName; |
28 | |
29 | public: |
30 | explicit GPUIndexIntrinsicOpLowering(LLVMTypeConverter &typeConverter) |
31 | : ConvertOpToLLVMPattern<Op>(typeConverter), |
32 | indexBitwidth(typeConverter.getIndexTypeBitwidth()), |
33 | boundsAttrName("" ) {} |
34 | |
35 | explicit GPUIndexIntrinsicOpLowering(LLVMTypeConverter &typeConverter, |
36 | StringRef boundsAttrName) |
37 | : ConvertOpToLLVMPattern<Op>(typeConverter), |
38 | indexBitwidth(typeConverter.getIndexTypeBitwidth()), |
39 | boundsAttrName(boundsAttrName) {} |
40 | |
41 | // Convert the kernel arguments to an LLVM type, preserve the rest. |
42 | LogicalResult |
43 | matchAndRewrite(Op op, typename Op::Adaptor adaptor, |
44 | ConversionPatternRewriter &rewriter) const override { |
45 | auto loc = op->getLoc(); |
46 | MLIRContext *context = rewriter.getContext(); |
47 | Operation *newOp; |
48 | switch (op.getDimension()) { |
49 | case gpu::Dimension::x: |
50 | newOp = rewriter.create<XOp>(loc, IntegerType::get(context, 32)); |
51 | break; |
52 | case gpu::Dimension::y: |
53 | newOp = rewriter.create<YOp>(loc, IntegerType::get(context, 32)); |
54 | break; |
55 | case gpu::Dimension::z: |
56 | newOp = rewriter.create<ZOp>(loc, IntegerType::get(context, 32)); |
57 | break; |
58 | } |
59 | |
60 | Operation *function; |
61 | if (auto gpuFunc = op->template getParentOfType<gpu::GPUFuncOp>()) |
62 | function = gpuFunc; |
63 | if (auto llvmFunc = op->template getParentOfType<LLVM::LLVMFuncOp>()) |
64 | function = llvmFunc; |
65 | if (!boundsAttrName.empty() && function) { |
66 | if (auto attr = function->template getAttrOfType<DenseI32ArrayAttr>( |
67 | boundsAttrName)) { |
68 | int32_t maximum = attr[static_cast<uint32_t>(op.getDimension())]; |
69 | newOp->setAttr("range" , rewriter.getDenseI32ArrayAttr({0, maximum})); |
70 | } |
71 | } |
72 | |
73 | if (indexBitwidth > 32) { |
74 | newOp = rewriter.create<LLVM::SExtOp>( |
75 | loc, IntegerType::get(context, indexBitwidth), newOp->getResult(0)); |
76 | } else if (indexBitwidth < 32) { |
77 | newOp = rewriter.create<LLVM::TruncOp>( |
78 | loc, IntegerType::get(context, indexBitwidth), newOp->getResult(0)); |
79 | } |
80 | |
81 | rewriter.replaceOp(op, newOp->getResults()); |
82 | return success(); |
83 | } |
84 | }; |
85 | |
86 | } // namespace mlir |
87 | |
88 | #endif // MLIR_CONVERSION_GPUCOMMON_INDEXINTRINSICSOPLOWERING_H_ |
89 | |