[tosa] Add option to disable tosa.apply_scale lowering in TosaToStandard

[rsuderman]

Summary

Apply scale should be optionally disabled when lowering via TosaToStandard.
In most cases it should persist until the lowering to specific backend.

[jpienaar]

mlir/include/mlir/Conversion/Passes.td

758

s/standard/arith/? Or are there other ones along? (I was mentioning to Mehdi I don't know how to rename the ShapeToStandard pass with Standard exploding :)) [you can also defer if this is multiple dialects being lowered to]

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

41

Could we add a comment somewhere in code (perhaps in pass description so that it makes it into the generated docs) as to the why.

mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir

1

Test with it false? (don't have to repeat all CHECK's, just a spot check)

[rsuderman]

Split TosaToStandard split to TosaToArith and TosaToTensor

[rsuderman]

Added the additional comments in.

[rsuderman]

mlir/include/mlir/Conversion/Passes.td

758

I split everything between TosaToArith and TosaToTensor as each defines a different party of the work.

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

41

Added details to the Passes.td file for while the ApplyScale is optionally included.

[jpienaar]

mlir/include/mlir/Conversion/Passes.td

714

Up to you, but I'd just drop ops and have everything on one line (or drop tosa. as this is a tosa-to-arith pass so the dialect is implicit for me)

mlir/include/mlir/Conversion/TosaToArith/TosaToArith.h

21 ↗

(On Diff #419908)

This and addTosaToArithPasses seem to overlap quite a bit, when would one use one vs the other?

mlir/include/mlir/Conversion/TosaToTensor/TosaToTensor.h

1 ↗

(On Diff #419908)

optimization pass? TOSA to Tensor legalization pass ?

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

1

Tensor

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

1 ↗

(On Diff #419908)

Tensor

9 ↗

(On Diff #419908)

Tensor

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

1 ↗

(On Diff #419908)

verify-diagnostics doesn't seem needed here

[rsuderman]

Updated for cleanup comments.

mlir/include/mlir/Conversion/Passes.td

def TosaToLinalg
 
   let constructor = "tosa::createTosaToLinalg()";
 }
 
 //===----------------------------------------------------------------------===//
 // TosaToLinalgNamed
 //===----------------------------------------------------------------------===//
 
 def TosaToLinalgNamed
     : InterfacePass<"tosa-to-linalg-named", "FunctionOpInterface"> {
   let summary = "Lower TOSA to LinAlg named operations";
   let description = [{
     Pass that converts TOSA operations to the equivalent operations using the
     Linalg named operations.
   }];
 
   let constructor = "tosa::createTosaToLinalgNamed()";
let dependentDialects = [
     "arith::ArithmeticDialect",
     "tensor::TensorDialect",
   ];
   let description = [{
     Pass that converts TOSA operations to the equivalent operations using the
     operations in the Standard dialect.
   }];
 
+  let options = [
+    Option<"includeApplyRescale", "include-apply-rescale",
+           "bool", /*default=*/"false",
+           "Whether to include the lowering for tosa.apply_rescale to standard"
+           " ops">
+  ];
+
   let constructor = "tosa::createTosaToStandard()";
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToGPU
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToGPU : Pass<"convert-vector-to-gpu"> {

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp

 //===- TosaToLinalgPass.cpp - Lowering Tosa to Linalg 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 Tosa operations to the Linalg dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
+#include "mlir/Conversion/TosaToStandard/TosaToStandard.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

LogicalResult matchAndRewrite(tosa::ApplyScaleOp op,
     return success();
   }
 };
 
 } // namespace
 
 void mlir::tosa::populateTosaToStandardConversionPatterns(
     RewritePatternSet *patterns) {
-  patterns->add<ApplyScaleOpConverter, ConstOpConverter, SliceOpConverter>(
-      patterns->getContext());
+  patterns->add<ConstOpConverter, SliceOpConverter>(patterns->getContext());
 }
 
 void mlir::tosa::populateTosaRescaleToStandardConversionPatterns(
     RewritePatternSet *patterns) {
   patterns->add<ApplyScaleOpConverter>(patterns->getContext());
 }

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

 //===- TosaToStandardPass.cpp - Lowering Tosa to Linalg 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 Tosa operations to the Standard dialect.

 namespace {
 struct TosaToStandard : public TosaToStandardBase<TosaToStandard> {
 public:
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     ConversionTarget target(getContext());
     target.addIllegalOp<tosa::ConstOp>();
     target.addIllegalOp<tosa::SliceOp>();
-    target.addIllegalOp<tosa::ApplyScaleOp>();
     target.addLegalDialect<arith::ArithmeticDialect>();
     target.addLegalDialect<tensor::TensorDialect>();
 
     mlir::tosa::populateTosaToStandardConversionPatterns(&patterns);
+
+    if (this->includeApplyRescale) {
+      mlir::tosa::populateTosaRescaleToStandardConversionPatterns(&patterns);
+      target.addIllegalOp<tosa::ApplyScaleOp>();
+    }
+
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::tosa::createTosaToStandard() {
   return std::make_unique<TosaToStandard>();
 }
 
 void mlir::tosa::addTosaToStandardPasses(OpPassManager &pm) {
   pm.addNestedPass<FuncOp>(createTosaToStandard());
 }

mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir

-// RUN: mlir-opt --split-input-file --tosa-to-standard %s -verify-diagnostics -o -| FileCheck %s
+// RUN: mlir-opt --split-input-file --tosa-to-standard="include-apply-rescale=true" %s -verify-diagnostics -o -| FileCheck %s
 
 // CHECK-LABEL: func @const_test
 func @const_test() -> (tensor<i32>) {
   // CHECK: [[C3:%.+]] = arith.constant dense<3> : tensor<i32>
   %0 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
 
   // CHECK: return [[C3]]
   return %0 : tensor<i32>