[Mlir-commits] [mlir] [mlir][spirv] Add TOSA graph constant marking (PR #201095)
Davide Grohmann
llvmlistbot at llvm.org
Tue Jun 2 04:13:15 PDT 2026
https://github.com/davidegrohmann created https://github.com/llvm/llvm-project/pull/201095
Add a TOSA to SPIR-V TOSA preprocessing pass that marks large tosa.const and tosa.const_shape operations for lowering to spirv.ARM.GraphConstant.
Keep small constants inline as spirv.Constant, assign graph constant IDs with a grapharm-prefixed marker attribute, and teach the existing constant conversion to use the marker when present.
Expose the grapharm source-side attribute names used for interface ABI annotations and graph constant IDs.
Add tests for marking large constants, leaving small constants unmarked, increasing graph constant IDs across mixed constants, and lowering pre-marked constants to spirv.ARM.GraphConstant.
>From 491f56544b2dcc93bc4f0320118ecbf09edcfe40 Mon Sep 17 00:00:00 2001
From: Davide Grohmann <davide.grohmann at arm.com>
Date: Mon, 1 Jun 2026 13:27:34 +0200
Subject: [PATCH] [mlir][spirv] Add TOSA graph constant marking
Add a TOSA to SPIR-V TOSA preprocessing pass that marks large
tosa.const and tosa.const_shape operations for lowering to
spirv.ARM.GraphConstant.
Keep small constants inline as spirv.Constant, assign graph constant
IDs with a grapharm-prefixed marker attribute, and teach the existing
constant conversion to use the marker when present.
Expose the grapharm source-side attribute names used for interface ABI
annotations and graph constant IDs.
Add tests for marking large constants, leaving small constants
unmarked, increasing graph constant IDs across mixed constants, and
lowering pre-marked constants to spirv.ARM.GraphConstant.
Change-Id: I5e7621c941e05834346db29e9aba4f7cd2441bb4
Signed-off-by: Davide Grohmann <davide.grohmann at arm.com>
---
mlir/include/mlir/Conversion/Passes.td | 16 +++-
.../TosaToSPIRVTosa/TosaToSPIRVTosa.h | 15 ++++
.../Conversion/TosaToSPIRVTosa/CMakeLists.txt | 1 +
.../TosaToSPIRVTosa/TosaToSPIRVTosa.cpp | 7 --
.../TosaToSPIRVTosaConstants.cpp | 76 +++++++++++++++++++
.../TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp | 7 ++
.../TosaToSPIRVTosa/graph-constant-mark.mlir | 48 ++++++++++++
.../TosaToSPIRVTosa/tosa-to-spirv.mlir | 18 +++++
8 files changed, 179 insertions(+), 9 deletions(-)
create mode 100644 mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaConstants.cpp
create mode 100644 mlir/test/Conversion/TosaToSPIRVTosa/graph-constant-mark.mlir
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index dda756ddab152..c30dd3b07d028 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1414,6 +1414,18 @@ def TosaToSCFPass : Pass<"tosa-to-scf"> {
// TOSA to SPIR-V Graph/TOSA
//===----------------------------------------------------------------------===//
+def TosaToSPIRVTosaMarkGraphConstants
+ : Pass<"tosa-to-spirv-tosa-mark-graph-constants", "func::FuncOp"> {
+ let summary = "Mark large TOSA constants as SPIR-V Graph constants";
+ let description = [{
+ Marks large `tosa.const` and `tosa.const_shape` operations so the
+ TOSA-to-SPIR-V Graph/TOSA conversion lowers them to
+ `spirv.ARM.GraphConstant` instead of inlining them as `spirv.Constant`.
+ }];
+
+ let constructor = "tosa::createTosaToSPIRVTosaMarkGraphConstants()";
+}
+
def TosaToSPIRVTosa : Pass<"tosa-to-spirv-tosa"> {
let summary = "Lower TOSA IR to SPIR-V Graph/TOSA operations";
let dependentDialects = [
@@ -1422,8 +1434,8 @@ def TosaToSPIRVTosa : Pass<"tosa-to-spirv-tosa"> {
let description = [{
Converts TOSA programs to the SPIR-V Graph/TOSA representation by
wrapping converted functions in `spirv.module` and `spirv.ARM.Graph`,
- and rewriting TOSA tensor and shape types to the corresponding SPIR-V ARM
- tensor types.
+ lowering supported TOSA ops to `spirv.Tosa.*`, and rewriting TOSA tensor
+ and shape types to the corresponding SPIR-V ARM tensor types.
}];
let constructor = "tosa::createTosaToSPIRVTosa()";
diff --git a/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h b/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h
index bcc1786642076..d03f09f9be78c 100644
--- a/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h
+++ b/mlir/include/mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h
@@ -16,14 +16,29 @@
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/StringRef.h"
namespace mlir {
+#define GEN_PASS_DECL_TOSATOSPIRVTOSAMARKGRAPHCONSTANTS
#define GEN_PASS_DECL_TOSATOSPIRVTOSA
#include "mlir/Conversion/Passes.h.inc"
namespace tosa {
+// Allows users to specify descriptor sets and binding ids on the source
+// function inputs and outputs. Use a source-side GraphARM attribute because
+// `spirv.interface_var_abi` is verified by the SPIR-V dialect before this
+// conversion runs, and result attrs are only accepted on `spirv.ARM.Graph`.
+constexpr llvm::StringLiteral graphARMInterfaceVarABIAttrName =
+ "grapharm.interface_var_abi";
+
+// Marks a `tosa.const` or `tosa.const_shape` as a SPIR-V Graph constant.
+// The conversion pass lowers marked constants to `spirv.ARM.GraphConstant`.
+constexpr llvm::StringLiteral graphARMGraphConstantIdAttrName =
+ "grapharm.graph_constant_id";
+
+std::unique_ptr<Pass> createTosaToSPIRVTosaMarkGraphConstants();
std::unique_ptr<Pass> createTosaToSPIRVTosa();
spirv::VerCapExtAttr getDefaultVerCapExtAttr(MLIRContext *context);
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt b/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt
index 9d00d4668b124..ea59ba5f24cb1 100644
--- a/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/CMakeLists.txt
@@ -1,5 +1,6 @@
add_mlir_conversion_library(MLIRTosaToSPIRVTosa
TosaToSPIRVTosa.cpp
+ TosaToSPIRVTosaConstants.cpp
TosaToSPIRVTosaOps.cpp
TosaToSPIRVTosaPass.cpp
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.cpp b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.cpp
index 92d2479fdf543..3ef57c1e2cd22 100644
--- a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.cpp
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.cpp
@@ -22,13 +22,6 @@
namespace mlir::tosa {
namespace {
-// Allows users to specify descriptor sets and binding ids on the source
-// function inputs and outputs. Use a source-side GraphARM attribute because
-// `spirv.interface_var_abi` is verified by the SPIR-V dialect before this
-// conversion runs, and result attrs are only accepted on `spirv.ARM.Graph`.
-constexpr StringLiteral graphARMInterfaceVarABIAttrName =
- "grapharm.interface_var_abi";
-
void copyFuncAttrsToGraph(func::FuncOp funcOp, func::FuncOpAdaptor adaptor,
spirv::GraphARMOp graphOp) {
for (NamedAttribute attr : adaptor.getAttributes()) {
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaConstants.cpp b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaConstants.cpp
new file mode 100644
index 0000000000000..8abc100042fa7
--- /dev/null
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaConstants.cpp
@@ -0,0 +1,76 @@
+//===- TosaToSPIRVTosaConstants.cpp - TOSA graph constants ---------------===//
+//
+// 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 preprocessing that marks TOSA constants that should be
+// lowered to SPIR-V Graph constants.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include <optional>
+
+namespace mlir {
+#define GEN_PASS_DEF_TOSATOSPIRVTOSAMARKGRAPHCONSTANTS
+#include "mlir/Conversion/Passes.h.inc"
+
+namespace tosa {
+namespace {
+
+constexpr uint32_t maxInlineConstElements = 16;
+constexpr uint32_t maxInlineConstShapeElements = 32;
+
+std::optional<ElementsAttr> getConstantValues(Operation *op) {
+ if (auto constOp = dyn_cast<tosa::ConstOp>(op))
+ return constOp.getValuesAttr();
+ if (auto constShapeOp = dyn_cast<tosa::ConstShapeOp>(op))
+ return constShapeOp.getValuesAttr();
+ return std::nullopt;
+}
+
+bool shouldMarkGraphConstant(Operation *op) {
+ if (op->use_empty())
+ return false;
+
+ std::optional<ElementsAttr> values = getConstantValues(op);
+ if (!values)
+ return false;
+
+ uint32_t maxInlineElements = isa<tosa::ConstOp>(op)
+ ? maxInlineConstElements
+ : maxInlineConstShapeElements;
+ return values->size() > maxInlineElements;
+}
+
+void setGraphConstantId(Operation *op, uint32_t id) {
+ auto i32Type = IntegerType::get(op->getContext(), 32);
+ op->setAttr(graphARMGraphConstantIdAttrName, IntegerAttr::get(i32Type, id));
+}
+
+struct TosaToSPIRVTosaMarkGraphConstants final
+ : impl::TosaToSPIRVTosaMarkGraphConstantsBase<
+ TosaToSPIRVTosaMarkGraphConstants> {
+ void runOnOperation() override {
+ uint32_t nextConstantId = 0;
+ getOperation().walk([&](Operation *op) {
+ if (shouldMarkGraphConstant(op))
+ setGraphConstantId(op, nextConstantId++);
+ });
+ }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createTosaToSPIRVTosaMarkGraphConstants() {
+ return std::make_unique<TosaToSPIRVTosaMarkGraphConstants>();
+}
+
+} // namespace tosa
+} // namespace mlir
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp
index 16ea459d5edbe..5b53806eea3b1 100644
--- a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp
@@ -450,6 +450,13 @@ LogicalResult replaceRescale(tosa::RescaleOp op, tosa::RescaleOpAdaptor adaptor,
template <typename SourceOp>
LogicalResult replaceConstant(SourceOp op, typename SourceOp::Adaptor adaptor,
Type type, ConversionPatternRewriter &rewriter) {
+ if (auto graphConstantId = op->template getAttrOfType<IntegerAttr>(
+ graphARMGraphConstantIdAttrName)) {
+ rewriter.replaceOpWithNewOp<spirv::GraphConstantARMOp>(
+ op, type, rewriter.getI32IntegerAttr(graphConstantId.getInt()));
+ return success();
+ }
+
auto convertedType = dyn_cast<ShapedType>(type);
auto values = dyn_cast<DenseElementsAttr>(adaptor.getValues());
if (!convertedType || !values)
diff --git a/mlir/test/Conversion/TosaToSPIRVTosa/graph-constant-mark.mlir b/mlir/test/Conversion/TosaToSPIRVTosa/graph-constant-mark.mlir
new file mode 100644
index 0000000000000..98baaeccd0c7b
--- /dev/null
+++ b/mlir/test/Conversion/TosaToSPIRVTosa/graph-constant-mark.mlir
@@ -0,0 +1,48 @@
+// RUN: mlir-opt --split-input-file --tosa-to-spirv-tosa-mark-graph-constants %s | FileCheck %s
+
+// CHECK-LABEL: func.func @large_const
+// CHECK: grapharm.graph_constant_id = 0 : i32
+func.func @large_const() -> tensor<17xi32> {
+ %res = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]> : tensor<17xi32>}> : () -> tensor<17xi32>
+ return %res : tensor<17xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @small_const
+// CHECK-NOT: grapharm.graph_constant_id
+func.func @small_const() -> tensor<16xi32> {
+ %res = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi32>}> : () -> tensor<16xi32>
+ return %res : tensor<16xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @large_const_shape
+// CHECK: grapharm.graph_constant_id = 0 : i32
+func.func @large_const_shape() -> !tosa.shape<33> {
+ %res = "tosa.const_shape"() <{values = dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]> : tensor<33xindex>}> : () -> !tosa.shape<33>
+ return %res : !tosa.shape<33>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @small_const_shape
+// CHECK-NOT: grapharm.graph_constant_id
+func.func @small_const_shape() -> !tosa.shape<32> {
+ %res = "tosa.const_shape"() <{values = dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32]> : tensor<32xindex>}> : () -> !tosa.shape<32>
+ return %res : !tosa.shape<32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @mixed_large_constants
+// CHECK: grapharm.graph_constant_id = 0 : i32
+// CHECK: grapharm.graph_constant_id = 1 : i32
+// CHECK: grapharm.graph_constant_id = 2 : i32
+func.func @mixed_large_constants() -> (tensor<17xi32>, tensor<18xi32>, !tosa.shape<33>) {
+ %const0 = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]> : tensor<17xi32>}> : () -> tensor<17xi32>
+ %const1 = "tosa.const"() <{values = dense<[20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37]> : tensor<18xi32>}> : () -> tensor<18xi32>
+ %shape = "tosa.const_shape"() <{values = dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]> : tensor<33xindex>}> : () -> !tosa.shape<33>
+ return %const0, %const1, %shape : tensor<17xi32>, tensor<18xi32>, !tosa.shape<33>
+}
diff --git a/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir b/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir
index 84edc10e8832e..a0792bf15b724 100644
--- a/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir
+++ b/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir
@@ -963,6 +963,15 @@ func.func @const_int() -> tensor<2x3xi8> {
// -----
+// CHECK-LABEL: spirv.ARM.Graph @graph_constant
+func.func @graph_constant() -> tensor<17xi32> {
+ // CHECK: %[[CONST:.*]] = spirv.ARM.GraphConstant {graph_constant_id = 7 : i32} : !spirv.arm.tensor<17xi32>
+ %res = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]> : tensor<17xi32>}> {grapharm.graph_constant_id = 7 : i32} : () -> tensor<17xi32>
+ return %res : tensor<17xi32>
+}
+
+// -----
+
// CHECK-LABEL: spirv.ARM.Graph @const_i4
func.func @const_i4() -> tensor<2xi4> {
// CHECK: %[[CONST:.*]] = spirv.Constant dense<[7, -8]> : !spirv.arm.tensor<2xi8>
@@ -990,6 +999,15 @@ func.func @const_shape_empty() -> !tosa.shape<0> {
// -----
+// CHECK-LABEL: spirv.ARM.Graph @graph_constant_shape
+func.func @graph_constant_shape() -> !tosa.shape<33> {
+ // CHECK: %[[SHAPE:.*]] = spirv.ARM.GraphConstant {graph_constant_id = 8 : i32} : !spirv.arm.tensor<33xi32>
+ %res = "tosa.const_shape"() <{values = dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]> : tensor<33xindex>}> {grapharm.graph_constant_id = 8 : i32} : () -> !tosa.shape<33>
+ return %res : !tosa.shape<33>
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// spirv.TOSA.Identity
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list