Migrate compare test to check framework

Vulkan Swiftshader runs of this test with non-opt builds are very slow (I'll file an issue with some traces once I can figure out how to run swiftshader in OSS). There seems to be a ton of overhead for PipelineExecutable::Create that scales with the size of the module. This doesn't come up in iree-run-mlir because it treats each test as a separate module (and actually recreates the driver per-test). After some discussion we decided to proceed with the migration and circle back to debug the swiftshader performance issues.

PiperOrigin-RevId: 306346899
diff --git a/iree/test/e2e/xla_ops/BUILD b/iree/test/e2e/xla_ops/BUILD
index 6b4dd0e..e5cdb85 100644
--- a/iree/test/e2e/xla_ops/BUILD
+++ b/iree/test/e2e/xla_ops/BUILD
@@ -34,6 +34,7 @@
     "batch_norm_inference.mlir",
     "broadcast.mlir",
     "broadcast_in_dim.mlir",
+    "compare.mlir",
     "concatenate.mlir",
     "constant.mlir",
     "convert_int.mlir",
@@ -75,6 +76,7 @@
     "add.mlir",
     "broadcast.mlir",
     "broadcast_in_dim.mlir",
+    "compare.mlir",
     "concatenate.mlir",
     "constant.mlir",
     "convert_int.mlir",
@@ -136,6 +138,7 @@
     "add.mlir",
     "batch_norm_inference.mlir",
     "broadcast_in_dim.mlir",
+    "compare.mlir",
     "constant.mlir",
     "cos.mlir",
     "dot.mlir",
@@ -163,7 +166,6 @@
 ]
 
 LEGACY_LIT_OP_TESTS = [
-    "compare.mlir",
     "conv.mlir",
 ]
 
diff --git a/iree/test/e2e/xla_ops/CMakeLists.txt b/iree/test/e2e/xla_ops/CMakeLists.txt
index 0a642e6..7cb4570 100644
--- a/iree/test/e2e/xla_ops/CMakeLists.txt
+++ b/iree/test/e2e/xla_ops/CMakeLists.txt
@@ -18,7 +18,6 @@
   NAME
     lit
   SRCS
-    "compare.mlir"
     "conv.mlir"
   DATA
     iree::tools::IreeFileCheck
@@ -36,6 +35,7 @@
     "batch_norm_inference.mlir"
     "broadcast.mlir"
     "broadcast_in_dim.mlir"
+    "compare.mlir"
     "concatenate.mlir"
     "constant.mlir"
     "convert_int.mlir"
@@ -84,6 +84,7 @@
     "add.mlir"
     "broadcast.mlir"
     "broadcast_in_dim.mlir"
+    "compare.mlir"
     "concatenate.mlir"
     "constant.mlir"
     "convert_int.mlir"
@@ -161,6 +162,7 @@
     "add.mlir"
     "batch_norm_inference.mlir"
     "broadcast_in_dim.mlir"
+    "compare.mlir"
     "constant.mlir"
     "cos.mlir"
     "dot.mlir"
diff --git a/iree/test/e2e/xla_ops/compare.mlir b/iree/test/e2e/xla_ops/compare.mlir
index af4552e..d6e4c78 100644
--- a/iree/test/e2e/xla_ops/compare.mlir
+++ b/iree/test/e2e/xla_ops/compare.mlir
@@ -1,210 +1,164 @@
-// RUN: iree-run-mlir -iree-hal-target-backends=vmla %s | IreeFileCheck %s
-// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -iree-hal-target-backends=llvm-ir %s | IreeFileCheck %s)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -iree-hal-target-backends=vulkan-spirv %s | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @compare_tensor
-func @compare_tensor() -> tensor<4xi8> {
+func @compare_tensor() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
-  return %output : tensor<4xi8>
+  check.expect_eq_const(%output, dense<[0, 1, 0, 1]> : tensor<4xi8>) : tensor<4xi8>
+  return
 }
-// CHECK: 4xi8=0 1 0 1
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_scalar
-func @compare_scalar() -> tensor<i8> {
+func @compare_scalar() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<1> : tensor<i32>
   %rhs = iree.unfoldable_constant dense<5> : tensor<i32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
-  return %output : tensor<i8>
+  check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+  return
 }
-// CHECK: i8=0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_i8
-func @compare_i8() -> tensor<i8> {
+func @compare_i8() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<1> : tensor<i8>
   %rhs = iree.unfoldable_constant dense<5> : tensor<i8>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i8>, tensor<i8>) -> tensor<i1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
-  return %output : tensor<i8>
+  check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+  return
 }
