| // Copyright 2019 The IREE Authors |
| // |
| // Licensed 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 |
| |
| #include "iree/compiler/Dialect/HAL/IR/HALOps.h" |
| #include "iree/compiler/Dialect/VM/Conversion/ImportUtils.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| |
| namespace mlir::iree_compiler { |
| |
| // Rewrites a hal.device.query of an i1/i16/i32 value to a hal.device.query of |
| // i64 with a truncation on the result. |
| class DeviceQueryCastOpConversion |
| : public OpConversionPattern<IREE::HAL::DeviceQueryOp> { |
| public: |
| DeviceQueryCastOpConversion(MLIRContext *context, |
| TypeConverter &typeConverter) |
| : OpConversionPattern(typeConverter, context) {} |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::DeviceQueryOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto targetType = op.getValue().getType(); |
| if (targetType.isInteger(64)) |
| return failure(); // handled natively |
| if (!targetType.isIntOrIndex()) |
| return rewriter.notifyMatchFailure(op, "unsupported result type"); |
| |
| // Query as i64. |
| // Note that due to type conversion we need to handle the default logic |
| // ourselves instead of allowing the i64 do the same. We could let it handle |
| // things but then we are generating more IR that may prevent other |
| // canonicalizations (a select of i1 to i1 is easier to handle). |
| auto queryOp = rewriter.create<IREE::HAL::DeviceQueryOp>( |
| op.getLoc(), rewriter.getI1Type(), rewriter.getI64Type(), |
| adaptor.getDevice(), op.getCategoryAttr(), op.getKeyAttr(), |
| TypedAttr{}); |
| auto ok = llvm::cast<Value>(queryOp.getOk()); |
| auto value = queryOp.getValue(); |
| |
| // Truncate or extend based on the target type. |
| if (targetType.isIndex()) { |
| // i64 -> index cast. |
| value = |
| rewriter.create<arith::IndexCastOp>(op.getLoc(), targetType, value); |
| } else if (targetType.isInteger(1)) { |
| // i64 -> i1 cast. |
| value = rewriter.create<IREE::VM::CmpNZI64Op>( |
| op.getLoc(), rewriter.getI32Type(), value); |
| } else { |
| // i64 -> {integer} cast. |
| if (targetType.getIntOrFloatBitWidth() < |
| value.getType().getIntOrFloatBitWidth()) { |
| // i64 -> narrowing cast. |
| value = |
| rewriter.create<arith::TruncIOp>(op.getLoc(), targetType, value); |
| } else { |
| // i64 -> widening cast. |
| value = rewriter.create<arith::ExtUIOp>(op.getLoc(), targetType, value); |
| } |
| } |
| |
| if (op.getDefaultValue().has_value()) { |
| // Select the default value based on the converted type as that's the type |
| // of the attribute we have is in. 'ok' result is set to true as we've |
| // already handled the error case. |
| value = rewriter.create<arith::SelectOp>( |
| op.getLoc(), ok, value, |
| rewriter.create<arith::ConstantOp>(op.getLoc(), |
| op.getDefaultValueAttr())); |
| ok = rewriter.create<IREE::VM::ConstI32Op>(op.getLoc(), 1); |
| } |
| |
| rewriter.replaceOp(op, {ok, value}); |
| return success(); |
| } |
| }; |
| |
| class DeviceQueryI64OpConversion |
| : public OpConversionPattern<IREE::HAL::DeviceQueryOp> { |
| public: |
| DeviceQueryI64OpConversion(MLIRContext *context, SymbolTable &importSymbols, |
| TypeConverter &typeConverter, StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::DeviceQueryOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| if (!op.getValue().getType().isInteger(64)) |
| return failure(); |
| auto results = |
| rewriteToCall(op, adaptor, importOp, *getTypeConverter(), rewriter); |
| if (!results.has_value()) |
| return failure(); |
| auto ok = results->front(); |
| auto value = results->back(); |
| if (op.getDefaultValue().has_value()) { |
| value = rewriter.createOrFold<arith::SelectOp>( |
| op.getLoc(), ok, value, |
| rewriter.createOrFold<IREE::VM::ConstI64Op>( |
| op.getLoc(), op.getDefaultValueAttr())); |
| ok = rewriter.createOrFold<IREE::VM::ConstI64Op>(op.getLoc(), 1); |
| } |
| rewriter.replaceOp(op, {ok, value}); |
| return success(); |
| } |
| |
| private: |
| mutable IREE::VM::ImportOp importOp; |
| }; |
| |
| class DeviceQueueFillOpConversion |
| : public OpConversionPattern<IREE::HAL::DeviceQueueFillOp> { |
| public: |
| DeviceQueueFillOpConversion(MLIRContext *context, SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::DeviceQueueFillOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| auto i64Type = rewriter.getI64Type(); |
| auto patternLength = rewriter.create<IREE::VM::ConstI32Op>( |
| op.getLoc(), |
| llvm::divideCeil(op.getPattern().getType().getIntOrFloatBitWidth(), 8)); |
| auto flags = |
| rewriter.create<IREE::VM::ConstI64Op>(op.getLoc(), op.getFlags()); |
| std::array<Value, 10> callOperands = { |
| adaptor.getDevice(), |
| castToImportType(adaptor.getQueueAffinity(), i64Type, rewriter), |
| adaptor.getWaitFence(), |
| adaptor.getSignalFence(), |
| adaptor.getTargetBuffer(), |
| castToImportType(adaptor.getTargetOffset(), i64Type, rewriter), |
| castToImportType(adaptor.getLength(), i64Type, rewriter), |
| castToImportType(adaptor.getPattern(), i64Type, rewriter), |
| patternLength, |
| flags, |
| }; |
| auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>( |
| op, SymbolRefAttr::get(importOp), importType.getResults(), |
| callOperands); |
| copyImportAttrs(importOp, callOp); |
| return success(); |
| } |
| |
| private: |
| mutable IREE::VM::ImportOp importOp; |
| }; |
| |
| class DeviceQueueExecuteIndirectOpConversion |
| : public OpConversionPattern<IREE::HAL::DeviceQueueExecuteIndirectOp> { |
| public: |
| DeviceQueueExecuteIndirectOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::DeviceQueueExecuteIndirectOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| auto i64Type = rewriter.getI64Type(); |
| |
| SmallVector<Value, 8> callOperands = { |
| adaptor.getDevice(), |
| castToImportType(adaptor.getQueueAffinity(), i64Type, rewriter), |
| adaptor.getWaitFence(), |
| adaptor.getSignalFence(), |
| adaptor.getCommandBuffer(), |
| }; |
| SmallVector<int16_t, 5> segmentSizes = { |
| /*device=*/-1, |
| /*queue_affinity=*/-1, |
| /*wait_fence=*/-1, |
| /*signal_fence=*/-1, |
| /*command_buffer=*/-1, |
| /*bindings=*/ |
| static_cast<int16_t>(adaptor.getBindingBuffers().size()), |
| }; |
| for (auto [bindingBuffer, bindingOffset, bindingLength] : llvm::zip_equal( |
| adaptor.getBindingBuffers(), adaptor.getBindingOffsets(), |
| adaptor.getBindingLengths())) { |
| callOperands.push_back(bindingBuffer); |
| callOperands.push_back( |
| castToImportType(bindingOffset, i64Type, rewriter)); |
| callOperands.push_back( |
| castToImportType(bindingLength, i64Type, rewriter)); |
| } |
| |
| auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallVariadicOp>( |
| op, SymbolRefAttr::get(importOp), importType.getResults(), segmentSizes, |
| importType.getInputs(), callOperands); |
| copyImportAttrs(importOp, callOp); |
| return success(); |
| } |
| |
| private: |
| mutable IREE::VM::ImportOp importOp; |
| }; |
| |
| void populateHALDeviceToVMPatterns(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| RewritePatternSet &patterns) { |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceAllocatorOp>>( |
| context, importSymbols, typeConverter, "hal.device.allocator"); |
| |
| patterns.insert<DeviceQueryCastOpConversion>(context, typeConverter); |
| patterns.insert<DeviceQueryI64OpConversion>( |
| context, importSymbols, typeConverter, "hal.device.query.i64"); |
| |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueAllocaOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.alloca"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueDeallocaOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.dealloca"); |
| patterns.insert<DeviceQueueFillOpConversion>( |
| context, importSymbols, typeConverter, "hal.device.queue.fill"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueUpdateOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.update"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueCopyOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.copy"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueReadOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.read"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueWriteOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.write"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueExecuteOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.execute"); |
| patterns.insert<DeviceQueueExecuteIndirectOpConversion>( |
| context, importSymbols, typeConverter, |
| "hal.device.queue.execute.indirect"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueFlushOp>>( |
| context, importSymbols, typeConverter, "hal.device.queue.flush"); |
| } |
| |
| } // namespace mlir::iree_compiler |
