[Mlir-commits] [mlir] f685c9a - [MLIR][TOSA] Lower tosa.identity and tosa.identitiyn to linalg

Rob Suderman llvmlistbot at llvm.org
Fri Feb 26 15:53:48 PST 2021


Author: Rob Suderman
Date: 2021-02-26T15:45:07-08:00
New Revision: f685c9ac8647daeb425b05ffb9c555ccfc7ec78b

URL: https://github.com/llvm/llvm-project/commit/f685c9ac8647daeb425b05ffb9c555ccfc7ec78b
DIFF: https://github.com/llvm/llvm-project/commit/f685c9ac8647daeb425b05ffb9c555ccfc7ec78b.diff

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

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.

Differential Revision: https://reviews.llvm.org/D97517

Added: 
    

Modified: 
    mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
    mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 75bbc46c4804..f6cf7ed709a3 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -442,6 +442,21 @@ class ReshapeOpConverter : public OpConversionPattern<tosa::ReshapeOp> {
   }
 };
 
+// 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(
@@ -462,5 +477,6 @@ void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
       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);
 }

diff  --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 985aec3d212b..829686712ca1 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -304,3 +304,16 @@ func @test_reshape_downrank_6D(%arg0: tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77
   %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>
+}


        


More information about the Mlir-commits mailing list