[MLIR][TOSA] Lower tosa.identity and tosa.identitiyn to linalg

[rsuderman]

Summary

Both identity ops can be loweried by replacing their results with their
inputs. We keep this as a linalg lowering as other backends may choose to
create copies.

[silvas]

Can you add a comment regarding which scenarios identity is load bearing and for which backends?

I was really hoping we could avoid the whole TF mess of "identity is load bearing".

[silvas]

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

259

can you check that the values were actually replaced properly? (probably a separate function is a good place to put them so you can check args/returns).

[rsuderman]

Added separate test for identitiy and added details for lowering.

[rsuderman]

Rebased

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

matchAndRewrite(tosa::ReshapeOp reshape, ArrayRef<Value> args,
     }
 
     rewriter.replaceOpWithNewOp<linalg::TensorReshapeOp>(
         reshape, resultTy, args[0], reassociationMap);
     return success();
   }
 };
 
+// At the codegen level any identity operations should be removed. Any cases
+// where identity is load-bearing (e.g. cross device computation) should be
+// handled before lowering to codegen.
+template <typename SrcOp>
+class IdentityNConverter : public OpRewritePattern<SrcOp> {
+public:
+  using OpRewritePattern<SrcOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SrcOp op,
+                                PatternRewriter &rewriter) const final {
+    rewriter.replaceOp(op, op.getOperation()->getOperands());
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
     MLIRContext *context, OwningRewritePatternList *patterns) {
   patterns->insert<
       PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
       PointwiseConverter<tosa::MulOp>, PointwiseConverter<tosa::NegateOp>,
       PointwiseConverter<tosa::PowOp>, PointwiseConverter<tosa::RsqrtOp>,
       PointwiseConverter<tosa::LogOp>, PointwiseConverter<tosa::ExpOp>,
       PointwiseConverter<tosa::AbsOp>, PointwiseConverter<tosa::TanhOp>,
       PointwiseConverter<tosa::BitwiseAndOp>,
       PointwiseConverter<tosa::BitwiseOrOp>,
       PointwiseConverter<tosa::BitwiseXorOp>,
       PointwiseConverter<tosa::LogicalLeftShiftOp>,
       PointwiseConverter<tosa::LogicalRightShiftOp>,
       PointwiseConverter<tosa::SelectOp>, PointwiseConverter<tosa::GreaterOp>,
       PointwiseConverter<tosa::GreaterEqualOp>,
       PointwiseConverter<tosa::MaximumOp>, PointwiseConverter<tosa::MinimumOp>,
       PointwiseConverter<tosa::CeilOp>, PointwiseConverter<tosa::FloorOp>,
       PointwiseConverter<tosa::ClampOp>, PointwiseConverter<tosa::ReluNOp>,
-      ReshapeOpConverter>(context);
+      IdentityNConverter<tosa::IdentityOp>,
+      IdentityNConverter<tosa::IdentityNOp>, ReshapeOpConverter>(context);
 }

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

func @test_simple_i32(%arg0: tensor<1xi32>) -> () {
   %14 = "tosa.clamp"(%0) {min_int = 1 : i64, max_int = 5 : i64, min_fp = 1.0 : f32, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 
   // CHECK: linalg.generic
   // CHECK: cmpi
   // CHECK: select
   %15 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 
   return
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: @test_reshape_downrank
 func @test_reshape_downrank(%arg0: tensor<2x3xf32>) -> tensor<6xf32> {
   // CHECK: [[RESHAPE:%.+]] = linalg.tensor_reshape %arg0 [#[[$MAP0]]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [6]} : (tensor<2x3xf32>) -> tensor<6xf32>

 // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
 
 // CHECK-LABEL: @test_reshape_downrank_6D
 func @test_reshape_downrank_6D(%arg0: tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32> {
   // CHECK: linalg.tensor_reshape %arg0 [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [2, 3]} : (tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32>
   return %0 : tensor<6x5x77xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @test_identity
+func @test_identity(%arg0: tensor<1xf32>, %arg1: tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>) {
+  %0 = "tosa.identity"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
+  %1 = "tosa.identity"(%arg1) : (tensor<1xi32>) -> tensor<1xi32>
+
+  %2:2 = "tosa.identityn"(%0, %1) : (tensor<1xf32>, tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>)
+
+  // CHECK: return %arg0, %arg1
+  return %2#0, %2#1 : tensor<1xf32>, tensor<1xi32>
+}