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 "llvm/ADT/StringSwitch.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 
28 public:
GPUIndexIntrinsicOpLoweringGPUIndexIntrinsicOpLowering29   explicit GPUIndexIntrinsicOpLowering(LLVMTypeConverter &typeConverter)
30       : ConvertOpToLLVMPattern<Op>(typeConverter),
31         indexBitwidth(typeConverter.getIndexTypeBitwidth()) {}
32 
33   // Convert the kernel arguments to an LLVM type, preserve the rest.
34   LogicalResult
matchAndRewriteGPUIndexIntrinsicOpLowering35   matchAndRewrite(Op op, typename Op::Adaptor adaptor,
36                   ConversionPatternRewriter &rewriter) const override {
37     auto loc = op->getLoc();
38     MLIRContext *context = rewriter.getContext();
39     Value newOp;
40     switch (op.dimension()) {
41     case gpu::Dimension::x:
42       newOp = rewriter.create<XOp>(loc, IntegerType::get(context, 32));
43       break;
44     case gpu::Dimension::y:
45       newOp = rewriter.create<YOp>(loc, IntegerType::get(context, 32));
46       break;
47     case gpu::Dimension::z:
48       newOp = rewriter.create<ZOp>(loc, IntegerType::get(context, 32));
49       break;
50     }
51 
52     if (indexBitwidth > 32) {
53       newOp = rewriter.create<LLVM::SExtOp>(
54           loc, IntegerType::get(context, indexBitwidth), newOp);
55     } else if (indexBitwidth < 32) {
56       newOp = rewriter.create<LLVM::TruncOp>(
57           loc, IntegerType::get(context, indexBitwidth), newOp);
58     }
59 
60     rewriter.replaceOp(op, {newOp});
61     return success();
62   }
63 };
64 
65 } // namespace mlir
66 
67 #endif // MLIR_CONVERSION_GPUCOMMON_INDEXINTRINSICSOPLOWERING_H_
68