[Mlir-commits] [mlir] [mlir][spirv] Add TOSA graph constant marking (PR #201095)

[Davide Grohmann]

================
@@ -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++);
+    });
+  }
----------------
davidegrohmann wrote:

The pass is intended to be used when there are no custom graph constant id marks already in place. So simply make the pass fail if there are already some marks present.

https://github.com/llvm/llvm-project/pull/201095