[Mlir-commits] [mlir] 7adb63d - [mlir][spirv] Lower mapped TOSA custom ops to ExperimentalML.Call (#202581)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 12 00:24:14 PDT 2026
Author: Davide Grohmann
Date: 2026-06-12T09:24:08+02:00
New Revision: 7adb63de37b95dfbbb67ff775ba5e2bc93f15458
URL: https://github.com/llvm/llvm-project/commit/7adb63de37b95dfbbb67ff775ba5e2bc93f15458
DIFF: https://github.com/llvm/llvm-project/commit/7adb63de37b95dfbbb67ff775ba5e2bc93f15458.diff
LOG: [mlir][spirv] Lower mapped TOSA custom ops to ExperimentalML.Call (#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.
Signed-off-by: Niklas Lithammer <niklas.lithammer at arm.com>
Signed-off-by: Davide Grohmann <davide.grohmann at arm.com>
Co-authored-by: Niklas Lithammer <niklas.lithammer at arm.com>
Added:
mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaCustom.cpp
mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode-invalid.mlir
mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode.mlir
Modified:
mlir/include/mlir/Conversion/Passes.td
mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h
mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt
mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaPass.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 07e0e0c4e29e8..4f8a66957d954 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1443,6 +1443,14 @@ 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>. The option can be "
+ "specified multiple times; if a domain appears more than once, "
+ "the last mapping wins.">
+ ];
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..14fa2fa12d2a8
--- /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 encodeStringAsI8Array(StringRef value, Location loc,
+ ConversionPatternRewriter &rewriter) {
+ Type i8Type = rewriter.getIntegerType(8);
+ // Empty strings are encoded as a single NULL 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 =
+ encodeStringAsI8Array(op.getOperatorName(), op.getLoc(), rewriter);
+ Value implementationAttrsBlob = encodeStringAsI8Array(
+ 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..15161f584d275
--- /dev/null
+++ b/mlir/test/Conversion/TosaToSPIRVTosa/custom-op-domain-to-opcode.mlir
@@ -0,0 +1,45 @@
+// 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
+// RUN: mlir-opt --split-input-file --pass-pipeline='builtin.module(tosa-to-spirv-tosa{custom-op-domain-to-opcode=my:custom:99,my:custom:7,com.example.accel:42})' %s | FileCheck %s --check-prefix=OVERRIDE
+
+//===----------------------------------------------------------------------===//
+// 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>
+ // OVERRIDE-LABEL: spirv.ARM.Graph @mapped_custom
+ // OVERRIDE: spirv.ExperimentalML.Call opcode = 7,
+ %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