[Mlir-commits] [mlir] [mlir][tosa][spirv] Add TOSA to SPIR-V TOSA pass plumbing (PR #196539)

[Igor Wodiany]

================
@@ -0,0 +1,191 @@
+//===- TosaToSPIRVTosa.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/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/STLExtras.h"
+
+#define DEBUG_TYPE "tosa-to-spirv-tosa-pattern"
+
+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()) {
+    StringRef attrName = attr.getName().getValue();
+    if (llvm::is_contained({SymbolTable::getSymbolAttrName(),
+                            funcOp.getFunctionTypeAttrName().getValue(),
+                            funcOp.getArgAttrsAttrName().getValue(),
+                            funcOp.getResAttrsAttrName().getValue(),
+                            graphOp.getEntryPointAttrName().getValue()},
+                           attrName))
+      continue;
+
+    graphOp->setAttr(attr.getName(), attr.getValue());
+  }
+}
+
+struct FuncGraphConvert final : OpConversionPattern<func::FuncOp> {
+  FuncGraphConvert(SPIRVTypeConverter &typeConverter, MLIRContext *context,
+                   spirv::TargetEnvAttr targetAttr)
+      : OpConversionPattern<func::FuncOp>(typeConverter, context),
+        targetAttr(targetAttr) {}
+
+private:
+  spirv::TargetEnvAttr targetAttr;
+
+  // Prefer an explicit source-side GraphARM ABI annotation, then preserve an
+  // already-canonical SPIR-V ABI annotation, and otherwise synthesize the
+  // default descriptor set and binding id.
+  void normalizeInterfaceVarABIAttr(spirv::GraphARMOp graphOp,
+                                    MLIRContext *context, unsigned index,
+                                    bool isResult,
+                                    uint32_t defaultDescriptorSet,
+                                    uint32_t defaultBinding) const {
+    auto abiInfo =
+        isResult ? graphOp.getResultAttrOfType<spirv::InterfaceVarABIAttr>(
+                       index, graphARMInterfaceVarABIAttrName)
+                 : graphOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
+                       index, graphARMInterfaceVarABIAttrName);
+
+    if (!abiInfo) {
+      abiInfo = isResult
+                    ? graphOp.getResultAttrOfType<spirv::InterfaceVarABIAttr>(
+                          index, spirv::getInterfaceVarABIAttrName())
+                    : graphOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
+                          index, spirv::getInterfaceVarABIAttrName());
+    }
+
+    if (!abiInfo) {
+      abiInfo = spirv::InterfaceVarABIAttr::get(
+          defaultDescriptorSet, defaultBinding, std::nullopt, context);
+    }
+
+    if (isResult) {
+      graphOp.setResultAttr(index, spirv::getInterfaceVarABIAttrName(),
+                            abiInfo);
+      graphOp.removeResultAttr(index, graphARMInterfaceVarABIAttrName);
+    } else {
+      graphOp.setArgAttr(index, spirv::getInterfaceVarABIAttrName(), abiInfo);
+      graphOp.removeArgAttr(index, graphARMInterfaceVarABIAttrName);
+    }
+  }
+
+  void normalizeInterfaceVarABIAttrs(spirv::GraphARMOp graphOp,
+                                     MLIRContext *context, unsigned inputs,
+                                     unsigned outputs) const {
+    constexpr uint32_t defaultDescriptorSet = 0;
+    for (auto argIndex : llvm::seq<unsigned>(0, inputs)) {
+      normalizeInterfaceVarABIAttr(graphOp, context, argIndex, false,
+                                   defaultDescriptorSet, argIndex);
+    }
+    for (auto resIndex : llvm::seq<unsigned>(0, outputs)) {
+      normalizeInterfaceVarABIAttr(graphOp, context, resIndex, true,
+                                   defaultDescriptorSet, resIndex + inputs);
+    }
+  }
+
+public:
+  LogicalResult
+  matchAndRewrite(func::FuncOp funcOp, func::FuncOpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    MLIRContext *context = rewriter.getContext();
+
+    StringRef name = adaptor.getSymName();
+
+    bool entryPoint = !isa<func::FuncOp>(funcOp->getParentOp());
+    if (entryPoint) {
----------------
IgWod wrote:

Yes, if it's not supported, I think rejecting it for now is a better solution.

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