Optional sync func conversion (#24747)
Adds an option to make async entry-point generation optional for torch
programs.
The coarse-fences ABI presumes queue-ordered execution and adds fence
handling into every entry point. When compiling for the inline HAL,
execution is synchronous and that machinery is not required.
Assisted-by: Claude code
---------
Signed-off-by: Jelle Schuhmacher <schuehmacher@roofline.ai>
Co-authored-by: ziereis <44057120+ziereis@users.noreply.github.com>
diff --git a/compiler/plugins/input/Torch/InputConversion/BUILD.bazel b/compiler/plugins/input/Torch/InputConversion/BUILD.bazel
index 3485639..d71c6b0 100644
--- a/compiler/plugins/input/Torch/InputConversion/BUILD.bazel
+++ b/compiler/plugins/input/Torch/InputConversion/BUILD.bazel
@@ -56,6 +56,7 @@
"FuncConversion.cpp",
"Passes.cpp",
"SetStrictSymbolicShapes.cpp",
+ "SyncFuncConversion.cpp",
],
hdrs = ["Passes.h"],
deps = [
diff --git a/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt b/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt
index 301f62c..0e494ab 100644
--- a/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt
+++ b/compiler/plugins/input/Torch/InputConversion/CMakeLists.txt
@@ -45,6 +45,7 @@
"FuncConversion.cpp"
"Passes.cpp"
"SetStrictSymbolicShapes.cpp"
+ "SyncFuncConversion.cpp"
DEPS
::PassHeaders
::PassesIncGen
diff --git a/compiler/plugins/input/Torch/InputConversion/Passes.cpp b/compiler/plugins/input/Torch/InputConversion/Passes.cpp
index ffd5c66..350885c 100644
--- a/compiler/plugins/input/Torch/InputConversion/Passes.cpp
+++ b/compiler/plugins/input/Torch/InputConversion/Passes.cpp
@@ -86,7 +86,13 @@
// differently and would not be subject to inlining.
pm.addPass(mlir::createInlinerPass());
- pm.addPass(createFuncConversionPass({options.externalizeTransients}));
+ if (options.emitAsyncEntryPoints) {
+ pm.addPass(createFuncConversionPass({options.externalizeTransients}));
+ } else {
+ // Sync-only entry points (e.g. for the inline HAL). Externalized
+ // transients require the coarse-fences ABI and do not apply here.
+ pm.addPass(createSyncFuncConversionPass());
+ }
pm.addNestedPass<IREE::Util::FuncOp>(createCanonicalizerPass());
pm.addPass(createSymbolDCEPass());
diff --git a/compiler/plugins/input/Torch/InputConversion/Passes.h b/compiler/plugins/input/Torch/InputConversion/Passes.h
index 5ec0799..bf9592a 100644
--- a/compiler/plugins/input/Torch/InputConversion/Passes.h
+++ b/compiler/plugins/input/Torch/InputConversion/Passes.h
@@ -52,6 +52,13 @@
Option<bool> enableShapeRefinement{*this, "enable-shape-refinement",
llvm::cl::desc("Enable shape refinement"),
llvm::cl::init(false)};
+ Option<bool> emitAsyncEntryPoints{
+ *this, "emit-async-entry-points",
+ llvm::cl::desc(
+ "Generate async functions with coarse-fences ABI and sync wrapper. "
+ "When false, generates only lightweight sync functions without "
+ "mutable tensor support."),
+ llvm::cl::init(true)};
};
// Creates a pipeline that lowers from the torch backend contract to IREE.
diff --git a/compiler/plugins/input/Torch/InputConversion/Passes.td b/compiler/plugins/input/Torch/InputConversion/Passes.td
index a868d4b..b74adf1 100644
--- a/compiler/plugins/input/Torch/InputConversion/Passes.td
+++ b/compiler/plugins/input/Torch/InputConversion/Passes.td
@@ -50,4 +50,17 @@
"memory and must be provided by the user.">];
}
+def SyncFuncConversionPass :
+ Pass<"torch-iree-sync-func-conversion", "ModuleOp"> {
+ let summary = "Finalizes conversion from torch to IREE with a sync-only ABI";
+ let description = [{
+ Synchronous-only alternative to torch-iree-func-conversion for targets
+ where the coarse-fences ABI is undesirable, such as programs compiled for
+ the inline HAL. Entry points keep their original name and a plain builtin
+ tensor ABI: no fences, no HAL imports/exports, and no $async variant.
+ Mutable tensor arguments require in-place aliasing through the HAL and are
+ not supported.
+ }];
+}
+
#endif // IREE_COMPILER_PLUGINS_INPUT_TORCH_INPUTCONVERSION_PASSES
diff --git a/compiler/plugins/input/Torch/InputConversion/SyncFuncConversion.cpp b/compiler/plugins/input/Torch/InputConversion/SyncFuncConversion.cpp
new file mode 100644
index 0000000..36605b5
--- /dev/null
+++ b/compiler/plugins/input/Torch/InputConversion/SyncFuncConversion.cpp
@@ -0,0 +1,415 @@
+// Copyright 2026 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/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.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_SYNCFUNCCONVERSIONPASS
+#include "compiler/plugins/input/Torch/InputConversion/Passes.h.inc"
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Synchronous-only variant of FuncConversion (see FuncConversion.cpp for the
+// coarse-fences ABI it is an alternative to). Functions are rewritten in
+// place to a plain builtin tensor ABI intended for synchronous invocation
+// (e.g. via the inline HAL): immutable tensors become builtin tensors and
+// torch primitives their builtin scalar equivalents, with torch_c conversion
+// ops materialized at the edges. No fences, no HAL imports/exports, and no
+// $async variant are created. Mutable tensors require in-place aliasing
+// through the HAL and are rejected.
+//===----------------------------------------------------------------------===//
+
+TensorType convertToBuiltinTensorType(OpBuilder &builder,
+ Torch::ValueTensorType vtensorType) {
+ TensorType builtinType = vtensorType.toBuiltinTensor();
+ if (auto intTy = dyn_cast<IntegerType>(builtinType.getElementType())) {
+ builtinType = builtinType.clone(
+ builder.getIntegerType(intTy.getIntOrFloatBitWidth()));
+ }
+ return builtinType;
+}
+
+Value convertToBuiltinTensor(OpBuilder &builder, Value possibleTorchTensor) {
+ Type ty = possibleTorchTensor.getType();
+ if (isa<TensorType>(ty)) {
+ return possibleTorchTensor;
+ }
+
+ if (auto defining = dyn_cast_if_present<TorchConversion::FromBuiltinTensorOp>(
+ possibleTorchTensor.getDefiningOp())) {
+ return defining.getOperand();
+ }
+
+ TensorType builtinType =
+ convertToBuiltinTensorType(builder, cast<Torch::ValueTensorType>(ty));
+ return TorchConversion::ToBuiltinTensorOp::create(
+ builder, possibleTorchTensor.getLoc(), builtinType, possibleTorchTensor);
+}
+
+enum class TypeDisposition {
+ IMMUTABLE_TENSOR,
+ TORCH_PRIMITIVE,
+ PASSTHROUGH,
+};
+
+struct ConvertedSyncFunctionInfo {
+ IREE::Util::FuncOp funcOp;
+ SmallVector<IREE::Util::ReturnOp> returnOps;
+ SmallVector<Type> torchInputTypes;
+ SmallVector<Type> torchResultTypes;
+ SmallVector<TypeDisposition> inputDispositions;
+ SmallVector<TypeDisposition> resultDispositions;
+
+ LogicalResult postProcess();
+ LogicalResult convertImmutableTensorArg(BlockArgument argValue,
+ Type torchType, OpBuilder &builder);
+};
+
+LogicalResult ConvertedSyncFunctionInfo::postProcess() {
+ if (funcOp.isExternal()) {
+ return success();
+ }
+
+ if (returnOps.size() != 1) {
+ // Multi-exit funcs could be supported by materializing the result
+ // conversions at each return; restricted to single-exit for parity with
+ // the coarse-fences conversion.
+ 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::TORCH_PRIMITIVE: {
+ Location loc = argValue.getLoc();
+ Operation *convertUser = nullptr;
+ Value convertResult;
+ if (isa<Torch::BoolType>(torchType)) {
+ convertUser =
+ TorchConversion::FromI1Op::create(preambleBuilder, loc, argValue);
+ convertResult = convertUser->getResult(0);
+ } else if (isa<Torch::FloatType>(torchType)) {
+ convertUser =
+ TorchConversion::FromF64Op::create(preambleBuilder, loc, argValue);
+ convertResult = convertUser->getResult(0);
+ } else if (isa<Torch::IntType>(torchType)) {
+ convertUser =
+ TorchConversion::FromI64Op::create(preambleBuilder, 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;
+ }
+ }
+
+ // Materialize result 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)) {
+ newReturnOperands.emplace_back(returnValue);
+ switch (disp) {
+ case TypeDisposition::IMMUTABLE_TENSOR: {
+ newReturnOperands.back() =
+ convertToBuiltinTensor(postambleBuilder, returnValue);
+ break;
+ }
+ case TypeDisposition::TORCH_PRIMITIVE: {
+ Location loc = returnValue.getLoc();
+ if (isa<Torch::BoolType>(torchType)) {
+ newReturnOperands.back() =
+ TorchConversion::ToI1Op::create(postambleBuilder, loc, returnValue);
+ } else if (isa<Torch::FloatType>(torchType)) {
+ newReturnOperands.back() = TorchConversion::ToF64Op::create(
+ postambleBuilder, loc, returnValue);
+ } else if (isa<Torch::IntType>(torchType)) {
+ newReturnOperands.back() = TorchConversion::ToI64Op::create(
+ postambleBuilder, loc, returnValue);
+ } else if (isa<Torch::GeneratorType>(torchType)) {
+ newReturnOperands.back() = TorchConversion::GeneratorToI64Op::create(
+ postambleBuilder, loc, returnValue);
+ } else {
+ emitError(loc) << "unhandled torch primitive materialization: "
+ << torchType;
+ return failure();
+ }
+ break;
+ }
+ case TypeDisposition::PASSTHROUGH:
+ // Do nothing.
+ break;
+ }
+ }
+ returnOp->setOperands(newReturnOperands);
+
+ return success();
+}
+
+LogicalResult ConvertedSyncFunctionInfo::convertImmutableTensorArg(
+ BlockArgument argValue, Type torchType, OpBuilder &builder) {
+ // Already a builtin tensor: nothing to materialize.
+ if (isa<TensorType>(torchType)) {
+ return success();
+ }
+
+ if (!isa<Torch::ValueTensorType>(torchType)) {
+ return emitError(argValue.getLoc())
+ << "unsupported immutable tensor argument: " << torchType;
+ }
+
+ // If the arg is just directly returned, then don't do anything special with
+ // it: the postamble passes it through unconverted.
+ bool hasNonTrivialUse = false;
+ for (auto *userOp : argValue.getUsers()) {
+ if (isa<IREE::Util::ReturnOp>(userOp)) {
+ continue;
+ }
+ hasNonTrivialUse = true;
+ }
+ if (!hasNonTrivialUse) {
+ return success();
+ }
+ Value converted = TorchConversion::FromBuiltinTensorOp::create(
+ builder, argValue.getLoc(), torchType, argValue);
+ argValue.replaceAllUsesExcept(converted, converted.getDefiningOp());
+ 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",
+ };
+ for (const char *attrName : kRetainedAttributes) {
+ if (Attribute attr = srcOp->getAttr(attrName)) {
+ destOp->setAttr(attrName, attr);
+ }
+ }
+}
+
+class SyncFuncConversionPass final
+ : public impl::SyncFuncConversionPassBase<SyncFuncConversionPass> {
+public:
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ 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<ConvertedSyncFunctionInfo> convertedFuncInfos;
+ for (auto funcOp : moduleOp.getOps<func::FuncOp>()) {
+ if (!shouldConvertFunc(funcOp)) {
+ continue;
+ }
+ ConvertedSyncFunctionInfo &convertedFuncInfo =
+ convertedFuncInfos.emplace_back();
+ if (failed(convertFuncOp(funcOp, convertedFuncInfo))) {
+ signalPassFailure();
+ return;
+ }
+ eraseFuncOps.push_back(funcOp);
+ }
+ for (auto op : eraseFuncOps) {
+ op->erase();
+ }
+
+ 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,
+ ConvertedSyncFunctionInfo &convertedFuncInfo) {
+ IRRewriter rewriter(torchFunc.getContext());
+ rewriter.setInsertionPoint(torchFunc);
+ Location loc = torchFunc.getLoc();
+
+ // Convert function signature.
+ FunctionType torchFuncType = torchFunc.getFunctionType();
+ convertedFuncInfo.torchInputTypes.append(torchFuncType.getInputs().begin(),
+ torchFuncType.getInputs().end());
+ convertedFuncInfo.torchResultTypes.append(
+ torchFuncType.getResults().begin(), torchFuncType.getResults().end());
+ 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(
+ rewriter, loc, convertedFuncInfo.torchInputTypes[i],
+ ireeInputTypes[i], convertedFuncInfo.inputDispositions[i]))) {
+ return failure();
+ }
+ }
+ for (size_t i = 0; i < convertedFuncInfo.torchResultTypes.size(); ++i) {
+ if (failed(convertType(
+ rewriter, 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 with the original name.
+ FunctionType syncFuncType =
+ FunctionType::get(loc.getContext(), ireeInputTypes, ireeResultTypes);
+ auto syncFuncOp = IREE::Util::FuncOp::create(
+ rewriter, torchFunc.getLoc(), torchFunc.getName(), syncFuncType,
+ tiedOperandsAttr);
+ convertedFuncInfo.funcOp = syncFuncOp;
+ syncFuncOp.setSymVisibilityAttr(torchFunc.getSymVisibilityAttr());
+ retainFunctionAttributes(torchFunc, syncFuncOp);
+ if (auto affinityAttr = torchFunc->getAttr("iree.abi.affinity")) {
+ syncFuncOp->setAttr("iree.abi.affinity", affinityAttr);
+ }
+ rewriter.inlineRegionBefore(torchFunc.getBody(),
+ syncFuncOp.getFunctionBody(), syncFuncOp.end());
+
+ // Convert block arguments.
+ Block *entryBlock = &syncFuncOp.getBlocks().front();
+ for (size_t i = 0; i < ireeInputTypes.size(); ++i) {
+ entryBlock->getArgument(i).setType(ireeInputTypes[i]);
+ }
+
+ // Replace return ops.
+ syncFuncOp->walk([&](func::ReturnOp returnOp) {
+ rewriter.setInsertionPoint(returnOp);
+ auto ireeReturnOp = rewriter.replaceOpWithNewOp<IREE::Util::ReturnOp>(
+ returnOp, returnOp.getOperands());
+ convertedFuncInfo.returnOps.push_back(ireeReturnOp);
+ });
+ return success();
+ }
+
+ LogicalResult convertType(OpBuilder &builder, Location loc, Type torchType,
+ Type &ireeType, TypeDisposition &disp) {
+ if (isa<TensorType>(torchType)) {
+ ireeType = torchType;
+ disp = TypeDisposition::IMMUTABLE_TENSOR;
+ return success();
+ }
+
+ if (auto vtType = dyn_cast<Torch::ValueTensorType>(torchType)) {
+ ireeType = convertToBuiltinTensorType(builder, vtType);
+ disp = TypeDisposition::IMMUTABLE_TENSOR;
+ return success();
+ }
+
+ if (isa<Torch::NonValueTensorType>(torchType)) {
+ return emitError(loc)
+ << "mutable tensors are not supported by sync func conversion: "
+ << torchType;
+ }
+
+ 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
+
+} // namespace mlir::iree_compiler::TorchInput
diff --git a/compiler/plugins/input/Torch/InputConversion/test/BUILD.bazel b/compiler/plugins/input/Torch/InputConversion/test/BUILD.bazel
index 9c11115..3ea5a7e 100644
--- a/compiler/plugins/input/Torch/InputConversion/test/BUILD.bazel
+++ b/compiler/plugins/input/Torch/InputConversion/test/BUILD.bazel
@@ -26,6 +26,8 @@
"bitcast_tensor.mlir",
"func_conversion.mlir",
"func_conversion_invalid.mlir",
+ "func_conversion_sync.mlir",
+ "func_conversion_sync_invalid.mlir",
"func_conversion_transients.mlir",
"scan.mlir",
"scatter.mlir",
diff --git a/compiler/plugins/input/Torch/InputConversion/test/CMakeLists.txt b/compiler/plugins/input/Torch/InputConversion/test/CMakeLists.txt
index 1364095..e5f5c2f 100644
--- a/compiler/plugins/input/Torch/InputConversion/test/CMakeLists.txt
+++ b/compiler/plugins/input/Torch/InputConversion/test/CMakeLists.txt
@@ -24,6 +24,8 @@
"bitcast_tensor.mlir"
"func_conversion.mlir"
"func_conversion_invalid.mlir"
+ "func_conversion_sync.mlir"
+ "func_conversion_sync_invalid.mlir"
"func_conversion_transients.mlir"
"scan.mlir"
"scatter.mlir"
diff --git a/compiler/plugins/input/Torch/InputConversion/test/func_conversion_sync.mlir b/compiler/plugins/input/Torch/InputConversion/test/func_conversion_sync.mlir
new file mode 100644
index 0000000..705aa9c
--- /dev/null
+++ b/compiler/plugins/input/Torch/InputConversion/test/func_conversion_sync.mlir
@@ -0,0 +1,194 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(torch-iree-sync-func-conversion)" --allow-unregistered-dialect --split-input-file %s | FileCheck %s
+
+// Canonical test of the immutable input->compute->return case. The function
+// keeps its original name and a plain builtin tensor ABI: no HAL imports or
+// exports, no fences, and no $async variant.
+// CHECK-LABEL: @immutable_import_export
+// CHECK: util.func public @main(
+// CHECK-SAME: %arg0: tensor<4x5xi32>, %arg1: tensor<5x4xf32>) ->
+// CHECK-SAME: (tensor<4x5xi32>, tensor<5x4xf32>)
+// CHECK-DAG: %[[TORCH_ARG0:.+]] = torch_c.from_builtin_tensor %arg0 : tensor<4x5xi32> -> !torch.vtensor<[4,5],si32>
+// CHECK-DAG: %[[TORCH_ARG1:.+]] = torch_c.from_builtin_tensor %arg1 : tensor<5x4xf32> -> !torch.vtensor<[5,4],f32>
+// CHECK-DAG: %[[TORCH_RESULT0:.+]] = torch.operator "foobar0"(%[[TORCH_ARG0]])
+// CHECK-DAG: %[[TORCH_RESULT1:.+]] = torch.operator "foobar1"(%[[TORCH_ARG1]])
+// CHECK-DAG: %[[TENSOR_RESULT0:.+]] = torch_c.to_builtin_tensor %[[TORCH_RESULT0]]
+// CHECK-DAG: %[[TENSOR_RESULT1:.+]] = torch_c.to_builtin_tensor %[[TORCH_RESULT1]]
+// CHECK: util.return %[[TENSOR_RESULT0]], %[[TENSOR_RESULT1]]
+// CHECK-NOT: hal.tensor
+// CHECK-NOT: hal.fence
+// CHECK-NOT: @main$async
+builtin.module @immutable_import_export {
+func.func @main(%arg0: !torch.vtensor<[4,5],si32>, %arg1: !torch.vtensor<[5,4],f32>)
+ -> (!torch.vtensor<[4,5],si32>, !torch.vtensor<[5,4],f32>) {
+ %0 = torch.operator "foobar0"(%arg0) : (!torch.vtensor<[4,5],si32>) -> !torch.vtensor<[4,5],si32>
+ %1 = torch.operator "foobar1"(%arg1) : (!torch.vtensor<[5,4],f32>) -> !torch.vtensor<[5,4],f32>
+ return %0, %1 : !torch.vtensor<[4,5],si32>, !torch.vtensor<[5,4],f32>
+}
+}
+
+// -----
+// A trivially returned argument needs no conversion ops at all.
+// CHECK-LABEL: @return_immutable_arg
+// CHECK: util.func public @main(
+// CHECK-SAME: %arg0: tensor<4x5xi32>) -> tensor<4x5xi32>
+// CHECK-NOT: torch_c.from_builtin_tensor
+// CHECK: util.return %arg0
+// CHECK-NOT: hal.fence
+// CHECK-NOT: @main$async
+builtin.module @return_immutable_arg {
+func.func @main(%arg0: !torch.vtensor<[4,5],si32>) -> !torch.vtensor<[4,5],si32> {
+ return %arg0 : !torch.vtensor<[4,5],si32>
+}
+}
+
+// -----
+// CHECK-LABEL: @retained_attribute_reflection
+// CHECK: util.func public @main(
+// CHECK-SAME: iree.reflection = {some.attr = 4 : index}
+// CHECK-NOT: @main$async
+builtin.module @retained_attribute_reflection {
+func.func @main(%arg0: !torch.vtensor<[4,5],si32>) -> !torch.vtensor<[4,5],si32>
+ attributes {
+ iree.reflection = {
+ some.attr = 4 : index
+ }
+ }
+{
+ return %arg0 : !torch.vtensor<[4,5],si32>
+}
+}
+
+// -----
+// CHECK-LABEL: @retained_attribute_ignored
+// CHECK: util.func public @main(
+// CHECK-NOT: iree.nonretained
+builtin.module @retained_attribute_ignored {
+func.func @main(%arg0: !torch.vtensor<[4,5],si32>) -> !torch.vtensor<[4,5],si32>
+ attributes {
+ iree.nonretained = "dummy"
+ }
+{
+ return %arg0 : !torch.vtensor<[4,5],si32>
+}
+}
+
+// -----
+// CHECK-LABEL: @private_visibility
+// CHECK: util.func private @main
+builtin.module @private_visibility {
+func.func private @main(%arg0: !torch.vtensor<[4,5],si32>) -> !torch.vtensor<[4,5],si32>
+{
+ return %arg0 : !torch.vtensor<[4,5],si32>
+}
+}
+
+// -----
+// CHECK-LABEL: @tied_operand
+// CHECK: util.func public @main(%arg0: tensor<4x5xi32>) -> %arg0
+// CHECK: util.return %arg0
+// CHECK-NOT: @main$async
+builtin.module @tied_operand {
+func.func @main(%arg0: !torch.vtensor<[4,5],si32>) ->
+ (!torch.vtensor<[4,5],si32> {iree.abi.tied = 0})
+{
+ return %arg0 : !torch.vtensor<[4,5],si32>
+}
+}
+
+// -----
+// Verify that dynamic dimensions convert.
+// CHECK-LABEL: @dynamic_dims
+// CHECK: util.func public @main(%arg0: tensor<4x?xi32>) -> tensor<4x?xi32>
+builtin.module @dynamic_dims {
+func.func @main(%arg0: !torch.vtensor<[4,?],si32>) -> !torch.vtensor<[4,?],si32> {
+ %0 = torch.operator "foobar0"(%arg0) : (!torch.vtensor<[4,?],si32>) -> !torch.vtensor<[4,?],si32>
+ return %0 : !torch.vtensor<[4,?],si32>
+}
+}
+
+// -----
+// CHECK-LABEL: @torch_bool_return
+// CHECK: torch_c.to_i1
+// CHECK: util.return {{.*}} : i1
+// CHECK-NOT: hal.fence
+module @torch_bool_return {
+ func.func @main() -> !torch.bool {
+ %0 = torch.operator "some.primitive"() : () -> !torch.bool
+ return %0 : !torch.bool
+ }
+}
+
+// -----
+// CHECK-LABEL: @torch_int_return
+// CHECK: torch_c.to_i64
+// CHECK: util.return {{.*}} : i64
+// CHECK-NOT: hal.fence
+module @torch_int_return {
+ func.func @main() -> !torch.int {
+ %0 = torch.operator "some.primitive"() : () -> !torch.int
+ return %0 : !torch.int
+ }
+}
+
+// -----
+// CHECK-LABEL: @torch_float_return
+// CHECK: torch_c.to_f64
+// CHECK: util.return {{.*}} : f64
+module @torch_float_return {
+ func.func @main() -> !torch.float {
+ %0 = torch.operator "some.primitive"() : () -> !torch.float
+ return %0 : !torch.float
+ }
+}
+
+// -----
+// CHECK-LABEL: @torch_generator_return
+// CHECK: torch_c.generator_to_i64
+// CHECK: util.return {{.*}} : i64
+module @torch_generator_return {
+ func.func @main() -> !torch.Generator {
+ %0 = torch.operator "some.primitive"() : () -> !torch.Generator
+ return %0 : !torch.Generator
+ }
+}
+
+// -----
+// CHECK-LABEL: @torch_bool_arg
+// CHECK: torch_c.from_i1 %arg0
+module @torch_bool_arg {
+ func.func @main(%arg0 : !torch.bool) -> (!torch.vtensor<[1],f32>) {
+ %0 = torch.operator "some.primitive"(%arg0) : (!torch.bool) -> (!torch.vtensor<[1],f32>)
+ return %0 : !torch.vtensor<[1],f32>
+ }
+}
+
+// -----
+// CHECK-LABEL: @torch_int_arg
+// CHECK: torch_c.from_i64 %arg0
+module @torch_int_arg {
+ func.func @main(%arg0 : !torch.int) -> (!torch.vtensor<[1],f32>) {
+ %0 = torch.operator "some.primitive"(%arg0) : (!torch.int) -> (!torch.vtensor<[1],f32>)
+ return %0 : !torch.vtensor<[1],f32>
+ }
+}
+
+// -----
+// CHECK-LABEL: @torch_float_arg
+// CHECK: torch_c.from_f64 %arg0
+module @torch_float_arg {
+ func.func @main(%arg0 : !torch.float) -> (!torch.vtensor<[1],f32>) {
+ %0 = torch.operator "some.primitive"(%arg0) : (!torch.float) -> (!torch.vtensor<[1],f32>)
+ return %0 : !torch.vtensor<[1],f32>
+ }
+}
+
+// -----
+// Builtin scalar arguments and results pass through unchanged.
+// CHECK-LABEL: @builtin_scalars
+// CHECK: util.func public @main(%arg0: index, %arg1: i32, %arg2: f32) -> i32
+module @builtin_scalars {
+ func.func @main(%arg0 : index, %arg1 : i32, %arg2 : f32) -> i32 {
+ %0 = "torch_test.operator"(%arg0, %arg1, %arg2) : (index, i32, f32) -> i32
+ return %0 : i32
+ }
+}
diff --git a/compiler/plugins/input/Torch/InputConversion/test/func_conversion_sync_invalid.mlir b/compiler/plugins/input/Torch/InputConversion/test/func_conversion_sync_invalid.mlir
new file mode 100644
index 0000000..c6ada1e
--- /dev/null
+++ b/compiler/plugins/input/Torch/InputConversion/test/func_conversion_sync_invalid.mlir
@@ -0,0 +1,24 @@
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(torch-iree-sync-func-conversion)" --verify-diagnostics %s
+
+// Externalized transients require the coarse-fences ABI; the flag combination
+// is rejected at plugin activation before any compilation runs.
+// RUN: not iree-compile --iree-torch-emit-async-entry-points=false --iree-torch-externalize-transients %s 2>&1 | FileCheck --check-prefix=CHECK-CONFLICT %s
+// CHECK-CONFLICT: iree-torch-externalize-transients requires async entry points
+
+// Mutable tensors require the coarse-fences ABI for in-place aliasing and are
+// not supported by the sync-only conversion.
+builtin.module @mutable_arg {
+// expected-error @+1 {{mutable tensors are not supported}}
+func.func @main(%arg0: !torch.tensor<[5,4],f32>) {
+ return
+}
+}
+
+// -----
+builtin.module @mutable_result {
+// expected-error @+1 {{mutable tensors are not supported}}
+func.func @main() -> !torch.tensor<[5,4],f32> {
+ %0 = torch.operator "some.mutable_producer"() : () -> !torch.tensor<[5,4],f32>
+ return %0 : !torch.tensor<[5,4],f32>
+}
+}
diff --git a/compiler/plugins/input/Torch/PluginRegistration.cpp b/compiler/plugins/input/Torch/PluginRegistration.cpp
index 2f934be..4dff2dd 100644
--- a/compiler/plugins/input/Torch/PluginRegistration.cpp
+++ b/compiler/plugins/input/Torch/PluginRegistration.cpp
@@ -31,6 +31,7 @@
bool decompose = true;
bool externalizeTransients = false;
bool enableShapeRefinement = false;
+ bool emitAsyncEntryPoints = true;
void bindOptions(OptionsBinder &binder) {
static llvm::cl::OptionCategory category("Torch Input");
binder.opt<bool>(
@@ -50,6 +51,13 @@
binder.opt<bool>("iree-torch-enable-shape-refinement",
enableShapeRefinement, llvm::cl::cat(category),
llvm::cl::desc("Enable shape refinement"));
+ binder.opt<bool>(
+ "iree-torch-emit-async-entry-points", emitAsyncEntryPoints,
+ llvm::cl::cat(category),
+ llvm::cl::desc(
+ "Generate async functions with coarse-fences ABI and sync wrapper. "
+ "When false, generates only lightweight sync functions without "
+ "mutable tensor support."));
}
};
@@ -65,6 +73,15 @@
TorchInput::registerTMTensorConversionPasses();
}
+ LogicalResult onActivate() override {
+ if (options.externalizeTransients && !options.emitAsyncEntryPoints) {
+ return emitError(UnknownLoc::get(context))
+ << "iree-torch-externalize-transients requires async entry "
+ "points";
+ }
+ return success();
+ }
+
void onRegisterDialects(DialectRegistry ®istry) override {
registry.insert<torch::Torch::TorchDialect>();
registry.insert<torch::TorchConversion::TorchConversionDialect>();
@@ -100,6 +117,7 @@
torchOptions.decompose = options.decompose;
torchOptions.externalizeTransients = options.externalizeTransients;
torchOptions.enableShapeRefinement = options.enableShapeRefinement;
+ torchOptions.emitAsyncEntryPoints = options.emitAsyncEntryPoints;
TorchInput::createTorchToIREEPipeline(passManager, torchOptions);
return true;
}
