[TOSA] Add TosaToMLProgram conversion (llvm#69787)

This patch adds a new pass to lower TOSA StatefulOps to corresponding ML
Program ops (https://mlir.llvm.org/docs/Dialects/MLProgramOps/).

Signed-off-by: Jerry Ge <[email protected]>

Author: Jerry-Ge

mlir/include/mlir/Conversion/Passes.h

@@ -60,6 +60,7 @@
 #include "mlir/Conversion/TensorToSPIRV/TensorToSPIRVPass.h"
 #include "mlir/Conversion/TosaToArith/TosaToArith.h"
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
+#include "mlir/Conversion/TosaToMLProgram/TosaToMLProgram.h"
 #include "mlir/Conversion/TosaToSCF/TosaToSCF.h"
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
 #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"

mlir/include/mlir/Conversion/Passes.td

@@ -1093,6 +1093,19 @@ def TosaToLinalgNamed
   let constructor = "tosa::createTosaToLinalgNamed()";
 }
 
+//===----------------------------------------------------------------------===//
+// TosaToMLProgram
+//===----------------------------------------------------------------------===//
+
+def TosaToMLProgram : Pass<"tosa-to-mlprogram", "ModuleOp"> {
+  let summary = "Lower TOSA to the MLProgram dialect";
+  let dependentDialects = ["ml_program::MLProgramDialect"];
+  let description = [{
+    Pass that converts TOSA's variable operator operations to the equivalent
+    MLProgram operations.
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // TosaToSCF
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Conversion/TosaToMLProgram/TosaToMLProgram.h

@@ -0,0 +1,30 @@
+//===-- TosaToMLProgram.h - TOSA to MLProgram dialect lowerings-*- C++ -*-===//
+//
+// 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 declares the passes for the TOSA to MLProgram Dialect conversion.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_TOSATOMLPROGRAM_TOSATOMLPROGRAM_H
+#define MLIR_CONVERSION_TOSATOMLPROGRAM_TOSATOMLPROGRAM_H
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+
+#define GEN_PASS_DECL_TOSATOMLPROGRAM
+
+namespace tosa {
+
+void populateTosaToMLProgramConversionPatterns(RewritePatternSet *patterns);
+
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_TOSATOMLPROGRAM_TOSATOMLPROGRAM_H

mlir/lib/Conversion/CMakeLists.txt

@@ -50,6 +50,7 @@ add_subdirectory(TensorToLinalg)
 add_subdirectory(TensorToSPIRV)
 add_subdirectory(TosaToArith)
 add_subdirectory(TosaToLinalg)
+add_subdirectory(TosaToMLProgram)
 add_subdirectory(TosaToSCF)
 add_subdirectory(TosaToTensor)
 add_subdirectory(UBToLLVM)

mlir/lib/Conversion/TosaToMLProgram/CMakeLists.txt

@@ -0,0 +1,19 @@
+add_mlir_conversion_library(MLIRTosaToMLProgram
+  TosaToMLProgram.cpp
+  TosaToMLProgramPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRMLProgramDialect
+  MLIRPass
+  MLIRTosaDialect
+  MLIRTosaTransforms
+  MLIRSupport
+  )

mlir/lib/Conversion/TosaToMLProgram/TosaToMLProgram.cpp

@@ -0,0 +1,76 @@
+//===- TosaToMLProgram.cpp - Lowering Tosa to MLProgram Dialect------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// These rewriters lower from the TOSA dialect to the MLProgram dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToMLProgram/TosaToMLProgram.h"
+#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+using namespace tosa;
+namespace {
+
+class VariableOpConverter : public OpRewritePattern<tosa::VariableOp> {
+public:
+  using OpRewritePattern<tosa::VariableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::VariableOp op,
+                                PatternRewriter &rewriter) const final {
+    auto newVariable = rewriter.create<mlir::ml_program::GlobalOp>(
+        op.getLoc(), op.getName(), op.getType(), /*is_mutable=*/true,
+        op.getInitialValueAttr(), /*sym_visibility=*/nullptr);
+    newVariable.setPrivate();
+    rewriter.replaceOp(op, newVariable);
+    return success();
+  }
+};
+
+class VariableWriteOpConverter
+    : public OpRewritePattern<tosa::VariableWriteOp> {
+public:
+  using OpRewritePattern<tosa::VariableWriteOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::VariableWriteOp op,
+                                PatternRewriter &rewriter) const final {
+    auto globalSymbolRef =
+        SymbolRefAttr::get(rewriter.getContext(), op.getName());
+    auto newVariableWrite = rewriter.create<ml_program::GlobalStoreOp>(
+        op.getLoc(), globalSymbolRef, op.getValue());
+    rewriter.replaceOp(op, newVariableWrite);
+    return success();
+  }
+};
+
+class VariableReadOpConverter : public OpRewritePattern<tosa::VariableReadOp> {
+public:
+  using OpRewritePattern<tosa::VariableReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::VariableReadOp op,
+                                PatternRewriter &rewriter) const final {
+    auto globalSymbolRef =
+        SymbolRefAttr::get(rewriter.getContext(), op.getName());
+    auto newVariableRead = rewriter.create<ml_program::GlobalLoadOp>(
+        op.getLoc(), op.getType(), globalSymbolRef);
+    rewriter.replaceOp(op, newVariableRead);
+
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::tosa::populateTosaToMLProgramConversionPatterns(
+    RewritePatternSet *patterns) {
+  patterns->add<VariableOpConverter, VariableWriteOpConverter,
+                VariableReadOpConverter>(patterns->getContext());
+}

mlir/lib/Conversion/TosaToMLProgram/TosaToMLProgramPass.cpp

@@ -0,0 +1,48 @@
+//===- TosaToMLProgramPass.cpp - Lowering Tosa to MLProgram Dialect--------===//
+//
+// 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 transformation pass legalizes the TOSA dialect to the MLProgram dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToMLProgram/TosaToMLProgram.h"
+#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_TOSATOMLPROGRAM
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+using namespace tosa;
+
+namespace {
+struct TosaToMLProgram : public impl::TosaToMLProgramBase<TosaToMLProgram> {
+public:
+  void runOnOperation() override {
+    auto *context = &getContext();
+    auto moduleOp = getOperation();
+
+    RewritePatternSet patterns(context);
+    ConversionTarget target(*context);
+    target.addIllegalOp<tosa::VariableOp, tosa::VariableReadOp,
+                        tosa::VariableWriteOp>();
+    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+
+    mlir::tosa::populateTosaToMLProgramConversionPatterns(&patterns);
+
+    if (failed(applyPartialConversion(moduleOp, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace

mlir/test/Conversion/TosaToMLProgram/tosa-to-mlprogram.mlir

@@ -0,0 +1,13 @@
+// RUN: mlir-opt --tosa-to-mlprogram %s -o -| FileCheck %s
+
+module {
+  // CHECK: ml_program.global private mutable @var_x(dense<7.000000e+00> : tensor<1xf32>) : tensor<1xf32>
+  tosa.variable @var_x = dense<7.000000e+00> : tensor<1xf32>
+  func.func @test_stateful_ops(%arg0: tensor<1xf32>) -> (tensor<1xf32>) {
+    // CHECK: ml_program.global_store @var_x = %arg0 : tensor<1xf32>
+    tosa.variable.write @var_x, %arg0 : tensor<1xf32>
+    // CHECK: %[[LOAD:.+]] = ml_program.global_load @var_x : tensor<1xf32>
+    %0 = tosa.variable.read @var_x : tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

@@ -3817,6 +3817,7 @@ cc_library(
         ":TensorToSPIRV",
         ":TosaToArith",
         ":TosaToLinalg",
+        ":TosaToMLProgram",
         ":TosaToSCF",
         ":TosaToTensor",
         ":UBToLLVM",
@@ -11212,6 +11213,30 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "TosaToMLProgram",
+    srcs = glob([
+        "lib/Conversion/TosaToMLProgram/*.cpp",
+        "lib/Conversion/TosaToMLProgram/*.h",
+    ]),
+    hdrs = glob([
+        "include/mlir/Conversion/TosaToMLProgram/*.h",
+    ]),
+    includes = [
+        "include",
+        "lib/Conversion/TosaToMLProgram",
+    ],
+    deps = [
+        ":ConversionPassIncGen",
+        ":FuncDialect",
+        ":IR",
+        ":Pass",
+        ":MLProgramDialect",
+        ":TosaDialect",
+        ":Transforms",
+    ],
+)
+
 cc_library(
     name = "TosaToSCF",
     srcs = glob([