| // Copyright 2024 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 "compiler/plugins/input/Torch/InputConversion/Passes.h" |
| #include "iree/compiler/Dialect/HAL/IR/HALDialect.h" |
| #include "iree/compiler/Dialect/HAL/IR/HALOps.h" |
| #include "iree/compiler/Dialect/HAL/IR/HALTypes.h" |
| #include "iree/compiler/Dialect/Util/IR/UtilDialect.h" |
| #include "iree/compiler/Dialect/Util/IR/UtilOps.h" |
| #include "mlir/Dialect/Func/IR/FuncOps.h" |
| #include "mlir/Dialect/Func/Transforms/FuncConversions.h" |
| #include "mlir/Dialect/Tensor/IR/Tensor.h" |
| #include "mlir/IR/BuiltinOps.h" |
| #include "torch-mlir/Dialect/Torch/IR/TorchOps.h" |
| #include "torch-mlir/Dialect/Torch/IR/TorchTypes.h" |
| #include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionDialect.h" |
| #include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h" |
| |
| namespace Torch = mlir::torch::Torch; |
| namespace TorchConversion = mlir::torch::TorchConversion; |
| |
| namespace mlir::iree_compiler::TorchInput { |
| |
| #define GEN_PASS_DEF_FUNCCONVERSIONPASS |
| #include "compiler/plugins/input/Torch/InputConversion/Passes.h.inc" |
| |
| namespace { |
| |
| //===----------------------------------------------------------------------===// |
| // Overall Approach |
| // ---------------- |
| // This pass converts from the "torch" programming model to the "iree" |
| // programming model by rewriting all functions and calls to operate on native |
| // IREE types. In the process, synchronization is added as appropriate for any |
| // mutable and immutable torch-level arguments. |
| // |
| // Currently, the result of this pass is that every torch-func is augmented |
| // to be implemented in terms of IREE's "coarse fences" ABI. In this ABI, |
| // there is a (wait, signal) fence pair added to the end of every function. |
| // Since torch functions are single-exit, practically, this involves: |
| // * Adding preambles to convert function arguments to `tensor` (and native |
| // torch types via `torch_c`), adding coarse synchronization on import |
| // (presently for all buffer arguments but in the future could only be |
| // those which are not tied to fine grained fences). |
| // * Adding a postamble with a synchronization barrier on any produced |
| // or mutated tensors and appropriate exports/in-place tieing to buffers. |
| // * Generation of a synchronous wrapper function with the original name |
| // (the async function is named with an `$async` suffix) which internally |
| // sets up/waits on fences while delegating to the async function. |
| // |
| // Immutable tensor types |
| // ---------------------- |
| // |
| // Immutable tensor types are mapped to a buffer_view and subject to |
| // `hal.tensor.import` on use. On return, they will be placed behind a |
| // synchronization barrier and exported. |
| // |
| // Mutable types |
| // ------------- |
| // Here we rely on the characteristic that at the torch level, conversion to |
| // and from the value domain is only legal at certain well defined points in |
| // the program (currently at graph edges but potentially in the future at |
| // various control flow points). These conversions are modeled by: |
| // * `torch.copy.to_vtensor`: Copy from a mutable tensor (torch.tensor) to |
| // an immutable value (torch.vtensor). |
| // * `torch.copy.to_tensor`: Allocates a new mutable tensor and initializes it |
| // with the value of the given immutable tensor. Presently un-used. |
| // * `torch.overwrite.tensor.contents`: Updates the contents of a mutable |
| // tensor from a given immutable tensor. |
| // |
| // Note that when importing from Torch, these ops cannot just be added at will, |
| // and they are only created as a result of structural conversions. Therefore, |
| // we can rely on these invariants and assume that usage outside of this is an |
| // invalid program. |
| //===----------------------------------------------------------------------===// |
| |
| std::optional<std::pair<Value, Value>> |
| getEnclosingWaitSignalFences(Operation *op) { |
| auto parentFuncOp = dyn_cast<IREE::Util::FuncOp>(op); |
| if (!parentFuncOp) { |
| parentFuncOp = parentFuncOp->getParentOfType<IREE::Util::FuncOp>(); |
| if (!parentFuncOp) |
| return {}; |
| } |
| Block *entryBlock = &parentFuncOp.front(); |
| auto numArguments = entryBlock->getNumArguments(); |
| Value coarseWaitFence = entryBlock->getArgument(numArguments - 2); |
| Value coarseSignalFence = entryBlock->getArgument(numArguments - 1); |
| return std::make_pair(coarseWaitFence, coarseSignalFence); |
| } |
| |
| std::optional<std::pair<Value, Value>> |
| getEnclosingWaitSignalFences(Value value) { |
| return getEnclosingWaitSignalFences(value.getParentRegion()->getParentOp()); |
| } |
| |
| Value convertToBuiltinTensor(OpBuilder &builder, Value possibleTorchTensor) { |
| Type ty = possibleTorchTensor.getType(); |
| if (isa<TensorType>(ty)) |
| return possibleTorchTensor; |
| |
| if (auto defining = dyn_cast_or_null<TorchConversion::FromBuiltinTensorOp>( |
| possibleTorchTensor.getDefiningOp())) { |
| return defining.getOperand(); |
| } |
| |
| Torch::ValueTensorType vtensorType = cast<Torch::ValueTensorType>(ty); |
| TensorType builtinTy = vtensorType.toBuiltinTensor(); |
| if (auto intTy = dyn_cast<IntegerType>(builtinTy.getElementType())) { |
| builtinTy = |
| builtinTy.clone(builder.getIntegerType(intTy.getIntOrFloatBitWidth())); |
| } |
| |
| return builder.create<TorchConversion::ToBuiltinTensorOp>( |
| possibleTorchTensor.getLoc(), builtinTy, possibleTorchTensor); |
| } |
| |
| enum class TypeDisposition { |
| IMMUTABLE_TENSOR, |
| MUTABLE_TENSOR, |
| TORCH_PRIMITIVE, |
| PASSTHROUGH, |
| FENCE, |
| }; |
| |
| struct BarrierResult { |
| BlockArgument storage; |
| Type torchType; |
| int returnIndex = -1; |
| }; |
| |
| struct ConvertedAsyncFunctionInfo { |
| IREE::Util::FuncOp funcOp; |
| SmallVector<IREE::Util::ReturnOp> returnOps; |
| SmallVector<DictionaryAttr> torchArgAttrs; |
| SmallVector<DictionaryAttr> torchResultAttrs; |
| SmallVector<Type> torchInputTypes; |
| SmallVector<Type> torchResultTypes; |
| SmallVector<TypeDisposition> inputDispositions; |
| SmallVector<TypeDisposition> resultDispositions; |
| |
| // Post processing state. |
| // Values that must be captured in the coarse barrier. |
| SmallVector<Value> barrierInputs; |
| // Meta data per barrier input: storage, torchType, returnIndex (or -1) |
| SmallVector<BarrierResult> barrierResultMeta; |
| |
| LogicalResult postProcess(); |
| LogicalResult convertImmutableTensorArg(BlockArgument argValue, |
| Type torchType, OpBuilder &builder); |
| LogicalResult convertMutableTensorArg(BlockArgument argValue, Type torchType, |
| OpBuilder &builder); |
| |
| void addBarrierInput(Value inputTensor, BlockArgument storage, Type torchType, |
| int returnIndex) { |
| barrierInputs.push_back(inputTensor); |
| barrierResultMeta.emplace_back(BarrierResult{ |
| storage, |
| torchType, |
| returnIndex, |
| }); |
| } |
| |
| Attribute getTorchArgAttr(BlockArgument argValue, StringRef attrName) { |
| return torchArgAttrs.empty() |
| ? Attribute{} |
| : torchArgAttrs[argValue.getArgNumber()].get(attrName); |
| } |
| Attribute getTorchResultAttr(int returnIndex, StringRef attrName) { |
| return torchResultAttrs.empty() |
| ? Attribute{} |
| : torchResultAttrs[returnIndex].get(attrName); |
| } |
| }; |
| |
| LogicalResult ConvertedAsyncFunctionInfo::postProcess() { |
| if (funcOp.isExternal()) |
| return success(); |
| |
| if (returnOps.size() != 1) { |
| // Multi-exit/CFG could be supported but requires more complicated dominance |
| // analysis with respect to where the exit happens relative to mutated |
| // buffers. |
| return emitError(funcOp.getLoc()) |
| << "currently only single exit torch funcs are supported"; |
| } |
| |
| Block *entryBlock = &funcOp.getBlocks().front(); |
| |
| // Materialize argument conversions. |
| OpBuilder preambleBuilder = OpBuilder::atBlockBegin(entryBlock); |
| auto entryArgs = entryBlock->getArguments(); |
| for (auto [disp, argValue, torchType] : |
| llvm::zip_equal(inputDispositions, entryArgs, torchInputTypes)) { |
| switch (disp) { |
| case TypeDisposition::IMMUTABLE_TENSOR: { |
| if (failed( |
| convertImmutableTensorArg(argValue, torchType, preambleBuilder))) |
| return failure(); |
| break; |
| } |
| case TypeDisposition::MUTABLE_TENSOR: { |
| if (failed(convertMutableTensorArg(argValue, torchType, preambleBuilder))) |
| return failure(); |
| break; |
| } |
| case TypeDisposition::TORCH_PRIMITIVE: { |
| Location loc = argValue.getLoc(); |
| Operation *convertUser = nullptr; |
| Value convertResult; |
| if (isa<Torch::BoolType>(torchType)) { |
| convertUser = |
| preambleBuilder.create<TorchConversion::FromI1Op>(loc, argValue); |
| convertResult = convertUser->getResult(0); |
| } else if (isa<Torch::FloatType>(torchType)) { |
| convertUser = |
| preambleBuilder.create<TorchConversion::FromF64Op>(loc, argValue); |
| convertResult = convertUser->getResult(0); |
| } else if (isa<Torch::IntType>(torchType)) { |
| convertUser = |
| preambleBuilder.create<TorchConversion::FromI64Op>(loc, argValue); |
| convertResult = convertUser->getResult(0); |
| } else { |
| emitError(loc) << "unhandled torch primitive materialization: " |
| << torchType; |
| return failure(); |
| } |
| argValue.replaceAllUsesExcept(convertResult, convertUser); |
| break; |
| } |
| case TypeDisposition::PASSTHROUGH: |
| // Do nothing. |
| break; |
| case TypeDisposition::FENCE: |
| // Do nothing. |
| break; |
| } |
| } |
| |
| // Materialize synchronization postamble and conversions. |
| IREE::Util::ReturnOp returnOp = returnOps.front(); |
| SmallVector<Value> newReturnOperands; |
| OpBuilder postambleBuilder(returnOp); |
| for (auto [disp, returnValue, torchType] : llvm::zip_equal( |
| resultDispositions, returnOp.getOperands(), torchResultTypes)) { |
| size_t returnIndex = newReturnOperands.size(); |
| newReturnOperands.emplace_back(returnValue); |
| switch (disp) { |
| case TypeDisposition::IMMUTABLE_TENSOR: { |
| bool needsBarrier = true; |
| if (auto blockArg = dyn_cast<BlockArgument>(returnValue)) { |
| // Trivial return of input. Just pass it through. |
| needsBarrier = blockArg.getOwner() != entryBlock; |
| } |
| if (needsBarrier) { |
| Value source = convertToBuiltinTensor(postambleBuilder, returnValue); |
| addBarrierInput(source, /*storage=*/BlockArgument{}, torchType, |
| returnIndex); |
| } |
| break; |
| } |
| case TypeDisposition::TORCH_PRIMITIVE: { |
| Location loc = returnValue.getLoc(); |
| if (isa<Torch::BoolType>(torchType)) { |
| newReturnOperands.back() = |
| postambleBuilder.create<TorchConversion::ToI1Op>(loc, returnValue); |
| } else if (isa<Torch::FloatType>(torchType)) { |
| newReturnOperands.back() = |
| postambleBuilder.create<TorchConversion::ToF64Op>(loc, returnValue); |
| } else if (isa<Torch::IntType>(torchType)) { |
| newReturnOperands.back() = |
| postambleBuilder.create<TorchConversion::ToI64Op>(loc, returnValue); |
| } else if (isa<Torch::GeneratorType>(torchType)) { |
| newReturnOperands.back() = |
| postambleBuilder.create<TorchConversion::GeneratorToI64Op>( |
| loc, returnValue); |
| } else { |
| emitError(loc) << "unhandled torch primitive materialization: " |
| << torchType; |
| return failure(); |
| } |
| break; |
| } |
| default: { |
| // Non-tensor/converting. Just preserve. |
| } |
| } |
| } |
| |
| // Emit the barrier and exports. |
| // If any of the exports are in-place we need to alias their storage to the |
| // provided buffers. |
| Value coarseSignalFence = |
| entryBlock->getArgument(entryBlock->getNumArguments() - 1); |
| if (barrierInputs.empty()) { |
| postambleBuilder.create<IREE::HAL::FenceSignalOp>(funcOp.getLoc(), |
| coarseSignalFence); |
| } else { |
| SmallVector<Value> aliasedResults; |
| for (auto [barrierInput, meta] : |
| llvm::zip_equal(barrierInputs, barrierResultMeta)) { |
| if (meta.storage) { |
| // Use the wait fence indicating when the storage is available for |
| // mutation. We need to ensure that no writes are made to the storage |
| // until it indicates it's safe to do so. |
| auto storageAffinityAttr = |
| getTorchArgAttr(meta.storage, "iree.abi.affinity"); |
| auto waitSignalFences = getEnclosingWaitSignalFences(meta.storage); |
| assert(waitSignalFences && "async function missing fences"); |
| Value waitFence = waitSignalFences->first; |
| auto barrierInputDims = IREE::Util::buildDynamicDimsForValue( |
| barrierInput.getLoc(), barrierInput, postambleBuilder); |
| aliasedResults.push_back( |
| postambleBuilder.create<IREE::HAL::TensorAliasOp>( |
| barrierInput.getLoc(), barrierInput.getType(), barrierInput, |
| barrierInputDims, meta.storage, waitFence, |
| storageAffinityAttr)); |
| } else { |
| aliasedResults.push_back(barrierInput); |
| } |
| } |
| auto barrierOp = postambleBuilder.create<IREE::HAL::TensorBarrierOp>( |
| funcOp.getLoc(), aliasedResults, coarseSignalFence); |
| for (auto [barrierResult, meta] : |
| llvm::zip_equal(barrierOp.getResults(), barrierResultMeta)) { |
| Attribute exportAffinityAttr; |
| if (meta.storage) { |
| exportAffinityAttr = getTorchArgAttr(meta.storage, "iree.abi.affinity"); |
| } else if (meta.returnIndex >= 0) { |
| exportAffinityAttr = |
| getTorchResultAttr(meta.returnIndex, "iree.abi.affinity"); |
| } |
| Value exportedValue = postambleBuilder.create<IREE::HAL::TensorExportOp>( |
| funcOp.getLoc(), |
| postambleBuilder.getType<IREE::HAL::BufferViewType>(), barrierResult, |
| TypeAttr::get(barrierResult.getType()), /*name=*/nullptr, |
| exportAffinityAttr); |
| if (meta.returnIndex >= 0) { |
| newReturnOperands[meta.returnIndex] = exportedValue; |
| } |
| } |
| } |
| |
| // New return operands are all collected. |
| returnOp->setOperands(newReturnOperands); |
| |
| return success(); |
| } |
| |
| class OriginalUses { |
| public: |
| OriginalUses(Value value) { |
| for (auto &use : value.getUses()) { |
| originalUses.push_back(&use); |
| } |
| } |
| |
| void assign(Value newValue) { |
| for (OpOperand *originalUse : originalUses) { |
| originalUse->assign(newValue); |
| } |
| } |
| |
| private: |
| SmallVector<OpOperand *> originalUses; |
| }; |
| |
| LogicalResult ConvertedAsyncFunctionInfo::convertImmutableTensorArg( |
| BlockArgument argValue, Type torchType, OpBuilder &builder) { |
| Location loc = argValue.getLoc(); |
| |
| // If the arg is just directly returned, then don't do anything special with |
| // it. |
| bool hasNonTrivialUse = false; |
| for (auto *userOp : argValue.getUsers()) { |
| if (isa<IREE::Util::ReturnOp>(userOp)) |
| continue; |
| hasNonTrivialUse = true; |
| } |
| if (!hasNonTrivialUse) |
| return success(); |
| |
| // Remember original uses so we can redirect them. |
| OriginalUses originalUses(argValue); |
| |
| // The type can either be a builtin TensorType or a Torch::ValueTensorType. |
| // OpBuilder |
| TensorType builtinTensorType; |
| if (auto tType = dyn_cast<TensorType>(torchType)) { |
| builtinTensorType = tType; |
| } else if (auto vtType = dyn_cast<Torch::ValueTensorType>(torchType)) { |
| builtinTensorType = vtType.toBuiltinTensor(); |
| if (auto intTy = |
| dyn_cast<IntegerType>(builtinTensorType.getElementType())) { |
| builtinTensorType = builtinTensorType.clone( |
| builder.getIntegerType(intTy.getIntOrFloatBitWidth())); |
| } |
| } else { |
| return emitError(loc) << "unsupported immutable tensor argument: " |
| << torchType; |
| } |
| |
| // Propagate explicit affinities to the read. |
| auto affinityAttr = getTorchArgAttr(argValue, "iree.abi.affinity"); |
| |
| auto waitSignalFences = getEnclosingWaitSignalFences(argValue); |
| assert(waitSignalFences && "async function missing fences"); |
| Value waitFence = waitSignalFences->first; |
| Value importedTensor = builder.create<IREE::HAL::TensorImportOp>( |
| loc, builtinTensorType, argValue, TypeAttr::get(builtinTensorType), |
| waitFence, |
| /*name=*/nullptr, affinityAttr); |
| if (builtinTensorType != torchType) { |
| importedTensor = builder.create<TorchConversion::FromBuiltinTensorOp>( |
| loc, torchType, importedTensor); |
| } |
| |
| originalUses.assign(importedTensor); |
| return success(); |
| } |
| |
| LogicalResult ConvertedAsyncFunctionInfo::convertMutableTensorArg( |
| BlockArgument argValue, Type torchType, OpBuilder &builder) { |
| Location loc = argValue.getLoc(); |
| auto fences = getEnclosingWaitSignalFences(argValue); |
| assert(fences && "could not find async fences on func"); |
| TensorType builtinTensorType; |
| if (auto t = dyn_cast<TensorType>(torchType)) { |
| builtinTensorType = t; |
| } else { |
| builtinTensorType = cast<Torch::NonValueTensorType>(torchType) |
| .getWithValueSemantics() |
| .toBuiltinTensor(); |
| } |
| |
| // Propagate explicit affinities to the read and write. |
| auto affinityAttr = getTorchArgAttr(argValue, "iree.abi.affinity"); |
| |
| // There are only a small set of possible users of a mutable tensor. |
| // Handle them by operation here. |
| SmallVector<Operation *> users(argValue.getUsers()); |
| for (auto *userOp : users) { |
| IRRewriter rewriter(loc.getContext()); |
| rewriter.setInsertionPoint(userOp); |
| if (auto copyToVtOp = dyn_cast<Torch::CopyToValueTensorOp>(userOp)) { |
| Value imported = rewriter.create<IREE::HAL::TensorImportOp>( |
| loc, builtinTensorType, argValue, |
| /*target_encoding=*/TypeAttr::get(builtinTensorType), |
| /*wait_fence*/ fences->first, |
| /*name=*/nullptr, affinityAttr); |
| rewriter.replaceOpWithNewOp<TorchConversion::FromBuiltinTensorOp>( |
| userOp, copyToVtOp.getResult().getType(), imported); |
| } else if (auto overwriteOp = |
| dyn_cast<Torch::OverwriteTensorContentsOp>(userOp)) { |
| Value overwriteValue = |
| convertToBuiltinTensor(rewriter, overwriteOp.getValue()); |
| addBarrierInput(overwriteValue, /*storage=*/argValue, torchType, |
| /*returnIndex=*/-1); |
| rewriter.eraseOp(overwriteOp); |
| } else { |
| return emitError(userOp->getLoc()) |
| << "unsupported operation on coarse signaling mutable tensor: " |
| << *userOp; |
| } |
| } |
| |
| return success(); |
| } |
| |
| void retainFunctionAttributes(Operation *srcOp, IREE::Util::FuncOp destOp) { |
| // Allowlist of function attributes to retain when importing funcs. |
| constexpr const char *kRetainedAttributes[] = { |
| "iree.reflection", |
| }; |
| auto retainedAttributes = ArrayRef<const char *>( |
| kRetainedAttributes, |
| sizeof(kRetainedAttributes) / sizeof(kRetainedAttributes[0])); |
| for (auto retainAttrName : retainedAttributes) { |
| StringRef attrName(retainAttrName); |
| Attribute attr = srcOp->getAttr(attrName); |
| if (attr) |
| destOp->setAttr(attrName, attr); |
| } |
| } |
| |
| void createCoarseFencesSyncWrapper(StringRef syncFunctionName, |
| IREE::Util::FuncOp asyncFuncOp, |
| IRRewriter &rewriter) { |
| Location loc = asyncFuncOp.getLoc(); |
| // The coarse fences wrapper has the same signature as the async variant |
| // but with the last two inputs (wait, signal fence) sliced off. |
| FunctionType asyncFuncType = asyncFuncOp.getFunctionType(); |
| SmallVector<Type> inputTypes(asyncFuncType.getInputs().begin(), |
| asyncFuncType.getInputs().end() - 2); |
| |
| // Create the function. |
| auto syncFuncType = rewriter.getType<mlir::FunctionType>( |
| inputTypes, asyncFuncType.getResults()); |
| auto syncFuncOp = |
| rewriter.create<IREE::Util::FuncOp>(loc, syncFunctionName, syncFuncType, |
| /*tiedOperandsAttr=*/nullptr); |
| syncFuncOp.setSymVisibilityAttr(asyncFuncOp.getSymVisibilityAttr()); |
| retainFunctionAttributes(asyncFuncOp, syncFuncOp); |
| syncFuncOp->setAttr("iree.abi.stub", rewriter.getUnitAttr()); |
| if (auto affinityAttr = asyncFuncOp->getAttr("iree.abi.affinity")) { |
| syncFuncOp->setAttr("iree.abi.affinity", affinityAttr); |
| } |
| Block *entryBlock = syncFuncOp.addEntryBlock(); |
| OpBuilder::InsertionGuard guard(rewriter); |
| rewriter.setInsertionPointToEnd(entryBlock); |
| |
| // HACK: this is relying on the fact that there's only one HAL device. |
| // We should instead have a way of creating fences on the device that |
| // is used to produce the tensors we're wrapping. |
| // |
| // TODO(multi-device): emit get with derived ordinal or lookup with attr. We |
| // could always say device 0 for now but could instead look for an |
| // iree.abi.affinity/iree.abi.device/etc. |
| Value timeoutMillis = rewriter.create<arith::ConstantIntOp>(loc, -1, 32); |
| Value device = IREE::HAL::DeviceType::resolveAny(loc, rewriter); |
| Value waitFence = rewriter.create<IREE::Util::NullOp>( |
| loc, rewriter.getType<IREE::HAL::FenceType>()); |
| Value signalFence = rewriter.create<IREE::HAL::FenceCreateOp>( |
| loc, rewriter.getType<IREE::HAL::FenceType>(), device, |
| IREE::HAL::FenceFlagBitfield::None); |
| |
| SmallVector<Value> callOperands(entryBlock->getArguments()); |
| callOperands.push_back(waitFence); |
| callOperands.push_back(signalFence); |
| std::optional<ArrayAttr> targetTiedOperands = asyncFuncOp.getTiedOperands(); |
| auto callResults = |
| rewriter |
| .create<IREE::Util::CallOp>(loc, asyncFuncOp, callOperands, |
| targetTiedOperands ? *targetTiedOperands |
| : ArrayAttr{}) |
| .getResults(); |
| |
| // Wait forever for signal. |
| rewriter.create<IREE::HAL::FenceAwaitOp>(loc, rewriter.getI32Type(), |
| timeoutMillis, signalFence); |
| |
| rewriter.create<IREE::Util::ReturnOp>(loc, callResults); |
| } |
| |
| } // namespace |
| |
| class FuncConversionPass final |
| : public impl::FuncConversionPassBase<FuncConversionPass> { |
| public: |
| void getDependentDialects(DialectRegistry ®istry) const override { |
| registry.insert<mlir::tensor::TensorDialect>(); |
| registry.insert<IREE::HAL::HALDialect>(); |
| registry.insert<IREE::Util::UtilDialect>(); |
| registry.insert<TorchConversion::TorchConversionDialect>(); |
| } |
| |
| void runOnOperation() override { |
| auto moduleOp = getOperation(); |
| |
| // Convert all functions in the module to IREE funcs. In this stage, |
| // we convert contained return ops and argument/result types, but we have |
| // not yet converted anything "on the inside". Therefore, it is pretty |
| // likely the functions are still illegal. |
| SmallVector<Operation *> eraseFuncOps; |
| std::vector<ConvertedAsyncFunctionInfo> convertedFuncInfos; |
| for (auto funcOp : moduleOp.getOps<func::FuncOp>()) { |
| if (!shouldConvertFunc(funcOp)) |
| continue; |
| ConvertedAsyncFunctionInfo &convertedFuncInfo = |
| convertedFuncInfos.emplace_back(); |
| if (failed(convertFuncOp(funcOp, convertedFuncInfo))) { |
| signalPassFailure(); |
| return; |
| } |
| eraseFuncOps.push_back(funcOp); |
| } |
| for (auto op : eraseFuncOps) { |
| op->erase(); |
| } |
| |
| // Now post-process async functions. |
| for (auto &info : convertedFuncInfos) { |
| if (failed(info.postProcess())) { |
| signalPassFailure(); |
| return; |
| } |
| } |
| } |
| |
| bool shouldConvertFunc(func::FuncOp torchFunc) { |
| // For now, we don't touch externals and assume they are in the proper |
| // calling convention. In the future, we may support "torch externals" |
| // which we convert to mate up with a torch module. We can remove/adapt |
| // this when that is elaborated. |
| if (torchFunc.isExternal()) |
| return false; |
| |
| // Something has already converted this and told us not to touch it. |
| if (torchFunc->hasAttr("iree.abi.stub")) |
| return false; |
| |
| return true; |
| } |
| |
| LogicalResult convertFuncOp(func::FuncOp torchFunc, |
| ConvertedAsyncFunctionInfo &convertedFuncInfo) { |
| IRRewriter rewriter(torchFunc.getContext()); |
| rewriter.setInsertionPoint(torchFunc); |
| Location loc = torchFunc.getLoc(); |
| // Determine whether to build pure async or async + sync wrapper. |
| bool generateSyncWrapper = true; |
| StringRef originalName = torchFunc.getName(); |
| std::string asyncFunctionName = originalName.str(); |
| if (generateSyncWrapper) { |
| asyncFunctionName.append("$async"); |
| } |
| |
| // Stash arg/result attrs so they can be referenced during conversion. |
| torchFunc.getAllArgAttrs(convertedFuncInfo.torchArgAttrs); |
| torchFunc.getAllResultAttrs(convertedFuncInfo.torchResultAttrs); |
| |
| // Convert function signature. |
| Type fenceType = rewriter.getType<IREE::HAL::FenceType>(); |
| FunctionType torchFuncType = torchFunc.getFunctionType(); |
| convertedFuncInfo.torchInputTypes.append(torchFuncType.getInputs().begin(), |
| torchFuncType.getInputs().end()); |
| convertedFuncInfo.torchResultTypes.append( |
| torchFuncType.getResults().begin(), torchFuncType.getResults().end()); |
| // For the coarse-fences ABI, we add two fences to the end. Treat these as |
| // original types so that the lists line up. |
| convertedFuncInfo.torchInputTypes.push_back(fenceType); |
| convertedFuncInfo.torchInputTypes.push_back(fenceType); |
| SmallVector<Type> ireeInputTypes(convertedFuncInfo.torchInputTypes); |
| SmallVector<Type> ireeResultTypes(convertedFuncInfo.torchResultTypes); |
| convertedFuncInfo.inputDispositions.resize(ireeInputTypes.size()); |
| convertedFuncInfo.resultDispositions.resize(ireeResultTypes.size()); |
| |
| for (size_t i = 0; i < convertedFuncInfo.torchInputTypes.size(); ++i) { |
| if (failed(convertType(loc, convertedFuncInfo.torchInputTypes[i], |
| ireeInputTypes[i], |
| convertedFuncInfo.inputDispositions[i]))) |
| return failure(); |
| } |
| for (size_t i = 0; i < convertedFuncInfo.torchResultTypes.size(); ++i) { |
| if (failed(convertType(loc, convertedFuncInfo.torchResultTypes[i], |
| ireeResultTypes[i], |
| convertedFuncInfo.resultDispositions[i]))) |
| return failure(); |
| } |
| |
| // Build tied operands index mapping results back to operands. |
| SmallVector<int64_t> tiedOperands; |
| bool anyTiedOperands = false; |
| for (unsigned i = 0; i < torchFuncType.getNumResults(); ++i) { |
| auto tiedAttr = |
| torchFunc.getResultAttrOfType<IntegerAttr>(i, "iree.abi.tied"); |
| if (tiedAttr) { |
| tiedOperands.push_back(tiedAttr.getInt()); |
| anyTiedOperands = true; |
| } else { |
| tiedOperands.push_back(-1); |
| } |
| } |
| auto tiedOperandsAttr = anyTiedOperands |
| ? rewriter.getIndexArrayAttr(tiedOperands) |
| : ArrayAttr{}; |
| |
| // Create new func. |
| FunctionType asyncFuncType = |
| FunctionType::get(loc.getContext(), ireeInputTypes, ireeResultTypes); |
| auto asyncFuncOp = rewriter.create<IREE::Util::FuncOp>( |
| torchFunc.getLoc(), asyncFunctionName, asyncFuncType, tiedOperandsAttr); |
| convertedFuncInfo.funcOp = asyncFuncOp; |
| asyncFuncOp.setSymVisibilityAttr(torchFunc.getSymVisibilityAttr()); |
| // Handle defacto attrs to specialized ones. |
| asyncFuncOp.setInliningPolicyAttr( |
| rewriter.getAttr<IREE::Util::InlineNeverAttr>()); |
| retainFunctionAttributes(torchFunc, asyncFuncOp); |
| asyncFuncOp->setAttr("iree.abi.stub", rewriter.getUnitAttr()); |
| asyncFuncOp->setAttr("iree.abi.model", |
| rewriter.getStringAttr("coarse-fences")); |
| if (auto affinityAttr = torchFunc->getAttr("iree.abi.affinity")) { |
| asyncFuncOp->setAttr("iree.abi.affinity", affinityAttr); |
| } |
| rewriter.inlineRegionBefore( |
| torchFunc.getBody(), asyncFuncOp.getFunctionBody(), asyncFuncOp.end()); |
| |
| // Convert block arguments. |
| Block *entryBlock = &asyncFuncOp.getBlocks().front(); |
| for (size_t i = 0; i < ireeInputTypes.size(); ++i) { |
| // Add if we have extended the list. |
| if (i >= entryBlock->getNumArguments()) { |
| entryBlock->addArgument(ireeInputTypes[i], loc); |
| continue; |
| } |
| // Convert. |
| entryBlock->getArgument(i).setType(ireeInputTypes[i]); |
| } |
| |
| // Replace return ops. |
| asyncFuncOp->walk([&](func::ReturnOp returnOp) { |
| rewriter.setInsertionPoint(returnOp); |
| auto ireeReturnOp = rewriter.replaceOpWithNewOp<IREE::Util::ReturnOp>( |
| returnOp, returnOp.getOperands()); |
| convertedFuncInfo.returnOps.push_back(ireeReturnOp); |
| }); |
| |
| // Create the sync variant. |
| rewriter.setInsertionPoint(torchFunc); |
| createCoarseFencesSyncWrapper(originalName, asyncFuncOp, rewriter); |
| return success(); |
| } |
| |
| LogicalResult convertType(Location loc, Type torchType, Type &ireeType, |
| TypeDisposition &disp) { |
| if (isa<TensorType, Torch::ValueTensorType>(torchType)) { |
| ireeType = IREE::HAL::BufferViewType::get(torchType.getContext()); |
| disp = TypeDisposition::IMMUTABLE_TENSOR; |
| return success(); |
| } |
| |
| if (isa<Torch::NonValueTensorType>(torchType)) { |
| ireeType = IREE::HAL::BufferViewType::get(torchType.getContext()); |
| disp = TypeDisposition::MUTABLE_TENSOR; |
| return success(); |
| } |
| |
| if (isa<IREE::HAL::FenceType>(torchType)) { |
| ireeType = torchType; |
| disp = TypeDisposition::FENCE; |
| return success(); |
| } |
| |
| if (isa<Torch::BoolType>(torchType)) { |
| ireeType = IntegerType::get(torchType.getContext(), 1); |
| disp = TypeDisposition::TORCH_PRIMITIVE; |
| return success(); |
| } |
| |
| if (isa<Torch::IntType, Torch::GeneratorType>(torchType)) { |
| ireeType = IntegerType::get(torchType.getContext(), 64); |
| disp = TypeDisposition::TORCH_PRIMITIVE; |
| return success(); |
| } |
| |
| if (isa<Torch::FloatType>(torchType)) { |
| ireeType = Float64Type::get(torchType.getContext()); |
| disp = TypeDisposition::TORCH_PRIMITIVE; |
| return success(); |
| } |
| |
| if (isa<IntegerType, FloatType, IndexType>(torchType)) { |
| ireeType = torchType; |
| disp = TypeDisposition::PASSTHROUGH; |
| return success(); |
| } |
| |
| return emitError(loc) << "unhandled torch type: " << torchType; |
| } |
| }; |
| |
| } // namespace mlir::iree_compiler::TorchInput |
