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

[Jakub Kuderski]

================
@@ -0,0 +1,158 @@
+//===- TosaToSPIRVTosaPass.cpp - Lower TOSA to SPIR-V Graph/TOSA ----------===//
+//
+// 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 pass lowers TOSA IR to the SPIR-V Graph/TOSA representation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include <algorithm>
+
+namespace mlir {
+#define GEN_PASS_DEF_TOSATOSPIRVTOSA
+#include "mlir/Conversion/Passes.h.inc"
+
+namespace tosa {
+namespace {
+
+spirv::VerCapExtAttr getVerCapExtAttr(MLIRContext *context) {
+  return spirv::VerCapExtAttr::get(
+      spirv::Version::V_1_5,
+      {
+          spirv::Capability::VulkanMemoryModel,
+          spirv::Capability::Shader,
+          spirv::Capability::Int8,
+          spirv::Capability::Int16,
+          spirv::Capability::Int64,
+          spirv::Capability::Float16,
+          spirv::Capability::BFloat16TypeKHR,
+          spirv::Capability::Float8EXT,
+          spirv::Capability::TensorsARM,
+          spirv::Capability::GraphARM,
+          spirv::Capability::ReplicatedCompositesEXT,
+      },
+      {
+          spirv::Extension::SPV_ARM_tensors,
+          spirv::Extension::SPV_ARM_graph,
+          spirv::Extension::SPV_KHR_vulkan_memory_model,
+          spirv::Extension::SPV_EXT_replicated_composites,
+          spirv::Extension::SPV_KHR_bfloat16,
+          spirv::Extension::SPV_EXT_float8,
+      },
+      context);
+}
+
+struct TosaToSPIRVTosa final : impl::TosaToSPIRVTosaBase<TosaToSPIRVTosa> {
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(context);
+    Operation *op = getOperation();
+
+    auto targetAttr = spirv::lookupTargetEnv(op);
+    if (!targetAttr) {
+      targetAttr = spirv::TargetEnvAttr::get(
+          getVerCapExtAttr(context), spirv::getDefaultResourceLimits(context),
+          spirv::ClientAPI::Unknown, spirv::Vendor::Unknown,
+          spirv::DeviceType::Unknown, spirv::TargetEnvAttr::kUnknownDeviceID);
+    }
+
+    std::unique_ptr<ConversionTarget> target =
+        SPIRVConversionTarget::get(targetAttr);
+
+    target->addLegalDialect<spirv::SPIRVDialect>();
+    target->addIllegalDialect<tosa::TosaDialect>();
+
+    SPIRVTypeConverter typeConverter(targetAttr);
+    typeConverter.addConversion([this](IntegerType integerType) {
+      return this->convertIntegerType(integerType);
+    });
+    typeConverter.addConversion([this](TensorType tensorType) {
+      return this->convertTensorType(tensorType);
+    });
+    typeConverter.addConversion([this](tosa::shapeType shapeType) {
+      return this->convertShapeType(shapeType);
+    });
+
+    populateTosaToSPIRVTosaConversionPatterns(typeConverter, patterns);
+
+    auto targetEnvAttrName = spirv::getTargetEnvAttrName();
+    auto symbolTableOp = SymbolTable::getNearestSymbolTable(op);
+    if (symbolTableOp && !symbolTableOp->hasAttrOfType<spirv::TargetEnvAttr>(
+                             targetEnvAttrName)) {
+      symbolTableOp->setAttr(targetEnvAttrName, targetAttr);
+    }
+
+    FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+
+    if (failed(applyPartialConversion(op, *target, frozenPatterns))) {
+      signalPassFailure();
+    }
+  }
+
+private:
+  IntegerType convertIntegerType(IntegerType integerType) {
+    if (integerType.getWidth() == 48) {
+      return IntegerType::get(&getContext(), 64, integerType.getSignedness());
+    }
+
+    if (integerType.getWidth() == 4) {
+      return IntegerType::get(&getContext(), 8, integerType.getSignedness());
+    }
+
+    return integerType;
+  }
+
+  SmallVector<int64_t> convertShape(ArrayRef<int64_t> shape) {
+    bool requiresRankConversion =
----------------
kuhar wrote:

Maybe explicitly add tensor of size zero as an illegal type to the type converter?

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