From 339b17e6aad6564d87480cefb209db570bc656a3 Mon Sep 17 00:00:00 2001
From: eedalong <yuanxiulong.yxl@alibaba-inc.com>
Date: Mon, 22 Jan 2024 14:13:42 +0800
Subject: [PATCH] add ut

---
 .../tests/disc-input-output-alias-inplace.mlir      | 11 +++++++++++
 .../tests/disc-input-output-alias-propogation.mlir  | 13 +++++++++++++
 2 files changed, 24 insertions(+)
 create mode 100755 tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-inplace.mlir
 create mode 100644 tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-propogation.mlir

diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-inplace.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-inplace.mlir
new file mode 100755
index 00000000000..e0248e27252
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-inplace.mlir
@@ -0,0 +1,11 @@
+// RUN: disc-opt -disc-input-output-alias \
+// RUN: %s -o - | FileCheck %s
+
+// CHECK-LABEL: main
+func.func @main(%arg0: tensor<200x200xf32>, %arg1: tensor<200x200xf32>) -> (tensor<200x200xf32>, tensor<200x200xf32>) attributes {tf.entry_function = {input_output_alias_outputs = "0,1", input_output_alias_params = "0,1", input_placements = "gpu,gpu", output_placements = "gpu,gpu"}} {
+    // CHECK: %0 = mhlo.add %arg1, %arg0 : tensor<200x200xf32>    
+    %0 = mhlo.add %arg1, %arg0 : tensor<200x200xf32>
+    // CHECK: "mhlo_disc.args_mutation"(%0, %arg1) : (tensor<200x200xf32>, tensor<200x200xf32>) -> ()
+    // CHECK: return %arg0, %0 : tensor<200x200xf32>, tensor<200x200xf32>
+    return %arg0, %0 : tensor<200x200xf32>, tensor<200x200xf32>
+}
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-propogation.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-propogation.mlir
new file mode 100644
index 00000000000..3bb32a22219
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-input-output-alias-propogation.mlir
@@ -0,0 +1,13 @@
+// RUN: disc-opt -disc-input-output-alias \
+// RUN: %s -o - | FileCheck %s
+
+// CHECK-LABEL: main
+func.func @main(%arg0: tensor<200x200xf32>, %arg1: tensor<200x200xf32>) -> (tensor<200x200xf32>, tensor<200x200xf32>) attributes {tf.entry_function = {input_output_alias_outputs = "0,1", input_output_alias_params = "0,1", input_placements = "gpu,gpu", output_placements = "gpu,gpu"}} {
+    // CHECK: %0 = mhlo.add %arg1, %arg0 : tensor<200x200xf32>    
+    %0 = mhlo.add %arg1, %arg0 : tensor<200x200xf32>
+    // CHECK: "mhlo_disc.args_mutation"(%0, %arg1) : (tensor<200x200xf32>, tensor<200x200xf32>) -> ()
+    %1 = mhlo.add %0, %arg0 : tensor<200x200xf32>
+    // CHECK: "mhlo_disc.args_mutation"(%1, %0) : (tensor<200x200xf32>, tensor<200x200xf32>) -> ()
+    // CHECK: return %arg0, %1 : tensor<200x200xf32>, tensor<200x200xf32>
+    return %arg0, %1 : tensor<200x200xf32>, tensor<200x200xf32>
+}
\ No newline at end of file