| // 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/Util/IR/UtilTypes.h" |
| #include "iree/compiler/Dialect/VM/Conversion/ImportUtils.h" |
| #include "iree/compiler/Dialect/VM/IR/VMOps.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| |
| namespace mlir::iree_compiler { |
| |
| namespace { |
| |
| // Returns a slot value and a buffer ref value. |
| // |bufferOrSlot| is intended to be a `AnyTypeOf<[Index, HAL_BufferType]>` in |
| // the op definition. |
| static std::tuple<Value, Value> |
| splitBufferSlot(Location loc, Value bufferOrSlot, OpBuilder &builder) { |
| if (!bufferOrSlot) { |
| return std::make_tuple( |
| builder.create<IREE::VM::ConstI32ZeroOp>(loc), |
| builder.create<IREE::VM::ConstRefZeroOp>( |
| loc, |
| IREE::VM::RefType::get(builder.getType<IREE::HAL::BufferType>()))); |
| } else if (isa<IREE::VM::RefType>(bufferOrSlot.getType())) { |
| // Direct buffer binding; pass 0 for table slot. |
| return std::make_tuple(builder.create<IREE::VM::ConstI32ZeroOp>(loc), |
| bufferOrSlot); |
| } else { |
| // Indirect binding table reference; pass null for the buffer. |
| return std::make_tuple( |
| castToImportType(bufferOrSlot, builder.getI32Type(), builder), |
| builder.create<IREE::VM::ConstRefZeroOp>( |
| loc, |
| IREE::VM::RefType::get(builder.getType<IREE::HAL::BufferType>()))); |
| } |
| } |
| |
| // TODO(benvanik): import op handling of optional values. |
| // It'd be nice if the std::optional<Index>:$binding_capacity could be emitted |
| // as 0 when not present; today it'll be omitted entirely (as it's not in the |
| // operand set) but we need it for the fixed call signature. |
| class CommandBufferCreateOpConversion |
| : public OpConversionPattern<IREE::HAL::CommandBufferCreateOp> { |
| public: |
| CommandBufferCreateOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::CommandBufferCreateOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| |
| SmallVector<Value, 8> callOperands = { |
| adaptor.getDevice(), |
| }; |
| auto modesValue = detail::rewriteAttrToOperands( |
| op.getLoc(), adaptor.getModesAttr(), rewriter.getI32Type(), rewriter); |
| if (!modesValue.has_value()) |
| return failure(); |
| callOperands.append(modesValue.value()); |
| auto categoriesValue = detail::rewriteAttrToOperands( |
| op.getLoc(), adaptor.getCommandCategoriesAttr(), rewriter.getI32Type(), |
| rewriter); |
| if (!categoriesValue.has_value()) |
| return failure(); |
| callOperands.append(categoriesValue.value()); |
| callOperands.push_back(adaptor.getQueueAffinity()); |
| if (adaptor.getBindingCapacity()) { |
| callOperands.push_back(castToImportType(adaptor.getBindingCapacity(), |
| rewriter.getI32Type(), rewriter)); |
| } else { |
| callOperands.push_back( |
| rewriter.create<IREE::VM::ConstI32ZeroOp>(op.getLoc())); |
| } |
| |
| 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 CommandBufferFillBufferOpConversion |
| : public OpConversionPattern<IREE::HAL::CommandBufferFillBufferOp> { |
| public: |
| CommandBufferFillBufferOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::CommandBufferFillBufferOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| |
| auto [targetBufferSlot, targetBuffer] = |
| splitBufferSlot(op.getLoc(), adaptor.getTargetBuffer(), rewriter); |
| SmallVector<Value, 8> callOperands = { |
| adaptor.getCommandBuffer(), |
| targetBuffer, |
| castToImportType(adaptor.getTargetOffset(), rewriter.getI64Type(), |
| rewriter), |
| castToImportType(adaptor.getLength(), rewriter.getI64Type(), rewriter), |
| targetBufferSlot, |
| }; |
| |
| // Record the original pattern length then extend it to a 32 bit integer. |
| auto originalPatternType = op.getPattern().getType(); |
| unsigned patternBitWidth = IREE::Util::getTypeBitWidth(originalPatternType); |
| // The pattern length (in bytes) will be used at runtime to issue the fill |
| // command. While the pattern itself will be stored in a 32 bit integer, |
| // the fill operation will use this length to slice a potentially smaller |
| // range of bits from the full pattern. |
| auto patternLengthBytes = |
| IREE::Util::getRoundedElementByteWidth(originalPatternType); |
| auto patternLengthConst = rewriter.createOrFold<mlir::arith::ConstantIntOp>( |
| op.getLoc(), patternLengthBytes, 32); |
| Value pattern = op.getPattern(); |
| if (patternBitWidth < 32) { |
| pattern = rewriter.createOrFold<arith::ExtUIOp>( |
| op.getLoc(), rewriter.getIntegerType(32), pattern); |
| } |
| callOperands.push_back(pattern); |
| callOperands.push_back(patternLengthConst); |
| |
| 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 CommandBufferUpdateBufferOpConversion |
| : public OpConversionPattern<IREE::HAL::CommandBufferUpdateBufferOp> { |
| public: |
| CommandBufferUpdateBufferOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::CommandBufferUpdateBufferOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| auto [targetBufferSlot, targetBuffer] = |
| splitBufferSlot(op.getLoc(), adaptor.getTargetBuffer(), rewriter); |
| SmallVector<Value, 8> callOperands = { |
| adaptor.getCommandBuffer(), |
| adaptor.getSourceBuffer(), |
| castToImportType(adaptor.getSourceOffset(), rewriter.getI64Type(), |
| rewriter), |
| targetBuffer, |
| castToImportType(adaptor.getTargetOffset(), rewriter.getI64Type(), |
| rewriter), |
| castToImportType(adaptor.getLength(), rewriter.getI64Type(), rewriter), |
| targetBufferSlot}; |
| 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 CommandBufferCopyBufferOpConversion |
| : public OpConversionPattern<IREE::HAL::CommandBufferCopyBufferOp> { |
| public: |
| CommandBufferCopyBufferOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::CommandBufferCopyBufferOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| auto [sourceBufferSlot, sourceBuffer] = |
| splitBufferSlot(op.getLoc(), adaptor.getSourceBuffer(), rewriter); |
| auto [targetBufferSlot, targetBuffer] = |
| splitBufferSlot(op.getLoc(), adaptor.getTargetBuffer(), rewriter); |
| SmallVector<Value, 8> callOperands = { |
| adaptor.getCommandBuffer(), |
| sourceBufferSlot, |
| targetBufferSlot, |
| sourceBuffer, |
| castToImportType(adaptor.getSourceOffset(), rewriter.getI64Type(), |
| rewriter), |
| targetBuffer, |
| castToImportType(adaptor.getTargetOffset(), rewriter.getI64Type(), |
| rewriter), |
| castToImportType(adaptor.getLength(), rewriter.getI64Type(), rewriter), |
| }; |
| 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 CommandBufferCollectiveOpConversion |
| : public OpConversionPattern<IREE::HAL::CommandBufferCollectiveOp> { |
| public: |
| CommandBufferCollectiveOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::CommandBufferCollectiveOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| |
| Value zeroI64; |
| auto getZeroI64 = [&]() { |
| if (!zeroI64) { |
| zeroI64 = rewriter.create<IREE::VM::ConstI64ZeroOp>(op.getLoc()); |
| } |
| return zeroI64; |
| }; |
| |
| // %command_buffer : !vm.ref<!hal.command_buffer>, |
| // %channel : !vm.ref<!hal.channel>, |
| // %op : i32, |
| // %param : i32, |
| // %send_buffer_slot : i32, |
| // %recv_buffer_slot : i32, |
| // %send_buffer : !vm.ref<!hal.buffer>, |
| // %recv_buffer : !vm.ref<!hal.buffer>, |
| // %send_offset : i64, |
| // %send_length : i64, |
| // %recv_offset : i64, |
| // %recv_length : i64, |
| // %element_count : i64 |
| SmallVector<Value> callOperands; |
| callOperands.push_back(adaptor.getCommandBuffer()); |
| callOperands.push_back(adaptor.getChannel()); |
| callOperands.push_back(rewriter.create<IREE::VM::ConstI32Op>( |
| op.getLoc(), adaptor.getOp().getEncodedValue())); |
| if (auto paramValue = adaptor.getParam()) { |
| callOperands.push_back(paramValue); |
| } else { |
| callOperands.push_back( |
| rewriter.create<IREE::VM::ConstI32ZeroOp>(op.getLoc())); |
| } |
| |
| auto [sendBufferSlot, sendBuffer] = |
| splitBufferSlot(op.getLoc(), adaptor.getSendBuffer(), rewriter); |
| auto [recvBufferSlot, recvBuffer] = |
| splitBufferSlot(op.getLoc(), adaptor.getRecvBuffer(), rewriter); |
| callOperands.push_back(sendBufferSlot); |
| callOperands.push_back(recvBufferSlot); |
| callOperands.push_back(sendBuffer); |
| callOperands.push_back(recvBuffer); |
| callOperands.push_back(adaptor.getSendOffset() ? adaptor.getSendOffset() |
| : getZeroI64()); |
| callOperands.push_back(adaptor.getSendLength() ? adaptor.getSendLength() |
| : getZeroI64()); |
| callOperands.push_back(adaptor.getRecvOffset() ? adaptor.getRecvOffset() |
| : getZeroI64()); |
| callOperands.push_back(adaptor.getRecvLength() ? adaptor.getRecvLength() |
| : getZeroI64()); |
| |
| callOperands.push_back(castToImportType(adaptor.getElementCount(), |
| rewriter.getI64Type(), rewriter)); |
| |
| 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 CommandBufferDispatchOpConversion |
| : public OpConversionPattern<IREE::HAL::CommandBufferDispatchOp> { |
| public: |
| CommandBufferDispatchOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::CommandBufferDispatchOp op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto importType = importOp.getFunctionType(); |
| |
| auto i32Type = rewriter.getI32Type(); |
| auto i64Type = rewriter.getI64Type(); |
| Value zeroI32 = rewriter.create<IREE::VM::ConstI32ZeroOp>(op.getLoc()); |
| |
| auto flags = adaptor.getFlagsAttr() |
| ? rewriter |
| .create<IREE::VM::ConstI64Op>( |
| op.getLoc(), adaptor.getFlagsAttr().getInt()) |
| .getResult() |
| : rewriter.create<IREE::VM::ConstI64ZeroOp>(op.getLoc()) |
| .getResult(); |
| SmallVector<Value, 8> callOperands = { |
| adaptor.getCommandBuffer(), |
| adaptor.getExecutable(), |
| castToImportType(adaptor.getEntryPoint(), i32Type, rewriter), |
| castToImportType(adaptor.getWorkgroupX(), i32Type, rewriter), |
| castToImportType(adaptor.getWorkgroupY(), i32Type, rewriter), |
| castToImportType(adaptor.getWorkgroupZ(), i32Type, rewriter), |
| flags, |
| }; |
| SmallVector<int16_t, 5> segmentSizes = { |
| /*command_buffer=*/-1, |
| /*executable=*/-1, |
| /*entry_point=*/-1, |
| /*workgroup_x=*/-1, |
| /*workgroup_y=*/-1, |
| /*workgroup_z=*/-1, |
| /*flags=*/-1, |
| /*constants=*/static_cast<int16_t>(adaptor.getConstants().size()), |
| /*bindings=*/ |
| static_cast<int16_t>(adaptor.getBindingBuffers().size()), |
| }; |
| llvm::append_range(callOperands, adaptor.getConstants()); |
| for (auto [bindingBufferOrSlot, bindingOffset, bindingLength] : |
| llvm::zip_equal(adaptor.getBindingBuffers(), |
| adaptor.getBindingOffsets(), |
| adaptor.getBindingLengths())) { |
| callOperands.push_back(zeroI32); |
| auto [bindingBufferSlot, bindingBuffer] = |
| splitBufferSlot(op.getLoc(), bindingBufferOrSlot, rewriter); |
| callOperands.push_back(bindingBufferSlot); |
| 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; |
| }; |
| |
| class CommandBufferDispatchIndirectOpConversion |
| : public OpConversionPattern<IREE::HAL::CommandBufferDispatchIndirectOp> { |
| public: |
| CommandBufferDispatchIndirectOpConversion(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| StringRef importName) |
| : OpConversionPattern(typeConverter, context) { |
| importOp = importSymbols.lookup<IREE::VM::ImportOp>(importName); |
| assert(importOp); |
| } |
| |
| LogicalResult |
| matchAndRewrite(IREE::HAL::CommandBufferDispatchIndirectOp op, |
| OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| |
| auto importType = importOp.getFunctionType(); |
| |
| auto i32Type = rewriter.getI32Type(); |
| auto i64Type = rewriter.getI64Type(); |
| Value zeroI32 = rewriter.create<IREE::VM::ConstI32ZeroOp>(op.getLoc()); |
| |
| auto [workgroupsBufferSlot, workgroupsBuffer] = |
| splitBufferSlot(op.getLoc(), adaptor.getWorkgroupsBuffer(), rewriter); |
| auto flags = adaptor.getFlagsAttr() |
| ? rewriter |
| .create<IREE::VM::ConstI64Op>( |
| op.getLoc(), adaptor.getFlagsAttr().getInt()) |
| .getResult() |
| : rewriter.create<IREE::VM::ConstI64ZeroOp>(op.getLoc()) |
| .getResult(); |
| SmallVector<Value, 8> callOperands = { |
| adaptor.getCommandBuffer(), |
| adaptor.getExecutable(), |
| castToImportType(adaptor.getEntryPoint(), i32Type, rewriter), |
| workgroupsBufferSlot, |
| workgroupsBuffer, |
| castToImportType(adaptor.getWorkgroupsOffset(), i64Type, rewriter), |
| flags, |
| }; |
| SmallVector<int16_t, 5> segmentSizes = { |
| /*command_buffer=*/-1, |
| /*executable=*/-1, |
| /*entry_point=*/-1, |
| /*workgroups_buffer_slot=*/-1, |
| /*workgroups_buffer=*/-1, |
| /*workgroups_offset=*/-1, |
| /*flags=*/-1, |
| /*constants=*/static_cast<int16_t>(adaptor.getConstants().size()), |
| /*bindings=*/ |
| static_cast<int16_t>(adaptor.getBindingBuffers().size()), |
| }; |
| llvm::append_range(callOperands, adaptor.getConstants()); |
| for (auto [bindingBufferOrSlot, bindingOffset, bindingLength] : |
| llvm::zip_equal(adaptor.getBindingBuffers(), |
| adaptor.getBindingOffsets(), |
| adaptor.getBindingLengths())) { |
| callOperands.push_back(zeroI32); |
| auto [bindingBufferSlot, bindingBuffer] = |
| splitBufferSlot(op.getLoc(), bindingBufferOrSlot, rewriter); |
| callOperands.push_back(bindingBufferSlot); |
| 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; |
| }; |
| |
| } // namespace |
| |
| void populateHALCommandBufferToVMPatterns(MLIRContext *context, |
| SymbolTable &importSymbols, |
| TypeConverter &typeConverter, |
| RewritePatternSet &patterns) { |
| patterns.insert<CommandBufferCreateOpConversion>( |
| context, importSymbols, typeConverter, "hal.command_buffer.create"); |
| patterns.insert<VMImportOpConversion<IREE::HAL::CommandBufferFinalizeOp>>( |
| context, importSymbols, typeConverter, "hal.command_buffer.finalize"); |
| patterns |
| .insert<VMImportOpConversion<IREE::HAL::CommandBufferBeginDebugGroupOp>>( |
| context, importSymbols, typeConverter, |
| "hal.command_buffer.begin_debug_group"); |
| patterns |
| .insert<VMImportOpConversion<IREE::HAL::CommandBufferEndDebugGroupOp>>( |
| context, importSymbols, typeConverter, |
| "hal.command_buffer.end_debug_group"); |
| patterns |
| .insert<VMImportOpConversion<IREE::HAL::CommandBufferExecutionBarrierOp>>( |
| context, importSymbols, typeConverter, |
| "hal.command_buffer.execution_barrier"); |
| patterns.insert<CommandBufferFillBufferOpConversion>( |
| context, importSymbols, typeConverter, "hal.command_buffer.fill_buffer"); |
| patterns.insert<CommandBufferUpdateBufferOpConversion>( |
| context, importSymbols, typeConverter, |
| "hal.command_buffer.update_buffer"); |
| patterns.insert<CommandBufferCopyBufferOpConversion>( |
| context, importSymbols, typeConverter, "hal.command_buffer.copy_buffer"); |
| patterns.insert<CommandBufferCollectiveOpConversion>( |
| context, importSymbols, typeConverter, "hal.command_buffer.collective"); |
| patterns.insert<CommandBufferDispatchOpConversion>( |
| context, importSymbols, typeConverter, "hal.command_buffer.dispatch"); |
| patterns.insert<CommandBufferDispatchIndirectOpConversion>( |
| context, importSymbols, typeConverter, |
| "hal.command_buffer.dispatch.indirect"); |
| } |
| |
| } // namespace mlir::iree_compiler |
