[Mlir-commits] [mlir] [mlir][spirv] Lower mapped TOSA custom ops to ExperimentalML.Call (PR #201848)
Davide Grohmann
llvmlistbot at llvm.org
Fri Jun 5 07:14:59 PDT 2026
https://github.com/davidegrohmann created https://github.com/llvm/llvm-project/pull/201848
Add the Arm.ExperimentalMLOperations.1 extended instruction set and a spirv.ExperimentalML.Call operation.
Reference:
https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extended/Arm.ExperimentalMLOperations.asciidoc
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 c9b5357fcbce3ee2cc4932a81e3ff5dfcab6041b 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
Add the Arm.ExperimentalMLOperations.1 extended instruction set and a
spirv.ExperimentalML.Call operation.
Reference:
https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extended/Arm.ExperimentalMLOperations.asciidoc
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 +
.../SPIRV/IR/SPIRVExperimentalMLOps.td | 52 +++++++++
.../include/mlir/Dialect/SPIRV/IR/SPIRVOps.td | 1 +
.../Conversion/TosaToSPIRVTosa/CMakeLists.txt | 1 +
.../TosaToSPIRVTosa/TosaToSPIRVTosaCustom.cpp | 107 ++++++++++++++++++
.../TosaToSPIRVTosa/TosaToSPIRVTosaPass.cpp | 47 ++++++++
.../custom-op-domain-to-opcode-invalid.mlir | 7 ++
.../custom-op-domain-to-opcode.mlir | 42 +++++++
.../Dialect/SPIRV/IR/experimental-ml-ops.mlir | 14 +++
.../Target/SPIRV/experimental-ml-ops.mlir | 25 ++++
11 files changed, 309 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/SPIRV/IR/SPIRVExperimentalMLOps.td
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
create mode 100644 mlir/test/Dialect/SPIRV/IR/experimental-ml-ops.mlir
create mode 100644 mlir/test/Target/SPIRV/experimental-ml-ops.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/include/mlir/Dialect/SPIRV/IR/SPIRVExperimentalMLOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVExperimentalMLOps.td
new file mode 100644
index 0000000000000..ddd478881958b
--- /dev/null
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVExperimentalMLOps.td
@@ -0,0 +1,52 @@
+//===- SPIRVExperimentalMLOps.td - Experimental ML ops ------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_IR_EXPERIMENTAL_ML_OPS
+#define MLIR_DIALECT_SPIRV_IR_EXPERIMENTAL_ML_OPS
+
+include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
+
+class SPIRV_ExperimentalMLOp<string mnemonic, int opcode,
+ list<Trait> traits = []> :
+ SPIRV_ExtInstOp<mnemonic, "ExperimentalML", "Arm.ExperimentalMLOperations.1",
+ opcode, traits> {
+ let availability = [
+ MinVersion<SPIRV_V_1_0>,
+ MaxVersion<SPIRV_V_1_6>,
+ Extension<[]>,
+ Capability<[]>
+ ];
+
+ let hasVerifier = 0;
+}
+
+def SPIRV_ExperimentalMLCallOp : SPIRV_ExperimentalMLOp<"Call", 0> {
+ let summary = "Call an Arm experimental ML operation.";
+
+ let description = [{
+ Calls an operation encoded using the Arm.ExperimentalMLOperations extended
+ instruction set. The `opcode` attribute is serialized as the operation
+ opcode literal integer operand of CALL.
+ }];
+
+ let arguments = (ins
+ I32Attr:$opcode,
+ Variadic<AnyType>:$parameters
+ );
+
+ let results = (outs
+ AnyType:$output
+ );
+
+ let assemblyFormat = [{
+ `opcode` `=` $opcode `,`
+ $parameters attr-dict `:` functional-type($parameters, $output)
+ }];
+}
+
+#endif // MLIR_DIALECT_SPIRV_IR_EXPERIMENTAL_ML_OPS
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
index 3ef9699154cd1..14dc76aef57db 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
@@ -30,6 +30,7 @@ include "mlir/Dialect/SPIRV/IR/SPIRVCastOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td"
+include "mlir/Dialect/SPIRV/IR/SPIRVExperimentalMLOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVIntelExtOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVGLOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVGraphOps.td"
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..abc1d5a4e2fc4 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
@@ -126,10 +162,17 @@ LogicalResult verifyGraphConstantIdAttrs(Operation *op) {
}
struct TosaToSPIRVTosa final : impl::TosaToSPIRVTosaBase<TosaToSPIRVTosa> {
+
void runOnOperation() override {
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 +207,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>
+}
diff --git a/mlir/test/Dialect/SPIRV/IR/experimental-ml-ops.mlir b/mlir/test/Dialect/SPIRV/IR/experimental-ml-ops.mlir
new file mode 100644
index 0000000000000..caccb043d3af2
--- /dev/null
+++ b/mlir/test/Dialect/SPIRV/IR/experimental-ml-ops.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spirv.ExperimentalML.Call
+//===----------------------------------------------------------------------===//
+
+spirv.ARM.Graph @experimental_ml_call(%arg0: !spirv.arm.tensor<1x16xf32>, %arg1: !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32> {
+ // CHECK: %[[NAME:.*]] = spirv.Constant dense<[83, 101, 108, 102, 65, 116, 116, 101, 110, 116, 105, 111, 110, 79, 112]> : tensor<15xi8> : !spirv.array<15 x i8>
+ %name = spirv.Constant dense<[83, 101, 108, 102, 65, 116, 116, 101, 110, 116, 105, 111, 110, 79, 112]> : tensor<15xi8> : !spirv.array<15 x i8>
+ // CHECK: {{%.*}} = spirv.ExperimentalML.Call opcode = 0, %[[NAME]], %arg0, %arg1 : (!spirv.array<15 x i8>, !spirv.arm.tensor<1x16xf32>, !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ %0 = spirv.ExperimentalML.Call opcode = 0, %name, %arg0, %arg1 : (!spirv.array<15 x i8>, !spirv.arm.tensor<1x16xf32>, !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ // CHECK: spirv.ARM.GraphOutputs {{%.*}} : !spirv.arm.tensor<1x16xf32>
+ spirv.ARM.GraphOutputs %0 : !spirv.arm.tensor<1x16xf32>
+}
diff --git a/mlir/test/Target/SPIRV/experimental-ml-ops.mlir b/mlir/test/Target/SPIRV/experimental-ml-ops.mlir
new file mode 100644
index 0000000000000..5876f7339c330
--- /dev/null
+++ b/mlir/test/Target/SPIRV/experimental-ml-ops.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip %s | FileCheck %s
+// RUN: %if spirv-tools %{ mlir-translate --no-implicit-module --serialize-spirv %s | spirv-val %}
+
+// CHECK: spirv.module Logical Vulkan requires
+spirv.module Logical Vulkan requires #spirv.vce<v1.3, [VulkanMemoryModel, Shader, Int8, Float16, TensorsARM, GraphARM], [SPV_ARM_tensors, SPV_ARM_graph, SPV_KHR_vulkan_memory_model]> {
+ // CHECK: spirv.GlobalVariable @main_arg_0 bind(0, 0) : !spirv.ptr<!spirv.arm.tensor<1x16xf32>, UniformConstant>
+ spirv.GlobalVariable @main_arg_0 bind(0, 0) : !spirv.ptr<!spirv.arm.tensor<1x16xf32>, UniformConstant>
+ // CHECK: spirv.GlobalVariable @main_arg_1 bind(0, 1) : !spirv.ptr<!spirv.arm.tensor<1x16xf32>, UniformConstant>
+ spirv.GlobalVariable @main_arg_1 bind(0, 1) : !spirv.ptr<!spirv.arm.tensor<1x16xf32>, UniformConstant>
+ // CHECK: spirv.GlobalVariable @main_res_0 bind(0, 2) : !spirv.ptr<!spirv.arm.tensor<1x16xf32>, UniformConstant>
+ spirv.GlobalVariable @main_res_0 bind(0, 2) : !spirv.ptr<!spirv.arm.tensor<1x16xf32>, UniformConstant>
+ // CHECK: spirv.ARM.GraphEntryPoint @main, @main_arg_0, @main_arg_1, @main_res_0
+ spirv.ARM.GraphEntryPoint @main, @main_arg_0, @main_arg_1, @main_res_0
+ // CHECK: spirv.ARM.Graph @main(%arg0: !spirv.arm.tensor<1x16xf32>, %arg1: !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ spirv.ARM.Graph @main(%arg0: !spirv.arm.tensor<1x16xf32>, %arg1: !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32> {
+ // CHECK: %[[NAME:.*]] = spirv.Constant [83 : i8, 101 : i8, 108 : i8, 102 : i8, 65 : i8, 116 : i8, 116 : i8, 101 : i8, 110 : i8, 116 : i8, 105 : i8, 111 : i8, 110 : i8, 79 : i8, 112 : i8] : !spirv.array<15 x i8>
+ %name = spirv.Constant dense<[83, 101, 108, 102, 65, 116, 116, 101, 110, 116, 105, 111, 110, 79, 112]> : tensor<15xi8> : !spirv.array<15 x i8>
+ // CHECK: %[[ATTRS:.*]] = spirv.Constant [123 : i8, 34 : i8, 111 : i8, 112 : i8, 101 : i8, 114 : i8, 97 : i8, 116 : i8, 111 : i8, 114 : i8, 95 : i8, 110 : i8, 97 : i8, 109 : i8, 101 : i8, 34 : i8, 58 : i8, 34 : i8, 83 : i8, 101 : i8, 108 : i8, 102 : i8, 65 : i8, 116 : i8, 116 : i8, 101 : i8, 110 : i8, 116 : i8, 105 : i8, 111 : i8, 110 : i8, 79 : i8, 112 : i8, 34 : i8, 125 : i8] : !spirv.array<35 x i8>
+ %attrs = spirv.Constant dense<[123, 34, 111, 112, 101, 114, 97, 116, 111, 114, 95, 110, 97, 109, 101, 34, 58, 34, 83, 101, 108, 102, 65, 116, 116, 101, 110, 116, 105, 111, 110, 79, 112, 34, 125]> : tensor<35xi8> : !spirv.array<35 x i8>
+ // CHECK: %[[CALL:.*]] = spirv.ExperimentalML.Call opcode = 0, %[[NAME]], %[[ATTRS]], %arg0, %arg1 : (!spirv.array<15 x i8>, !spirv.array<35 x i8>, !spirv.arm.tensor<1x16xf32>, !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ %0 = spirv.ExperimentalML.Call opcode = 0, %name, %attrs, %arg0, %arg1 : (!spirv.array<15 x i8>, !spirv.array<35 x i8>, !spirv.arm.tensor<1x16xf32>, !spirv.arm.tensor<1x16xf32>) -> !spirv.arm.tensor<1x16xf32>
+ // CHECK: spirv.ARM.GraphOutputs %[[CALL]] : !spirv.arm.tensor<1x16xf32>
+ spirv.ARM.GraphOutputs %0 : !spirv.arm.tensor<1x16xf32>
+ }
+}
More information about the Mlir-commits
mailing list