-// CHECK: i8=0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_i16
-func @compare_i16() -> tensor<i8> {
+func @compare_i16() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<1> : tensor<i16>
   %rhs = iree.unfoldable_constant dense<5> : tensor<i16>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i16>, tensor<i16>) -> tensor<i1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
-  return %output : tensor<i8>
+  check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+  return
 }
-// CHECK: i8=0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_i32
-func @compare_i32() -> tensor<i8> {
+func @compare_i32() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<1> : tensor<i32>
   %rhs = iree.unfoldable_constant dense<5> : tensor<i32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
-  return %output : tensor<i8>
+  check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+  return
 }
-// CHECK: i8=0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_i64
-func @compare_i64() -> tensor<i8> {
+func @compare_i64() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<1> : tensor<i64>
   %rhs = iree.unfoldable_constant dense<5> : tensor<i64>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i64>, tensor<i64>) -> tensor<i1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
-  return %output : tensor<i8>
+  check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+  return
 }
-// CHECK: i8=0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_f32
-func @compare_f32() -> tensor<i8> {
+func @compare_f32() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<1.0> : tensor<f32>
   %rhs = iree.unfoldable_constant dense<5.0> : tensor<f32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
-  return %output : tensor<i8>
+  check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+  return
 }
-// CHECK: i8=0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_f64
-func @compare_f64() -> tensor<i8> {
+func @compare_f64() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<1.0> : tensor<f64>
   %rhs = iree.unfoldable_constant dense<5.0> : tensor<f64>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<f64>, tensor<f64>) -> tensor<i1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
-  return %output : tensor<i8>
+  check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+  return
 }
-// CHECK: i8=0
-// -----
 
-// CHECK-LABEL: EXEC @compare_tensor_odd_length
-func @compare_tensor_odd_length() -> tensor<3xi8> {
+func @compare_tensor_odd_length() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7]> : tensor<3xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3]> : tensor<3xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<3xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<3xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<3xi1>, tensor<3xi8>, tensor<3xi8>) -> tensor<3xi8>
-  return %output : tensor<3xi8>
+  check.expect_eq_const(%output, dense<[0, 1, 0]> : tensor<3xi8>) : tensor<3xi8>
+  return
 }
-// CHECK: 3xi8=0 1 0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_eq
-func @compare_eq() -> tensor<4xi8> {
+func @compare_eq() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
-  return %output : tensor<4xi8>
+  check.expect_eq_const(%output, dense<[0, 1, 0, 1]> : tensor<4xi8>) : tensor<4xi8>
+  return
 }
-// CHECK: 4xi8=0 1 0 1
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_ne
-func @compare_ne() -> tensor<4xi8> {
+func @compare_ne() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "NE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
-  return %output : tensor<4xi8>
+  check.expect_eq_const(%output, dense<[1, 0, 1, 0]> : tensor<4xi8>) : tensor<4xi8>
+  return
 }
-// CHECK: 4xi8=1 0 1 0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_lt
-func @compare_lt() -> tensor<4xi8> {
+func @compare_lt() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "LT"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
-  return %output : tensor<4xi8>
+  check.expect_eq_const(%output, dense<[1, 0, 0, 0]> : tensor<4xi8>) : tensor<4xi8>
+  return
 }
-// CHECK: 4xi8=1 0 0 0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_le
-func @compare_le() -> tensor<4xi8> {
+func @compare_le() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "LE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
-  return %output : tensor<4xi8>
+  check.expect_eq_const(%output, dense<[1, 1, 0, 1]> : tensor<4xi8>) : tensor<4xi8>
+  return
 }
-// CHECK: 4xi8=1 1 0 1
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_gt
-func @compare_gt() -> tensor<4xi8> {
+func @compare_gt() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "GT"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
-  return %output : tensor<4xi8>
+  check.expect_eq_const(%output, dense<[0, 0, 1, 0]> : tensor<4xi8>) : tensor<4xi8>
+  return
 }
-// CHECK: 4xi8=0 0 1 0
 
-// -----
-
-// CHECK-LABEL: EXEC @compare_ge
-func @compare_ge() -> tensor<4xi8> {
+func @compare_ge() attributes { iree.module.export } {
   %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
   %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
   %result = "xla_hlo.compare"(%lhs, %rhs) {comparison_direction = "GE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
   %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
   %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
   %output = "xla_hlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
-  return %output : tensor<4xi8>
+  check.expect_eq_const(%output, dense<[0, 1, 1, 1]> : tensor<4xi8>) : tensor<4xi8>
+  return
 }
-// CHECK: 4xi8=0 1 1 1