Migrate constant test to check framework

PiperOrigin-RevId: 306327708

iree/test/e2e/xla_ops/BUILD

diff --git a/iree/test/e2e/xla_ops/BUILD b/iree/test/e2e/xla_ops/BUILD
index 87d9b30..6b4dd0e 100644
--- a/iree/test/e2e/xla_ops/BUILD
+++ b/iree/test/e2e/xla_ops/BUILD
@@ -35,6 +35,7 @@
     "broadcast.mlir",
     "broadcast_in_dim.mlir",
     "concatenate.mlir",
+    "constant.mlir",
     "convert_int.mlir",
     "cos.mlir",
     "dot.mlir",
@@ -75,6 +76,7 @@
     "broadcast.mlir",
     "broadcast_in_dim.mlir",
     "concatenate.mlir",
+    "constant.mlir",
     "convert_int.mlir",
     "cos.mlir",
     "dot.mlir",
@@ -111,6 +113,7 @@
     "add.mlir",
     "batch_norm_inference.mlir",
     "broadcast_in_dim.mlir",
+    "constant.mlir",
     "cos.mlir",
     "exp.mlir",
     "gemm.mlir",
@@ -133,6 +136,7 @@
     "add.mlir",
     "batch_norm_inference.mlir",
     "broadcast_in_dim.mlir",
+    "constant.mlir",
     "cos.mlir",
     "dot.mlir",
     "exp.mlir",
@@ -160,7 +164,6 @@
 
 LEGACY_LIT_OP_TESTS = [
     "compare.mlir",
-    "constants.mlir",
     "conv.mlir",
 ]
 

iree/test/e2e/xla_ops/CMakeLists.txt

diff --git a/iree/test/e2e/xla_ops/CMakeLists.txt b/iree/test/e2e/xla_ops/CMakeLists.txt
index 5c4c4f8..0a642e6 100644
--- a/iree/test/e2e/xla_ops/CMakeLists.txt
+++ b/iree/test/e2e/xla_ops/CMakeLists.txt
@@ -19,7 +19,6 @@
     lit
   SRCS
     "compare.mlir"
-    "constants.mlir"
     "conv.mlir"
   DATA
     iree::tools::IreeFileCheck
@@ -38,6 +37,7 @@
     "broadcast.mlir"
     "broadcast_in_dim.mlir"
     "concatenate.mlir"
+    "constant.mlir"
     "convert_int.mlir"
     "cos.mlir"
     "dot.mlir"
@@ -85,6 +85,7 @@
     "broadcast.mlir"
     "broadcast_in_dim.mlir"
     "concatenate.mlir"
+    "constant.mlir"
     "convert_int.mlir"
     "cos.mlir"
     "dot.mlir"
@@ -128,6 +129,7 @@
     "add.mlir"
     "batch_norm_inference.mlir"
     "broadcast_in_dim.mlir"
+    "constant.mlir"
     "cos.mlir"
     "exp.mlir"
     "gemm.mlir"
@@ -159,6 +161,7 @@
     "add.mlir"
     "batch_norm_inference.mlir"
     "broadcast_in_dim.mlir"
+    "constant.mlir"
     "cos.mlir"
     "dot.mlir"
     "exp.mlir"

iree/test/e2e/xla_ops/constant.mlir

