[MLIR][TOSA] First lowerings from Tosa to Linalg

[rsuderman]

Summary

Initial commit to add support for lowering from TOSA to Linalg. The focus is on
the essential infrastructure for these lowerings and integration with existing
passes.

Includes lowerings for a subset of operations including:

abs, add, sub, pow, and, or, xor, left shift, right shift, tanh

Lit tests are used to validate correctness.

[rsuderman]

Some minor fixes before code review.

[mehdi_amini]

Should this be in the lib/Conversion directory with all the other dialect conversions?

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h

30 ↗

(On Diff #315185)

I don't know what a "dispatch function" is, if this is in the context of IREE please remove as it won't help anyone working upstream without this context.

30 ↗

(On Diff #315185)

Is it true? You're using a partial conversion, not a full one which would enforce this.

32 ↗

(On Diff #315185)

Remove reference to IREE from upstream please.

mlir/lib/Dialect/Tosa/Transforms/TosaToLinalgOnTensors.cpp

35 ↗

(On Diff #315185)

Nit: I think function are starting with a lower case.

41 ↗

(On Diff #315185)

The TODO here looks like a copy paste, that seems like something to refactor instead of copying, feel free to split this refactoring in a NFC commit that you can rebase on.

56 ↗

(On Diff #315185)

I'm a bit concerned about the code bloat for the template here.

Using some sort of table / interface may just be more appropriate.

Alternatively, refactoring the matchAndRewrite to make it "trivial" and mostly calling non-templated helper can probably work as well.

65 ↗

(On Diff #315185)

Nit: in general, please use auto when it make the type is obvious and/or the code more clear (everywhere else as well).

78 ↗

(On Diff #315185)

Seem like you could do:

Type uniqueTy = operation->getOperand(0).getType();

bool all_input_types_equal = llvm::all_of(operation->getOperandTypes(), [] (Type operandTy) { return operandTy == uniqueTy; });
if (!all_input_types_equal) 
  return notifyMatchFailure("All operands must have the same type");
bool all_result_types_equal = llvm::all_of(operation->getResultTypes(), [] (Type resultTy) { return resultTy == uniqueTy; });
if (!all_result_types_equal) 
  return notifyMatchFailure("All results must have the same type as the input");

Note that return notifyMatchFailure is useful when debugging failed conversion (it shows up in tracing mode), it can be a good thing to use in place of returning failure().
(and anectodically the code here isn't O(n2) like the nested loops).

81 ↗

(On Diff #315185)

Please name the lambda according to what it is checking, like checkValidShapedType or something like that.

93 ↗

(On Diff #315185)

Please use notifyMatchFailure, conversion in general are supposed to fail gracefully because other patterns could kick in.

102 ↗

(On Diff #315185)

Nit: remove trivial braces.

105 ↗

(On Diff #315185)

124 ↗

(On Diff #315185)

The rewriter.getMultiDimIdentityMap(nloops) does not work with the 0 case?

140 ↗

(On Diff #315185)

Remove this, failed is never modified

198 ↗

(On Diff #315185)

patterns->insert is a variadic template you can do:

patterns->insert<PointwiseConverter<tosa::AbsOp, FloatType, mlir::AbsFOp>,

PointwiseConverter<tosa::AddOp, FloatType, mlir::AddFOp>,
PointwiseConverter<tosa::PowOp, FloatType, mlir::PowFOp>,
...>(context);

mlir/test/Dialect/Tosa/to_linalg_on_tensors.mlir

223 ↗

(On Diff #315185)

Can you add variety of rank and test for the non-conversion of the non-supported cases?

[mehdi_amini]

Looks like I commented on the original diff and you updated in-between, my comments aren't necessarily on the right line, ping me on chat if some of them don't make sense!

(They seem mostly offset by 5 lines down)

[silvas]

mlir/lib/Dialect/Tosa/Transforms/TosaToLinalgOnTensors.cpp

47 ↗

(On Diff #315187)

I would recommend to constrain this pass to only work for static shapes for now.

The code that this pass currently produces will miscompile in the case of dynamic shapes. E.g.

"tosa.add"(%arg0, %arg1) : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>

If the dynamic sizes are tensor<1xf32>, tensor<2xf32> -> tensor<2xf32> then the code emitted by the current pass here will miscompile.

78 ↗

(On Diff #315187)

I belive that tosa ops are always ranked. Also I think the verifier for those ops ensure that the rank of all inputs and outputs match.

mlir/test/Dialect/Tosa/to_linalg_on_tensors.mlir

2 ↗

(On Diff #315187)

Put this pass / test in ConvertTosaToLinalg and name the pass convert-tosa-to-linalg

48 ↗

(On Diff #315187)

Add a test exhibiting that we don't handle broadcasting.

[rsuderman]

Made corresponding changes for the code review comments.

ExtractDynamicDimensions is still TBD in another CL.

[rriddle]

mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h

18

This include doesn't look necessary.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

24

Please make sure that only classes are in namespaces, functions should be at the top level.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp

27

Same here.

[rriddle]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

40

Remove user names from TODO comments.

47–51

70

nit: Only prefer auto if it makes the code more readable. Many of this uses could just use the proper type instead.

101

nit: Can you move this functor out-of-line into a variable before the if? Makes the control flow here pretty hard to read.

150

Can you just replace the use of operation here with op directly?

176

Is the to_vector actually necessary? Does the builder of ScalarOp not take a ValueRange?

[silvas]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

30

use getParallelIteratorTypeName() instead of kParallelIterType

61

Just merge SimplePointwiseConverter into this pattern.

70

please restrict this to the static shape case for now.

95

hasRank and getRank != nloops is redundant

97

do we actually need ComplexType? I don't see any test covering it.

103

move the dyn_cast into checkValidShapedType

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

3

This testing seems really redundant. I think that things like:

// CHECK: absf
%0 = "tosa.abs"(%arg0) : (tensor<f32>) -> tensor<f32>

is sufficient. (with one test that does the more in-depth check of the full linalg op.

7

also check the indexing maps.

And check higher-rank cases as well.

319

test this error in only one place.

[rsuderman]

Handled remaining comments from Mehdi/Sean

[silvas]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

28

these should be static functions not inside an anonymous namespace

45

Does this have to be a conversion pattern? It isn't doing a type conversion, so a regular RewritePattern would work and reduce complexity I think.

50

I don't think there's any value in having this be virtual.

57

can this be a static helper function outside the class?

Also, you might want to comment that it is just this way due to code the amount of templated code.

84

you can use cast here instead of dyn_cast.

93

can either of these checks ever fail? I think all these conditions are guaranteed by the tosa op verifiers. If not, can you add a test case that hits this?

108

hasStaticShape is already guaranteed by the above code.

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

49

check at least the full op name (addf in this case)

also can just put all the tosa.* ops in a single function to reduce boilerplate

[rsuderman]

Change from OpConversionPattern to RewritePattern

[mehdi_amini]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

81

Nit: if t can be null, you can't call cast or dyn_cast on it right above I believe, so I guess this test can be removed?

[sjarus]

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td

29 ↗

(On Diff #315560)

Shouldn't this go in include/mlir/Conversion/Passes.td ?

[rsuderman]

Additional round of comments.

[rsuderman]

Updates to tosa to linalg testing

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

45

Made a RewritePattern.

47–51

Removed entirely.

50

The main advantage is the rewrite helper is identical between instantiations and won't be duplicated decreasing binary size.

61

The intention is for multi op versions to instead implement the virtual function directly. I have some example lowerings in the pipeline to do that.

93

I am attempting to defensively code these practices given how early TOSA support is. Better to verify before lowering than assume valid lowerings.

[rriddle]

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td

9 ↗

(On Diff #315969)

Can changes to this file be reverted now?

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

34

Remove the outer namespaces here. Anonymous namespaces should be top-level.

62

nit: Can you use Type here instead.

167

Prefer using full name resolution for methods. Only classes should really be in namespaces in .cpp files.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp

44

nit: auto -> FuncOp

[rsuderman]

River Riddle comments

[silvas]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

61

I think you'd be better served with a single helper function like this:

https://github.com/llvm/mlir-npcomp/blob/ceab22cf90667f03b3f502f6967392fad8913716/lib/Dialect/TCP/Transforms/Bufferize.cpp#L151

[rsuderman]

Changed to a single function for generating generic op contents.

[silvas]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

100

I think we can hit this with the way we have the code structured now. E.g. an integer add op would currently slip through. Maybe a rewriter.notifyMatchFailure instead?

113

Can you put this in a static helper? I think it will still be duplicated code-size-wise due to the enclosing template.

120

You can remove this check and use cast<ShapedType>. The verifier for tosa ops guarantees it.

144

!st is redundant
!st.hasRank() is also redundant (TOSA ops require it)
!st.getElementType().isSignlessIntOrFloat() is redundant I think (TOSA ops require it)

191

I think you need to handle the case where createLinalgBodyCalculationForElementwiseOp doesn't succeed. (I think just set a bool in the enclosing function and return rewriter.notifyMatchFailure)

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

125

Can you test some cases don't get converted?

• mismatching ranks (rank broadcasting)
• dynamic shapes

[rsuderman]

Removed redundant checks, added more failure test cases.

[rsuderman]

Added missing static qualifier.

[rsuderman]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

113

I moved the whole thing to a static helper. Should guarantee no replication.

[silvas]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

137

nit: remove mention of "buffers" here. initTensors seems like a better name?

149

nit: auto bodyResultTypes = and remove the declaration above.

167

This name "error" seems very confusing. I read it as "did encounter an error", so setting it to true here is confusing. Suggest calling it didEncounterError and setting it to false here.

[mehdi_amini]

LGTM in general.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

23

(FYI: We don't need to specify the SmallVector size anymore.)

32

This kind of helper make me think that they look like a missing build method to be added on the InitTensorOp instead.

45

Remove trivial braces, here and elsewhere.

[rsuderman]

Nit fixes.

[rriddle]

Awesome!

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRVPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/StandardToSPIRVPass.h"
+#include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Conversion/VectorToSPIRV/VectorToSPIRVPass.h"
 
 namespace mlir {
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_PASSES_H

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertStandardToSPIRV : Pass<"convert-std-to-spirv", "ModuleOp"> {
   let summary = "Convert Standard dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertStandardToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 //===----------------------------------------------------------------------===//
+// TosaToLinalg
+//===----------------------------------------------------------------------===//
+
+def TosaToLinalgOnTensors : FunctionPass<"tosa-to-linalg-on-tensors"> {
+  let summary = "Lower TOSA to LinAlg on tensors";
+  let description = [{
+    Pass that converts TOSA operations to the equivalent operations using the
+    tensor operations in LinAlg.
+  }];
+
+  let constructor = "tosa::createTosaToLinalgOnTensors()";
+}
+
+//===----------------------------------------------------------------------===//
 // VectorToSCF
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToSCF : FunctionPass<"convert-vector-to-scf"> {
   let summary = "Lower the operations from the vector dialect into the SCF "
                 "dialect";
   let constructor = "mlir::createConvertVectorToSCFPass()";
   let dependentDialects = ["AffineDialect", "scf::SCFDialect"];

mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h

+//===-- TosaToLinalg.h - TOSA optimization pass declarations ----------*- 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 Linalg Dialect in MLIR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_TOSATOLINALG_TOSATOLINALG_H
+#define MLIR_CONVERSION_TOSATOLINALG_TOSATOLINALG_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace tosa {
+
+std::unique_ptr<Pass> createTosaToLinalgOnTensors();
+
+/// Populates passes to convert from TOSA to Linalg on buffers. At the end of
+/// the pass, the function will only contain linalg ops or standard ops if the
+/// pipeline succeeds.
+void addTosaToLinalgOnTensorsPasses(OpPassManager &pm);
+
+/// Populates conversion passes from TOSA dialect to Linalg dialect.
+void populateTosaToLinalgOnTensorsConversionPatterns(
+    MLIRContext *context, OwningRewritePatternList *patterns);
+
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_TOSATOLINALG_TOSATOLINALG_H

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToOpenMP)
 add_subdirectory(SCFToSPIRV)
 add_subdirectory(SCFToStandard)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
+add_subdirectory(TosaToLinalg)
 add_subdirectory(ArmSVEToLLVM)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)
 add_subdirectory(VectorToSPIRV)

mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt

+add_mlir_conversion_library(MLIRTosaToLinalg
+  TosaToLinalg.cpp
+  TosaToLinalgPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLinalg
+  MLIRLinalgUtils
+  MLIRPass
+  MLIRTosa
+  MLIRTosaTransforms
+  MLIRSupport
+  )

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

+//===- TosaToLinalg.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
+//
+//===----------------------------------------------------------------------===//
+//
+// These rewriters lower from the Tosa to the Linalg dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/Passes.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+static SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops) {
+  return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName());
+}
+
+static Value
+createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
+                                            ArrayRef<Type> resultTypes,
+                                            PatternRewriter &rewriter) {
+  Location loc = op->getLoc();
+  auto elementTy =
+      op->getResult(0).getType().cast<ShapedType>().getElementType();
+
+  // tosa::AbsOp
+  if (isa<tosa::AbsOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.create<mlir::AbsFOp>(loc, resultTypes, args);
+
+  // tosa::AddOp
+  if (isa<tosa::AddOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.create<mlir::AddFOp>(loc, resultTypes, args);
+
+  if (isa<tosa::AddOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.create<mlir::AddIOp>(loc, resultTypes, args);
+
+  // tosa::BitwiseAndOp
+  if (isa<tosa::BitwiseAndOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.create<mlir::AndOp>(loc, resultTypes, args);
+
+  // tosa::BitwiseOrOp
+  if (isa<tosa::BitwiseOrOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.create<mlir::OrOp>(loc, resultTypes, args);
+
+  // tosa::BitwiseXOrOp
+  if (isa<tosa::BitwiseXorOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.create<mlir::XOrOp>(loc, resultTypes, args);
+
+  // tosa::LogicalLeftShiftOp
+  if (isa<tosa::LogicalLeftShiftOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.create<mlir::ShiftLeftOp>(loc, resultTypes, args);
+
+  // tosa::LogicalrightShiftOp
+  if (isa<tosa::LogicalRightShiftOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.create<mlir::UnsignedShiftRightOp>(loc, resultTypes, args);
+
+  // tosa::PowOp
+  if (isa<tosa::PowOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.create<mlir::PowFOp>(loc, resultTypes, args);
+
+  // tosa::SubOp
+  if (isa<tosa::SubOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.create<mlir::SubFOp>(loc, resultTypes, args);
+
+  if (isa<tosa::SubOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.create<mlir::SubIOp>(loc, resultTypes, args);
+
+  // tosa::TanhOp
+  if (isa<tosa::TanhOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.create<mlir::TanhOp>(loc, resultTypes, args);
+
+  rewriter.notifyMatchFailure(
+      op, "unhandled op for linalg body calculation for elementwise op");
+  return nullptr;
+}
+
+static LogicalResult
+elementwiseMatchAndRewriteHelper(Operation *operation,
+                                 PatternRewriter &rewriter) {
+  auto loc = operation->getLoc();
+  auto results = operation->getResults();
+  auto t0 = operation->getOperand(0).getType().template dyn_cast<ShapedType>();
+  if (!t0)
+    return rewriter.notifyMatchFailure(operation,
+                                       "All results must be a shaped type");
+
+  // For now require no broadcasting. Consider making it support broadcasting
+  // operations.
+  Type uniqueTy = operation->getOperand(0).getType();
+  bool allInputTypesEqual =
+      llvm::all_of(operation->getOperandTypes(),
+                   [&](Type operandTy) { return operandTy == uniqueTy; });
+  if (!allInputTypesEqual)
+    return rewriter.notifyMatchFailure(operation,
+                                       "All operands must have the same type");
+  bool allResultTypesEqual =
+      llvm::all_of(operation->getResultTypes(),
+                   [&](Type resultTy) { return resultTy == uniqueTy; });
+  if (!allResultTypesEqual)
+    return rewriter.notifyMatchFailure(
+        operation, "All results must have the same type as the input");
+
+  // Construct the indexing maps needed for linalg.generic ops.
+  SmallVector<Type> bodyArgTypes;
+
+  for (Value in : operation->getOperands())
+    bodyArgTypes.emplace_back(getElementTypeOrSelf(in.getType()));
+
+  SmallVector<Type> opResultTypes;
+  SmallVector<Value> initTensors;
+  for (auto result : results) {
+    auto resultType = result.getType().template cast<ShapedType>();
+    if (!resultType.hasStaticShape())
+      return rewriter.notifyMatchFailure(
+          operation,
+          "tosa to linalg conversion expects statically shaped tensors");
+
+    initTensors.push_back(rewriter.create<linalg::InitTensorOp>(
+        loc, ArrayRef<Value>({}), resultType.getShape(),
+        resultType.getElementType()));
+    opResultTypes.push_back(result.getType());
+  }
+
+  auto bodyResultTypes = llvm::to_vector<4>(llvm::map_range(
+      initTensors, [](Value v) { return getElementTypeOrSelf(v); }));
+
+  // Supports only non-broadcasted operation. Shoudl consider update indexing
+  // map to be multidimensional.
+  unsigned nloops = t0.getRank();
+  AffineMap commonIndexingMap = rewriter.getMultiDimIdentityMap(nloops);
+  SmallVector<AffineMap, 2> indexingMaps(
+      operation->getNumOperands() + bodyResultTypes.size(), commonIndexingMap);
+
+  bool didEncounterError = false;
+  auto linalgOp = rewriter.create<linalg::GenericOp>(
+      loc, opResultTypes, operation->getOperands(), initTensors, indexingMaps,
+      getNParallelLoopsAttrs(nloops),
+      [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) {
+        Value opResult = createLinalgBodyCalculationForElementwiseOp(
+            operation, blockArgs.take_front(operation->getNumOperands()),
+            bodyResultTypes, rewriter);
+        if (opResult) {
+          didEncounterError = true;
+        }
+        nestedBuilder.create<linalg::YieldOp>(loc, opResult);
+      });
+
+  if (!didEncounterError)
+    return failure();
+
+  rewriter.replaceOp(operation, linalgOp->getResults());
+  return success();
+}
+
+namespace {
+
+template <typename SrcOp>
+class PointwiseConverter : public OpRewritePattern<SrcOp> {
+public:
+  using OpRewritePattern<SrcOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SrcOp op,
+                                PatternRewriter &rewriter) const final {
+    return elementwiseMatchAndRewriteHelper(op, rewriter);
+  }
+};
+
+} // namespace
+
+void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
+    MLIRContext *context, OwningRewritePatternList *patterns) {
+  patterns->insert<
+      PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
+      PointwiseConverter<tosa::PowOp>, PointwiseConverter<tosa::AbsOp>,
+      PointwiseConverter<tosa::TanhOp>, PointwiseConverter<tosa::BitwiseAndOp>,
+      PointwiseConverter<tosa::BitwiseOrOp>,
+      PointwiseConverter<tosa::BitwiseXorOp>,
+      PointwiseConverter<tosa::LogicalLeftShiftOp>,
+      PointwiseConverter<tosa::LogicalRightShiftOp>>(context);
+}

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/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TosaToLinalgOnTensors
+    : public TosaToLinalgOnTensorsBase<TosaToLinalgOnTensors> {
+public:
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect, StandardOpsDialect>();
+  }
+
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    ConversionTarget target(getContext());
+    target.addLegalDialect<linalg::LinalgDialect, StandardOpsDialect>();
+    target.addIllegalDialect<tosa::TosaDialect>();
+    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+
+    FuncOp func = getFunction();
+    mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
+        func.getContext(), &patterns);
+    if (failed(applyFullConversion(func, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::tosa::createTosaToLinalgOnTensors() {
+  return std::make_unique<TosaToLinalgOnTensors>();
+}
+
+void mlir::tosa::addTosaToLinalgOnTensorsPasses(OpPassManager &pm) {
+  pm.addNestedPass<FuncOp>(createTosaMakeBroadcastablePass());
+  pm.addNestedPass<FuncOp>(createTosaToLinalgOnTensors());
+}

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

+// RUN: mlir-opt --split-input-file --tosa-to-linalg-on-tensors %s -verify-diagnostics -o -| FileCheck %s
+
+// CHECK: #map = affine_map<() -> ()>
+
+// CHECK-LABEL: @test_abs
+func @test_abs(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [] : tensor<f32>
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = []} ins(%arg0 : tensor<f32>) outs([[INIT]] : tensor<f32>) {
+  // CHECK: ^bb0(%arg1: f32, %arg2: f32):
+  // CHECK:   [[ELEMENT:%.+]] = absf %arg1
+  // CHECK:   linalg.yield [[ELEMENT]] : f32
+  // CHECK: } -> tensor<f32>
+
+  %0 = "tosa.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
+
+  // CHECK: return [[GENERIC]]
+  return %0 : tensor<f32>
+}
+
+// -----
+
+// CHECK: #map = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: @test_abs
+func @test_abs(%arg0: tensor<1xf32>) -> tensor<1xf32> {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [1] : tensor<1xf32>
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<1xf32>) outs([[INIT]] : tensor<1xf32>) {
+  // CHECK: ^bb0(%arg1: f32, %arg2: f32):
+  // CHECK:   [[ELEMENT:%.+]] = absf %arg1
+  // CHECK:   linalg.yield [[ELEMENT]] : f32
+  // CHECK: } -> tensor<1xf32>
+  %0 = "tosa.abs"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
+
+  // CHECK: return [[GENERIC]]
+  return %0 : tensor<1xf32>
+}
+
+// -----
+
+// CHECK: #map = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @test_abs
+func @test_abs(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 2] : tensor<1x2xf32>
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<1x2xf32>) outs([[INIT]] : tensor<1x2xf32>) {
+  // CHECK: ^bb0(%arg1: f32, %arg2: f32):
+  // CHECK:   [[ELEMENT:%.+]] = absf %arg1
+  // CHECK:   linalg.yield [[ELEMENT]] : f32
+  // CHECK: } -> tensor<1x2xf32>
+  %0 = "tosa.abs"(%arg0) : (tensor<1x2xf32>) -> tensor<1x2xf32>
+
+  // CHECK: return [[GENERIC]]
+  return %0 : tensor<1x2xf32>
+}
+
+// -----
+
+func @test_add(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
+  // expected-error @+1 {{failed to legalize operation 'tosa.add'}}
+  %0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<2xf32>) -> tensor<2xf32>
+  return %0 : tensor<2xf32>
+}
+
+// -----
+
+func @test_add(%arg0: tensor<1xf32>, %arg1: tensor<f32>) -> tensor<1xf32> {
+  // expected-error @+1 {{failed to legalize operation 'tosa.add'}}
+  %0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<f32>) -> tensor<1xf32>
+  return %0 : tensor<1xf32>
+}
+
+// -----
+
+func @test_abs(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  // expected-error @+1 {{failed to legalize operation 'tosa.abs'}}
+  %0 = "tosa.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_simple_f32
+func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
+  // CHECK: linalg.generic
+  // CHECK: tanh
+  %0 = "tosa.tanh"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
+
+  // CHECK: linalg.generic
+  // CHECK: absf
+  %1 = "tosa.abs"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
+
+  // CHECK: linalg.generic
+  // CHECK: addf
+  %2 = "tosa.add"(%0, %0) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
+
+  // CHECK: linalg.generic
+  // CHECK: subf
+  %3 = "tosa.sub"(%0, %1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
+
+  // CHECK: linalg.generic
+  // CHECK: pow
+  %4 = "tosa.pow"(%1, %2) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @test_simple_i32
+func @test_simple_i32(%arg0: tensor<1xi32>) -> () {
+  // CHECK: linalg.generic
+  // CHECK: addi
+  %0 = "tosa.add"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: subi
+  %1 = "tosa.sub"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: and
+  %2 = "tosa.bitwise_and"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: or
+  %3 = "tosa.bitwise_or"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: xor
+  %4 = "tosa.bitwise_xor"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: shift_left
+  %5 = "tosa.logical_left_shift"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: shift_right_unsigned
+  %6 = "tosa.logical_right_shift"(%arg0, %arg0) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+
+  return
+}
+