Skip to content

Commit

Permalink
add ut
Browse files Browse the repository at this point in the history
  • Loading branch information
eedalong committed Jan 22, 2024
1 parent 9f78a57 commit 339b17e
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -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>
}
Original file line number Diff line number Diff line change
@@ -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>
}

0 comments on commit 339b17e

Please sign in to comment.