5446ec8507
Historical modeling of the LLVM dialect types had been wrapping LLVM IR types and therefore needed access to the instance of LLVMContext stored in the LLVMDialect. The new modeling does not rely on that and only needs the MLIRContext that is used for uniquing, similarly to other MLIR types. Change LLVMType::get<Kind>Ty functions to take `MLIRContext *` instead of `LLVMDialect *` as first argument. This brings the code base closer to completely removing the dependence on LLVMContext from the LLVMDialect, together with additional support for thread-safety of its use. Depends On D85371 Reviewed By: rriddle Differential Revision: https://reviews.llvm.org/D85372
81 lines
2.9 KiB
C++
81 lines
2.9 KiB
C++
//===- IndexIntrinsicsOpLowering.h - GPU IndexOps Lowering class *- C++ -*-===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
#ifndef MLIR_CONVERSION_GPUCOMMON_INDEXINTRINSICSOPLOWERING_H_
|
|
#define MLIR_CONVERSION_GPUCOMMON_INDEXINTRINSICSOPLOWERING_H_
|
|
|
|
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
|
|
#include "mlir/Dialect/GPU/GPUDialect.h"
|
|
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
|
|
#include "llvm/ADT/StringSwitch.h"
|
|
|
|
namespace mlir {
|
|
|
|
// Rewriting that replaces Op with XOp, YOp, or ZOp depending on the dimension
|
|
// that Op operates on. Op is assumed to return an `std.index` value and
|
|
// XOp, YOp and ZOp are assumed to return an `llvm.i32` value. Depending on
|
|
// `indexBitwidth`, sign-extend or truncate the resulting value to match the
|
|
// bitwidth expected by the consumers of the value.
|
|
template <typename Op, typename XOp, typename YOp, typename ZOp>
|
|
struct GPUIndexIntrinsicOpLowering : public ConvertToLLVMPattern {
|
|
private:
|
|
enum dimension { X = 0, Y = 1, Z = 2, invalid };
|
|
unsigned indexBitwidth;
|
|
|
|
static dimension dimensionToIndex(Op op) {
|
|
return llvm::StringSwitch<dimension>(op.dimension())
|
|
.Case("x", X)
|
|
.Case("y", Y)
|
|
.Case("z", Z)
|
|
.Default(invalid);
|
|
}
|
|
|
|
public:
|
|
explicit GPUIndexIntrinsicOpLowering(LLVMTypeConverter &typeConverter)
|
|
: ConvertToLLVMPattern(Op::getOperationName(),
|
|
typeConverter.getDialect()->getContext(),
|
|
typeConverter),
|
|
indexBitwidth(typeConverter.getIndexTypeBitwidth()) {}
|
|
|
|
// Convert the kernel arguments to an LLVM type, preserve the rest.
|
|
LogicalResult
|
|
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
auto loc = op->getLoc();
|
|
MLIRContext *context = rewriter.getContext();
|
|
Value newOp;
|
|
switch (dimensionToIndex(cast<Op>(op))) {
|
|
case X:
|
|
newOp = rewriter.create<XOp>(loc, LLVM::LLVMType::getInt32Ty(context));
|
|
break;
|
|
case Y:
|
|
newOp = rewriter.create<YOp>(loc, LLVM::LLVMType::getInt32Ty(context));
|
|
break;
|
|
case Z:
|
|
newOp = rewriter.create<ZOp>(loc, LLVM::LLVMType::getInt32Ty(context));
|
|
break;
|
|
default:
|
|
return failure();
|
|
}
|
|
|
|
if (indexBitwidth > 32) {
|
|
newOp = rewriter.create<LLVM::SExtOp>(
|
|
loc, LLVM::LLVMType::getIntNTy(context, indexBitwidth), newOp);
|
|
} else if (indexBitwidth < 32) {
|
|
newOp = rewriter.create<LLVM::TruncOp>(
|
|
loc, LLVM::LLVMType::getIntNTy(context, indexBitwidth), newOp);
|
|
}
|
|
|
|
rewriter.replaceOp(op, {newOp});
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_CONVERSION_GPUCOMMON_INDEXINTRINSICSOPLOWERING_H_
|