| //===- ACCExecutableDirectivePatterns.cpp - ACC exec patterns ---*- 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 |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Lowers OpenACC executable directives (init, shutdown, wait, set) to calls to |
| // an OpenACC offloading runtime compiler interface. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Conversion/OpenACCToLLVM/ACCToLLVM.h" |
| #include "mlir/Conversion/OpenACCToLLVM/ACCToLLVMUtils.h" |
| |
| #include "mlir/Conversion/LLVMCommon/Pattern.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
| #include "mlir/Dialect/OpenACC/OpenACC.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include "llvm/ADT/STLFunctionalExtras.h" |
| |
| #include <cstdint> |
| #include <iterator> |
| |
| using namespace mlir; |
| using namespace mlir::acc; |
| |
| namespace { |
| static Value castToI64(Location loc, Value value, |
| ConversionPatternRewriter &rewriter) { |
| Type i64Ty = IntegerType::get(rewriter.getContext(), 64); |
| unsigned bitwidth = value.getType().getIntOrFloatBitWidth(); |
| if (bitwidth > 64) |
| return arith::TruncIOp::create(rewriter, loc, i64Ty, value); |
| if (bitwidth < 64) |
| return arith::ExtSIOp::create(rewriter, loc, i64Ty, value); |
| return value; |
| } |
| |
| static Value getAsyncQueue(WaitOp op, ConversionPatternRewriter &rewriter, |
| const ACCRuntimeCallConfig &config) { |
| Location loc = op->getLoc(); |
| Type i64Ty = IntegerType::get(rewriter.getContext(), 64); |
| if (op.getAsync()) |
| return LLVM::ConstantOp::create(rewriter, loc, i64Ty, |
| config.getAsyncNoValueRuntimeValue()); |
| if (Value asyncValue = op.getAsyncOperand()) { |
| asyncValue = rewriter.getRemappedValue(asyncValue); |
| return castToI64(loc, asyncValue, rewriter); |
| } |
| return LLVM::ConstantOp::create(rewriter, loc, i64Ty, |
| config.getAsyncSyncRuntimeValue()); |
| } |
| |
| static LogicalResult createIfThen(Location loc, Value ifCond, |
| ConversionPatternRewriter &rewriter, |
| function_ref<LogicalResult()> thenFn) { |
| Block *parentBlock = rewriter.getInsertionBlock(); |
| Block *continueBlock = |
| rewriter.splitBlock(parentBlock, rewriter.getInsertionPoint()); |
| Block *thenBlock = rewriter.createBlock( |
| parentBlock->getParent(), std::next(Region::iterator(parentBlock))); |
| |
| rewriter.setInsertionPointToEnd(parentBlock); |
| LLVM::CondBrOp::create(rewriter, loc, ifCond, thenBlock, ValueRange{}, |
| continueBlock, ValueRange{}); |
| |
| rewriter.setInsertionPointToStart(thenBlock); |
| LogicalResult result = thenFn(); |
| rewriter.setInsertionPointToEnd(thenBlock); |
| LLVM::BrOp::create(rewriter, loc, ValueRange{}, continueBlock); |
| rewriter.setInsertionPointToStart(continueBlock); |
| return result; |
| } |
| |
| /// Run \p emitFn, guarded by a branch on \p ifCond when it is present. |
| static LogicalResult emitGuardedByIfCond(Location loc, Value ifCond, |
| ConversionPatternRewriter &rewriter, |
| function_ref<LogicalResult()> emitFn) { |
| if (ifCond) |
| return createIfThen(loc, ifCond, rewriter, emitFn); |
| return emitFn(); |
| } |
| |
| template <typename OpTy> |
| struct ACCExecutableDirectivePattern : public ConvertOpToLLVMPattern<OpTy> { |
| ACCExecutableDirectivePattern(const LLVMTypeConverter &converter, |
| const ACCRuntimeCallConfig &config, |
| PatternBenefit benefit = 1) |
| : ConvertOpToLLVMPattern<OpTy>(converter, benefit), config(config) {} |
| |
| ACCRuntimeCallConfig config; |
| }; |
| |
| struct WaitOpLowering : public ACCExecutableDirectivePattern<WaitOp> { |
| using ACCExecutableDirectivePattern<WaitOp>::ACCExecutableDirectivePattern; |
| |
| LogicalResult |
| matchAndRewrite(WaitOp op, WaitOp::Adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| Location loc = op->getLoc(); |
| ModuleOp module = op->getParentOfType<ModuleOp>(); |
| Type i32Ty = rewriter.getI32Type(); |
| Type i64Ty = rewriter.getI64Type(); |
| Type ptrTy = LLVM::LLVMPointerType::get(rewriter.getContext()); |
| |
| auto emitWait = [&]() -> LogicalResult { |
| Value asyncQueue = getAsyncQueue(op, rewriter, config); |
| SmallVector<Value> waitValues; |
| for (Value operand : op.getWaitOperands()) |
| waitValues.push_back( |
| castToI64(loc, rewriter.getRemappedValue(operand), rewriter)); |
| |
| unsigned size = waitValues.size(); |
| Value waitNum = LLVM::ConstantOp::create(rewriter, loc, i32Ty, size); |
| Value waitList; |
| if (size == 0) { |
| waitList = LLVM::ZeroOp::create(rewriter, loc, ptrTy); |
| } else { |
| waitList = LLVM::AllocaOp::create(rewriter, loc, ptrTy, i64Ty, waitNum); |
| for (auto [index, waitValue] : llvm::enumerate(waitValues)) { |
| Value idx = LLVM::ConstantOp::create(rewriter, loc, i32Ty, |
| static_cast<int64_t>(index)); |
| Value elementPtr = LLVM::GEPOp::create( |
| rewriter, loc, ptrTy, i64Ty, waitList, ArrayRef<Value>{idx}); |
| LLVM::StoreOp::create(rewriter, loc, waitValue, elementPtr); |
| } |
| } |
| |
| StringRef functionName = getParentFunctionName(waitValues); |
| if (functionName.empty()) |
| functionName = getParentFunctionName(op); |
| Value ident = createIdent(loc, functionName, rewriter, module, config); |
| Value flags = LLVM::ConstantOp::create(rewriter, loc, i64Ty, 0); |
| Value deviceType = LLVM::ConstantOp::create( |
| rewriter, loc, i64Ty, |
| config.getDeviceTypeRuntimeValue(DeviceType::None)); |
| Value deviceNum = LLVM::ConstantOp::create(rewriter, loc, i32Ty, 0); |
| |
| return createRuntimeCall( |
| loc, rewriter, module, RuntimeFunction::ACCRTL_tgt_acc_wait, config, |
| {ident, flags, deviceType, deviceNum, waitNum, waitList, asyncQueue}); |
| }; |
| |
| if (failed(emitGuardedByIfCond(loc, op.getIfCond(), rewriter, emitWait))) |
| return failure(); |
| |
| rewriter.eraseOp(op); |
| return success(); |
| } |
| }; |
| |
| /// Emit a call to a runtime entry point taking |
| /// `(ident, flags, deviceType, deviceNum)`. A null `deviceNum` selects the |
| /// current device. |
| static LogicalResult |
| emitDeviceOperationCall(Location loc, RuntimeFunction fn, DeviceType deviceType, |
| Value deviceNum, StringRef functionName, |
| ModuleOp module, ConversionPatternRewriter &rewriter, |
| const ACCRuntimeCallConfig &config) { |
| Type i64Ty = rewriter.getI64Type(); |
| Value deviceTypeValue = LLVM::ConstantOp::create( |
| rewriter, loc, i64Ty, config.getDeviceTypeRuntimeValue(deviceType)); |
| Value ident = createIdent(loc, functionName, rewriter, module, config); |
| Value flags = LLVM::ConstantOp::create(rewriter, loc, i64Ty, 0); |
| Value deviceNumValue = |
| deviceNum ? castToI64(loc, deviceNum, rewriter) |
| : LLVM::ConstantOp::create(rewriter, loc, i64Ty, -1); |
| return createRuntimeCall(loc, rewriter, module, fn, config, |
| {ident, flags, deviceTypeValue, deviceNumValue}); |
| } |
| |
| static LogicalResult rewriteInitOrShutdown(Operation *op, Value deviceNum, |
| ArrayAttr deviceTypesAttr, |
| Value ifCond, bool isInit, |
| ConversionPatternRewriter &rewriter, |
| const ACCRuntimeCallConfig &config) { |
| ModuleOp module = op->getParentOfType<ModuleOp>(); |
| Location loc = op->getLoc(); |
| |
| auto emitCalls = [&]() -> LogicalResult { |
| StringRef functionName = deviceNum ? getParentFunctionName(deviceNum) |
| : getParentFunctionName(op); |
| RuntimeFunction fn = isInit ? RuntimeFunction::ACCRTL_tgt_acc_init |
| : RuntimeFunction::ACCRTL_tgt_acc_shutdown; |
| |
| auto emitOne = [&](DeviceType deviceType) { |
| return emitDeviceOperationCall(loc, fn, deviceType, deviceNum, |
| functionName, module, rewriter, config); |
| }; |
| |
| if (!deviceTypesAttr) |
| return emitOne(DeviceType::None); |
| |
| for (Attribute attr : deviceTypesAttr) { |
| if (auto typeAttr = dyn_cast<DeviceTypeAttr>(attr)) |
| if (failed(emitOne(typeAttr.getValue()))) |
| return failure(); |
| } |
| return success(); |
| }; |
| |
| if (failed(emitGuardedByIfCond(loc, ifCond, rewriter, emitCalls))) |
| return failure(); |
| |
| rewriter.eraseOp(op); |
| return success(); |
| } |
| |
| struct InitOpLowering : public ACCExecutableDirectivePattern<InitOp> { |
| using ACCExecutableDirectivePattern<InitOp>::ACCExecutableDirectivePattern; |
| |
| LogicalResult |
| matchAndRewrite(InitOp op, InitOp::Adaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| return rewriteInitOrShutdown(op, adaptor.getDeviceNum(), |
| op.getDeviceTypesAttr(), op.getIfCond(), |
| /*isInit=*/true, rewriter, config); |
| } |
| }; |
| |
| struct ShutdownOpLowering : public ACCExecutableDirectivePattern<ShutdownOp> { |
| using ACCExecutableDirectivePattern< |
| ShutdownOp>::ACCExecutableDirectivePattern; |
| |
| LogicalResult |
| matchAndRewrite(ShutdownOp op, ShutdownOp::Adaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| return rewriteInitOrShutdown(op, adaptor.getDeviceNum(), |
| op.getDeviceTypesAttr(), op.getIfCond(), |
| /*isInit=*/false, rewriter, config); |
| } |
| }; |
| |
| struct SetOpLowering : public ACCExecutableDirectivePattern<SetOp> { |
| using ACCExecutableDirectivePattern<SetOp>::ACCExecutableDirectivePattern; |
| |
| LogicalResult |
| matchAndRewrite(SetOp op, SetOp::Adaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| ModuleOp module = op->getParentOfType<ModuleOp>(); |
| Location loc = op.getLoc(); |
| Type i64Ty = rewriter.getI64Type(); |
| |
| auto emitSet = [&]() -> LogicalResult { |
| if (Value asyncValue = adaptor.getDefaultAsync()) { |
| asyncValue = castToI64(loc, asyncValue, rewriter); |
| Value ident = createIdent(loc, getParentFunctionName(asyncValue), |
| rewriter, module, config); |
| if (failed(createRuntimeCall( |
| loc, rewriter, module, |
| RuntimeFunction::ACCRTL_tgt_acc_set_default_async, config, |
| {ident, asyncValue}))) |
| return failure(); |
| } |
| |
| if (op.getDeviceNum()) { |
| Value deviceNum = adaptor.getDeviceNum(); |
| DeviceType deviceType = DeviceType::None; |
| if (auto deviceTypeAttr = op.getDeviceTypeAttr()) |
| deviceType = deviceTypeAttr.getValue(); |
| return emitDeviceOperationCall( |
| loc, RuntimeFunction::ACCRTL_tgt_acc_set_device_num, deviceType, |
| deviceNum, getParentFunctionName(deviceNum), module, rewriter, |
| config); |
| } |
| if (auto deviceTypeAttr = op.getDeviceTypeAttr()) { |
| Value deviceTypeValue = LLVM::ConstantOp::create( |
| rewriter, loc, i64Ty, |
| config.getDeviceTypeRuntimeValue(deviceTypeAttr.getValue())); |
| Value ident = createIdent(loc, StringRef(), rewriter, module, config); |
| Value flags = LLVM::ConstantOp::create(rewriter, loc, i64Ty, 0); |
| return createRuntimeCall( |
| loc, rewriter, module, |
| RuntimeFunction::ACCRTL_tgt_acc_set_device_type, config, |
| {ident, flags, deviceTypeValue}); |
| } |
| return success(); |
| }; |
| |
| if (failed(emitGuardedByIfCond(loc, op.getIfCond(), rewriter, emitSet))) |
| return failure(); |
| |
| rewriter.eraseOp(op); |
| return success(); |
| } |
| }; |
| |
| } // namespace |
| |
| void mlir::configureACCExecutableDirectiveConversionLegality( |
| ConversionTarget &target) { |
| target.addIllegalOp<acc::InitOp, acc::ShutdownOp, acc::WaitOp, acc::SetOp>(); |
| } |
| |
| void mlir::populateACCExecutableDirectivePatterns( |
| LLVMTypeConverter &converter, RewritePatternSet &patterns, |
| const acc::ACCRuntimeCallConfig &config) { |
| patterns |
| .add<WaitOpLowering, InitOpLowering, ShutdownOpLowering, SetOpLowering>( |
| converter, config); |
| } |