diff --git a/iree/test/e2e/xla_ops/constant.mlir b/iree/test/e2e/xla_ops/constant.mlir
new file mode 100644
index 0000000..9cbd098
--- /dev/null
+++ b/iree/test/e2e/xla_ops/constant.mlir
@@ -0,0 +1,26 @@
+func @high_rank () attributes { iree.module.export } {
+  %dense = xla_hlo.constant dense<[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]> : tensor<2x2x3xi32>
+  check.expect_eq_const(%dense, dense<[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]> : tensor<2x2x3xi32>) : tensor<2x2x3xi32>
+
+  %splat = xla_hlo.constant dense<1> : tensor<2x2x3xi32>
+  check.expect_eq_const(%splat, dense<1> : tensor<2x2x3xi32>) : tensor<2x2x3xi32>
+  return
+}
+
+func @i8() attributes { iree.module.export } {
+  %c = xla_hlo.constant dense<[1, 2]> : tensor<2xi8>
+  check.expect_eq_const(%c, dense<[1, 2]> : tensor<2xi8>) : tensor<2xi8>
+  return
+}
+
+func @i32 () attributes { iree.module.export } {
+  %c = xla_hlo.constant dense<[1, 2]> : tensor<2xi32>
+  check.expect_eq_const(%c,  dense<[1, 2]> : tensor<2xi32>) : tensor<2xi32>
+  return
+}
+
+func @f32 () attributes { iree.module.export } {
+  %c = xla_hlo.constant dense<[1.1, 2.1]> : tensor<2xf32>
+  check.expect_almost_eq_const(%c, dense<[1.1, 2.1]> : tensor<2xf32>) : tensor<2xf32>
+  return
+}

iree/test/e2e/xla_ops/constants.mlir

diff --git a/iree/test/e2e/xla_ops/constants.mlir b/iree/test/e2e/xla_ops/constants.mlir
deleted file mode 100644
index 1ad2dae..0000000
--- a/iree/test/e2e/xla_ops/constants.mlir
+++ /dev/null
@@ -1,33 +0,0 @@
-// 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)
-// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -iree-hal-target-backends=vulkan-spirv -iree-use-linalg-to-spirv-path %s | IreeFileCheck %s)
-
-// CHECK-LABEL: EXEC @xla_constant_i32
-func @xla_constant_i32 () -> (tensor<2x2x3xi32>, tensor<2x2x3xi32>) {
-  %0 = iree.unfoldable_constant dense<[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]> : tensor<2x2x3xi32>
-  %1 = iree.unfoldable_constant dense<1> : tensor<2x2x3xi32>
-  return %0, %1: tensor<2x2x3xi32>, tensor<2x2x3xi32>
-}
-// CHECK: 2x2x3xi32={{\[}}[1 2 3][4 5 6]]{{\[}}[7 8 9][10 11 12]]
-// CHECK-NEXT: 2x2x3xi32={{\[}}[1 1 1][1 1 1]]{{\[}}[1 1 1][1 1 1]]
-
-// -----
-
-// CHECK-LABEL: EXEC @xla_constant_f32
-func @xla_constant_f32 () -> (tensor<2x2x3xf32>, tensor<2x2x3xf32>) {
-  %0 = iree.unfoldable_constant dense<[[[1.1e0, 2.1e0, 3.1e0], [4.1e0, 5.1e0, 6.1e0]], [[7.1e0, 8.1e0, 9.1e0], [10.1e0, 11.1e0, 12.1e0]]]> : tensor<2x2x3xf32>
-  %1 = iree.unfoldable_constant dense<1.1e0> : tensor<2x2x3xf32>
-  return %0, %1: tensor<2x2x3xf32>, tensor<2x2x3xf32>
-}
-// CHECK: 2x2x3xf32={{\[}}[1.1 2.1 3.1][4.1 5.1 6.1]]{{\[}}[7.1 8.1 9.1][10.1 11.1 12.1]]
-// CHECK-NEXT: 2x2x3xf32={{\[}}[1.1 1.1 1.1][1.1 1.1 1.1]]{{\[}}[1.1 1.1 1.1][1.1 1.1 1.1]]
-
-// -----
-
-// CHECK-LABEL: EXEC @xla_constant_i8
-func @xla_constant_i8() -> tensor<3xi8> {
-  %0 = iree.unfoldable_constant dense<1> : tensor<3xi8>
-  return %0 : tensor<3xi8>
-}
-// CHECK: 3xi8=1 1 1