[Mlir-commits] [mlir] [mlir][spirv] Lower mapped TOSA custom ops to ExperimentalML.Call (PR #202581)
Davide Grohmann
llvmlistbot at llvm.org
Tue Jun 9 04:08:04 PDT 2026
https://github.com/davidegrohmann created https://github.com/llvm/llvm-project/pull/202581
Extend the TOSA to SPIR-V Graph conversion with an optional custom-op-domain-to-opcode mapping. TOSA custom ops whose domain appears in this mapping are lowered to spirv.ExperimentalML.Call using the mapped CALL opcode. Later mappings for the same domain override earlier ones, matching command-line option precedence.
For this TOSA lowering, CALL operands start with a spirv.array<N x i8> carrying the operator_name byte blob. This is followed by another i8 array for implementation_attrs, then the original tensor inputs. Empty strings are encoded as a single NUL byte because SPIR-V array types require at least one element.
Use existing SPIR-V array constants for the metadata operands so the target test stays on the SPIR-V binary round-trip path.
>From a3844b0838d4fb30fa963f30b09aa38f9f2fbdfe Mon Sep 17 00:00:00 2001
From: Niklas Lithammer <niklas.lithammer at arm.com>
Date: Tue, 28 Apr 2026 11:08:14 +0200
Subject: [PATCH] [mlir][spirv] Lower mapped TOSA custom ops to
ExperimentalML.Call
Extend the TOSA to SPIR-V Graph conversion with an optional
custom-op-domain-to-opcode mapping. TOSA custom ops whose domain appears in
this mapping are lowered to spirv.ExperimentalML.Call using the mapped CALL
opcode. Later mappings for the same domain override earlier ones, matching
command-line option precedence.
For this TOSA lowering, CALL operands start with a spirv.array<N x i8>
carrying the operator_name byte blob. This is followed by another i8 array for
implementation_attrs, then the original tensor inputs. Empty strings are
encoded as a single NUL byte because SPIR-V array types require at least one
element.
Use existing SPIR-V array constants for the metadata operands so the target
test stays on the SPIR-V binary round-trip path.
Signed-off-by: Niklas Lithammer <niklas.lithammer at arm.com>
Signed-off-by: Davide Grohmann <davide.grohmann at arm.com>
Change-Id: I80fd566f455a9684ad13998c3a061f1e0da5bf1f
---
mlir/include/mlir/Conversion/Passes.td | 7 ++
.../TosaToSPIRVTosa/TosaToSPIRVTosa.h | 6 +
.../Conversion/TosaToSPIRVTosa/CMakeLists.txt | 1 +
.../TosaToSPIRVTosa/TosaToSPIRVTosaCustom.cpp | 107 ++++++++++++++++++
.../TosaToSPIRVTosa/TosaToSPIRVTosaPass.cpp | 46 ++++++++
.../custom-op-domain-to-opcode-invalid.mlir | 7 ++
.../custom-op-domain-to-opcode.mlir | 42 +++++++
7 files changed, 216 insertions(+)
create mode 100644 mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaCustom.cpp
create mode 100644 mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode-invalid.mlir
create mode 100644 mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode.mlir
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index c30dd3b07d028..c5eebb1fee978 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1437,6 +1437,13 @@ def TosaToSPIRVTosa : Pass<"tosa-to-spirv-tosa"> {
lowering supported TOSA ops to `spirv.Tosa.*`, and rewriting TOSA tensor
and shape types to the corresponding SPIR-V ARM tensor types.
}];
+ let options = [
+ ListOption<"customOpDomainToOpcode", "custom-op-domain-to-opcode",
+ "std::pair<std::string, int32_t>",
+ "Map TOSA custom op domains to spirv.ExperimentalML.Call "
+ "opcodes. Entries use <domain>:<opcode> and can be specified "
+ "multiple times.">
+ ];
let constructor = "tosa::createTosaToSPIRVTosa()";
}
diff --git a/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h b/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h
index d03f09f9be78c..9a42cf75f187b 100644
--- a/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h
+++ b/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h
@@ -16,8 +16,11 @@
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
+#include <cstdint>
+
namespace mlir {
#define GEN_PASS_DECL_TOSATOSPIRVTOSAMARKGRAPHCONSTANTS
@@ -55,6 +58,9 @@ void populateTosaToSPIRVTosaConversionPatterns(
spirv::TargetEnvAttr targetAttr);
void populateTosaToSPIRVTosaOpsConversionPatterns(
SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns);
+void populateTosaToSPIRVTosaCustomConversionPatterns(
+ SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns,
+ llvm::StringMap<int32_t> domainToOpcode);
} // namespace tosa
} // namespace mlir
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt b/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt
index ea59ba5f24cb1..c4b71ed61f39b 100644
--- a/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_conversion_library(MLIRTosaToSPIRVTosa
TosaToSPIRVTosa.cpp
TosaToSPIRVTosaConstants.cpp
TosaToSPIRVTosaOps.cpp
+ TosaToSPIRVTosaCustom.cpp
TosaToSPIRVTosaPass.cpp
ADDITIONAL_HEADER_DIRS
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaCustom.cpp b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaCustom.cpp
new file mode 100644
index 0000000000000..db97b97d57714
--- /dev/null
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaCustom.cpp
@@ -0,0 +1,107 @@
+//===- TosaToSPIRVTosaCustom.cpp - TOSA to SPIR-V Graph/TOSA patterns -----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements patterns to convert TOSA IR to SPIR-V Graph/TOSA.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
+
+#define DEBUG_TYPE "tosa-to-spirv-tosa-custom-pattern"
+
+namespace mlir::tosa {
+namespace {
+
+Value getI8ArrayConstant(StringRef value, Location loc,
+ ConversionPatternRewriter &rewriter) {
+ auto i8Type = rewriter.getIntegerType(8);
+ // Empty strings are encoded as a single NUL byte because SPIR-V array
+ // types require at least one element.
+ StringRef encodedValue = value.empty() ? StringRef("\0", 1) : value;
+
+ SmallVector<Attribute> bytes;
+ bytes.reserve(encodedValue.size());
+ llvm::transform(
+ encodedValue, std::back_inserter(bytes),
+ [&](unsigned char byte) { return IntegerAttr::get(i8Type, byte); });
+
+ auto arrayType =
+ spirv::ArrayType::get(i8Type, static_cast<unsigned>(bytes.size()));
+ auto arrayValue = ArrayAttr::get(rewriter.getContext(), bytes);
+ return spirv::ConstantOp::create(rewriter, loc, arrayType, arrayValue);
+}
+
+struct TosaCustomOpConvert final : public OpConversionPattern<tosa::CustomOp> {
+ TosaCustomOpConvert(const TypeConverter &typeConverter, MLIRContext *context,
+ llvm::StringMap<int32_t> domainToOpcode)
+ : OpConversionPattern<tosa::CustomOp>(typeConverter, context),
+ domainToOpcode(std::move(domainToOpcode)) {}
+
+ LogicalResult
+ matchAndRewrite(tosa::CustomOp op, tosa::CustomOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto opCode = domainToOpcode.find(op.getDomainName());
+ if (opCode == domainToOpcode.end())
+ return failure();
+
+ if (op->getResultTypes().empty())
+ return op.emitOpError("with mapped domain requires at least one result");
+
+ SmallVector<Type> types;
+ if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(),
+ types)))
+ return rewriter.notifyMatchFailure(op, "type conversion failed");
+
+ Type resultType =
+ types.size() == 1 ? types.front() : spirv::StructType::get(types);
+
+ Value operatorName =
+ getI8ArrayConstant(op.getOperatorName(), op.getLoc(), rewriter);
+ Value implementationAttrsBlob =
+ getI8ArrayConstant(op.getImplementationAttrs(), op.getLoc(), rewriter);
+
+ SmallVector<Value> inputs = {operatorName, implementationAttrsBlob};
+ inputs.append(adaptor.getInputList().begin(), adaptor.getInputList().end());
+
+ Value result = spirv::ExperimentalMLCallOp::create(
+ rewriter, op.getLoc(), resultType,
+ rewriter.getI32IntegerAttr(opCode->second), inputs);
+
+ if (types.size() == 1) {
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+
+ SmallVector<Value> results;
+ for (auto index : llvm::seq<int32_t>(0, types.size())) {
+ results.push_back(spirv::CompositeExtractOp::create(rewriter, op.getLoc(),
+ result, {index}));
+ }
+ rewriter.replaceOp(op, results);
+ return success();
+ }
+
+private:
+ llvm::StringMap<int32_t> domainToOpcode;
+};
+
+} // namespace
+
+void populateTosaToSPIRVTosaCustomConversionPatterns(
+ SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns,
+ llvm::StringMap<int32_t> domainToOpcode) {
+ patterns.add<TosaCustomOpConvert>(typeConverter, patterns.getContext(),
+ std::move(domainToOpcode));
+}
+
+} // namespace mlir::tosa
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaPass.cpp b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaPass.cpp
index 3224cfe5f11c1..815f2db74ffa6 100644
--- a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaPass.cpp
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaPass.cpp
@@ -19,8 +19,44 @@
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/CommandLine.h"
#include <algorithm>
+#include <string>
+#include <utility>
+
+namespace llvm::cl {
+template <>
+class parser<std::pair<std::string, int32_t>>
+ : public basic_parser<std::pair<std::string, int32_t>> {
+public:
+ parser(Option &option) : basic_parser(option) {}
+
+ bool parse(Option &option, StringRef argName, StringRef arg,
+ std::pair<std::string, int32_t> &value) {
+ auto [domain, opcodeString] = arg.rsplit(":");
+ if (domain.empty() || opcodeString.empty())
+ return option.error("expected <domain>:<opcode>", argName);
+
+ int32_t opcode;
+ if (opcodeString.getAsInteger(0, opcode))
+ return option.error("invalid opcode in custom op domain mapping",
+ argName);
+
+ value = {domain.str(), opcode};
+ return false;
+ }
+
+ StringRef getValueName() const override { return "domain:opcode"; }
+
+ static void print(raw_ostream &os,
+ const std::pair<std::string, int32_t> &value) {
+ os << value.first << ":" << value.second;
+ }
+};
+} // namespace llvm::cl
namespace mlir {
#define GEN_PASS_DEF_TOSATOSPIRVTOSA
@@ -130,6 +166,12 @@ struct TosaToSPIRVTosa final : impl::TosaToSPIRVTosaBase<TosaToSPIRVTosa> {
MLIRContext *context = &getContext();
RewritePatternSet patterns(context);
Operation *op = getOperation();
+ llvm::StringMap<int32_t> domainToOpcode;
+ for (const auto &[domain, opcode] : customOpDomainToOpcode) {
+ // Allow later entries to override earlier ones, matching command-line
+ // option precedence when the same key is specified multiple times.
+ domainToOpcode[domain] = opcode;
+ }
spirv::TargetEnvAttr targetAttr = spirv::lookupTargetEnv(op);
if (!targetAttr) {
@@ -164,6 +206,10 @@ struct TosaToSPIRVTosa final : impl::TosaToSPIRVTosaBase<TosaToSPIRVTosa> {
targetAttr);
populateTosaToSPIRVTosaOpsConversionPatterns(typeConverter, patterns);
+ if (!domainToOpcode.empty())
+ populateTosaToSPIRVTosaCustomConversionPatterns(
+ typeConverter, patterns, std::move(domainToOpcode));
+
FrozenRewritePatternSet frozenPatterns(std::move(patterns));
if (failed(applyPartialConversion(op, *target, frozenPatterns))) {
diff --git a/mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode-invalid.mlir b/mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode-invalid.mlir
new file mode 100644
index 0000000000000..e300603d4844a
--- /dev/null
+++ b/mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode-invalid.mlir
@@ -0,0 +1,7 @@
+// RUN: not mlir-opt --split-input-file --pass-pipeline='builtin.module(tosa-to-spirv-tosa{custom-op-domain-to-opcode=test:7})' %s 2>&1 | FileCheck %s
+
+func.func @zero_results(%arg0: tensor<1x16xf32>) {
+ // CHECK: 'tosa.custom' op with mapped domain requires at least one result
+ tosa.custom %arg0 {domain_name = "test", implementation_attrs = "{}", operator_name = "NoResult"} : (tensor<1x16xf32>) -> ()
+ return
+}
diff --git a/mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode.mlir b/mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode.mlir
new file mode 100644
index 0000000000000..b20b866d7eb58
--- /dev/null
+++ b/mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode.mlir
@@ -0,0 +1,42 @@
+// RUN: mlir-opt --split-input-file --pass-pipeline='builtin.module(tosa-to-spirv-tosa{custom-op-domain-to-opcode=my:custom:0,com.example.accel:42})' %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// Mapped tosa.custom
+//===----------------------------------------------------------------------===//
+
+// CHECK: spirv.module @_spirv_tosa_mapped_custom Logical Vulkan
+// CHECK: spirv.ARM.Graph @mapped_custom(%arg0: !spirv.arm.tensor<1x16xf32> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}, %arg1: !spirv.arm.tensor<1x16xf32> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}) -> (!spirv.arm.tensor<1x16xf32> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 2)>}) attributes {entry_point = true} {
+func.func @mapped_custom(%arg0: tensor<1x16xf32>, %arg1: tensor<1x16xf32>) -> tensor<1x16xf32> {
+ // CHECK: %[[OP_NAME:.*]] = spirv.Constant [84 : i8, 101 : i8, 115 : i8, 116 : i8, 79 : i8, 112 : i8] : !spirv.array<6 x i8>
+ // CHECK: %[[IMPLEMENTATION_ATTRS:.*]] = spirv.Constant [123 : i8, 34 : i8, 112 : i8, 97 : i8, 114 : i8, 97 : i8, 109 : i8, 34 : i8, 58 : i8, 34 : i8, 118 : i8, 97 : i8, 108 : i8, 117 : i8, 101 : i8, 34 : i8, 125 : i8] : !spirv.array<17 x i8>
+ // CHECK: %[[CALL:.*]] = spirv.ExperimentalML.Call opcode = 0, %[[OP_NAME]], %[[IMPLEMENTATION_ATTRS]], %arg0, %arg1 : (!spirv.array<6 x i8>, !spirv.array<17 x i8>, !spirv.arm.tensor<1x16xf32>, !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ %0 = tosa.custom %arg0, %arg1 {domain_name = "my:custom", implementation_attrs = "{\"param\":\"value\"}", operator_name = "TestOp"} : (tensor<1x16xf32>, tensor<1x16xf32>) -> tensor<1x16xf32>
+ // CHECK: spirv.ARM.GraphOutputs %[[CALL]] : !spirv.arm.tensor<1x16xf32>
+ return %0 : tensor<1x16xf32>
+}
+
+// -----
+
+// CHECK: spirv.module @_spirv_tosa_other_custom Logical Vulkan
+// CHECK: spirv.ARM.Graph @other_custom(%arg0: !spirv.arm.tensor<1x16xf32> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}) -> (!spirv.arm.tensor<1x16xf32> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}) attributes {entry_point = true} {
+func.func @other_custom(%arg0: tensor<1x16xf32>) -> tensor<1x16xf32> {
+ // CHECK: %[[OP_NAME:.*]] = spirv.Constant [69 : i8, 120 : i8, 97 : i8, 109 : i8, 112 : i8, 108 : i8, 101 : i8, 79 : i8, 112 : i8] : !spirv.array<9 x i8>
+ // CHECK: %[[IMPLEMENTATION_ATTRS:.*]] = spirv.Constant [123 : i8, 125 : i8] : !spirv.array<2 x i8>
+ // CHECK: %[[CALL:.*]] = spirv.ExperimentalML.Call opcode = 42, %[[OP_NAME]], %[[IMPLEMENTATION_ATTRS]], %arg0 : (!spirv.array<9 x i8>, !spirv.array<2 x i8>, !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ %0 = tosa.custom %arg0 {domain_name = "com.example.accel", implementation_attrs = "{}", operator_name = "ExampleOp"} : (tensor<1x16xf32>) -> tensor<1x16xf32>
+ // CHECK: spirv.ARM.GraphOutputs %[[CALL]] : !spirv.arm.tensor<1x16xf32>
+ return %0 : tensor<1x16xf32>
+}
+
+// -----
+
+// CHECK: spirv.module @_spirv_tosa_empty_strings Logical Vulkan
+// CHECK: spirv.ARM.Graph @empty_strings(%arg0: !spirv.arm.tensor<1x16xf32> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}) -> (!spirv.arm.tensor<1x16xf32> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}) attributes {entry_point = true} {
+func.func @empty_strings(%arg0: tensor<1x16xf32>) -> tensor<1x16xf32> {
+ // CHECK: %[[OP_NAME:.*]] = spirv.Constant [0 : i8] : !spirv.array<1 x i8>
+ // CHECK: %[[IMPLEMENTATION_ATTRS:.*]] = spirv.Constant [0 : i8] : !spirv.array<1 x i8>
+ // CHECK: %[[CALL:.*]] = spirv.ExperimentalML.Call opcode = 0, %[[OP_NAME]], %[[IMPLEMENTATION_ATTRS]], %arg0 : (!spirv.array<1 x i8>, !spirv.array<1 x i8>, !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ %0 = tosa.custom %arg0 {domain_name = "my:custom", implementation_attrs = "", operator_name = ""} : (tensor<1x16xf32>) -> tensor<1x16xf32>
+ // CHECK: spirv.ARM.GraphOutputs %[[CALL]] : !spirv.arm.tensor<1x16xf32>
+ return %0 : tensor<1x16xf32>
+}
More information about the Mlir-commits
mailing list