Merge pull request #189 from Xilinx/christopher.fix_clamp_canonicalization_loc

fix(TosaCanonicalize): create FusedLoc when folding clamp

Author: ttjost

mlir/include/mlir/IR/PatternMatch.h

@@ -532,14 +532,24 @@ class RewriterBase : public OpBuilder {
   virtual void replaceOp(Operation *op, Operation *newOp);
 
   /// Replaces the result op with a new op that is created without verification.
+  /// Use a given list of locations to generate a FusedLoc for the new op.
   /// The result values of the two ops must be the same types.
   template <typename OpTy, typename... Args>
-  OpTy replaceOpWithNewOp(Operation *op, Args &&...args) {
-    auto newOp = create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
+  OpTy replaceOpWithNewOp(Operation *op, ArrayRef<Location> locs,
+                          Args &&...args) {
+    auto newOp = create<OpTy>(getFusedLoc(locs), std::forward<Args>(args)...);
     replaceOp(op, newOp.getOperation());
     return newOp;
   }
 
+  /// Replaces the result op with a new op that is created without verification.
+  /// The result values of the two ops must be the same types.
+  template <typename OpTy, typename... Args>
+  OpTy replaceOpWithNewOp(Operation *op, Args &&...args) {
+    return replaceOpWithNewOp<OpTy, Args...>(op, {op->getLoc()},
+                                             std::forward<Args>(args)...);
+  }
+
   /// This method erases an operation that is known to have no uses.
   virtual void eraseOp(Operation *op);
 

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

@@ -404,9 +404,9 @@ struct ClampClampOptimization : public OpRewritePattern<tosa::ClampOp> {
       auto minInt = std::max(op.getMinInt(), clampOp.getMinInt());
       auto maxInt = std::min(op.getMaxInt(), clampOp.getMaxInt());
 
-      rewriter.replaceOpWithNewOp<tosa::ClampOp>(
-          op, op.getType(), clampOp.getInput(),
-          rewriter.getI64IntegerAttr(minInt),
+      rewriter.replaceOpWithNewOp<ClampOp>(
+          op, {op->getLoc(), clampOp->getLoc()}, op.getType(),
+          clampOp.getInput(), rewriter.getI64IntegerAttr(minInt),
           rewriter.getI64IntegerAttr(maxInt), rewriter.getF32FloatAttr(minFp),
           rewriter.getF32FloatAttr(maxFp));
       return success();

mlir/test/Dialect/Tosa/canonicalize_with_debuginfo.mlir

@@ -0,0 +1,14 @@
+// RUN: mlir-opt -mlir-print-debuginfo -canonicalize="test-convergence" %s | FileCheck %s
+
+// CHECK-LABEL: @clamp_twice_is_single_clamp
+func.func @clamp_twice_is_single_clamp(%arg0: tensor<4xi8>) -> tensor<4xi8> {
+  // CHECK: tosa.clamp %arg0 {max_fp = 3.000000e+00 : f32, max_int = 2 : i64, min_fp = -3.000000e+00 : f32, min_int = -2 : i64} {{.*}} loc(#[[FUSED:.*]])
+  // CHECK-DAG: #[[A:.*]] = loc("Clamp_A")
+  // CHECK-DAG: #[[B:.*]] = loc("Clamp_B")
+  // CHECK:     #[[FUSED]] = loc(fused[#[[B]], #[[A]]])
+  %0 = tosa.clamp %arg0 {max_fp = 3.0 : f32, max_int = 4 : i64, min_fp = -5.0 : f32, min_int = -2 : i64} :  (tensor<4xi8>) -> tensor<4xi8> loc(#loc0)
+  %1 = tosa.clamp %0 {max_fp = 5.0 : f32, max_int = 2 : i64, min_fp = -3.0 : f32, min_int = -4 : i64} :  (tensor<4xi8>) -> tensor<4xi8> loc(#loc1)
+  return %1 : tensor<4xi8>
+}
+#loc0 = loc("Clamp_A")
+#loc1 = loc("Clamp_B")