Adding mhlo.rng_uniform e2e tests. (#6996)
The tests are derived from https://github.com/google/iree/pull/6676
We can enable the tests because the lowering was landed upstream.
diff --git a/iree/test/e2e/xla_ops/BUILD b/iree/test/e2e/xla_ops/BUILD
index c74fc1e..9a4e62b 100644
--- a/iree/test/e2e/xla_ops/BUILD
+++ b/iree/test/e2e/xla_ops/BUILD
@@ -62,6 +62,7 @@
"remainder.mlir",
"reshape.mlir",
"reverse.mlir",
+ "rng_uniform.mlir",
"rsqrt.mlir",
"scatter.mlir",
"scatter_dynamic.mlir",
@@ -207,6 +208,7 @@
"remainder.mlir",
"reshape.mlir",
"reverse.mlir",
+ "rng_uniform.mlir",
"rsqrt.mlir",
"scatter.mlir",
"scatter_dynamic.mlir",
@@ -291,6 +293,7 @@
"scatter.mlir", # TODO(GH-6601): Enable the test.
"scatter_dynamic.mlir", # TODO(GH-6601): Enable the test.
"sort.mlir",
+ "rng_uniform.mlir", # TODO(GH-6995): Enable the test.
],
),
compiler_flags = ["-iree-input-type=mhlo"],
@@ -340,6 +343,7 @@
"remainder.mlir",
"reshape.mlir",
"reverse.mlir",
+ "rng_uniform.mlir",
"rsqrt.mlir",
"scatter.mlir",
"scatter_dynamic.mlir",
diff --git a/iree/test/e2e/xla_ops/CMakeLists.txt b/iree/test/e2e/xla_ops/CMakeLists.txt
index 268830b..392dfa4 100644
--- a/iree/test/e2e/xla_ops/CMakeLists.txt
+++ b/iree/test/e2e/xla_ops/CMakeLists.txt
@@ -52,6 +52,7 @@
"remainder.mlir"
"reshape.mlir"
"reverse.mlir"
+ "rng_uniform.mlir"
"rsqrt.mlir"
"scatter.mlir"
"scatter_dynamic.mlir"
@@ -121,6 +122,7 @@
"remainder.mlir"
"reshape.mlir"
"reverse.mlir"
+ "rng_uniform.mlir"
"rsqrt.mlir"
"scatter.mlir"
"scatter_dynamic.mlir"
@@ -242,6 +244,7 @@
"remainder.mlir"
"reshape.mlir"
"reverse.mlir"
+ "rng_uniform.mlir"
"rsqrt.mlir"
"scatter.mlir"
"scatter_dynamic.mlir"
diff --git a/iree/test/e2e/xla_ops/rng_uniform.mlir b/iree/test/e2e/xla_ops/rng_uniform.mlir
new file mode 100644
index 0000000..7e496db
--- /dev/null
+++ b/iree/test/e2e/xla_ops/rng_uniform.mlir
@@ -0,0 +1,34 @@
+// Note that they are stateless random generators, so they have fixed results.
+func @rng_uniform_1d() {
+ %min = util.unfoldable_constant dense<-10.0> : tensor<f32>
+ %max = util.unfoldable_constant dense<10.0> : tensor<f32>
+ %shape = util.unfoldable_constant dense<[10]> : tensor<1xi32>
+ %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<1xi32>) -> tensor<10xf32>
+ check.expect_almost_eq_const(%res, dense<[
+ -9.99994, -4.8613, 0.277344, 5.41599, -9.44537, -4.30673, 0.831918, 5.97056, -8.8908, -3.75215
+ ]> : tensor<10xf32>) : tensor<10xf32>
+ return
+}
+
+func @rng_uniform_2d() {
+ %min = util.unfoldable_constant dense<-10.0> : tensor<f32>
+ %max = util.unfoldable_constant dense<10.0> : tensor<f32>
+ %shape = util.unfoldable_constant dense<[3, 3]> : tensor<2xi32>
+ %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<3x3xf32>
+ check.expect_almost_eq_const(%res, dense<[
+ [6.55154, -8.30982, -3.17117],
+ [1.75741, 6.89606, -7.9653],
+ [-3.03671, 2.10193, 7.24057]]> : tensor<3x3xf32>) : tensor<3x3xf32>
+ return
+}
+
+func @rng_uniform_3d() {
+ %min = util.unfoldable_constant dense<-10.0> : tensor<f32>
+ %max = util.unfoldable_constant dense<10.0> : tensor<f32>
+ %shape = util.unfoldable_constant dense<[2, 2, 2]> : tensor<3xi32>
+ %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<3xi32>) -> tensor<2x2x2xf32>
+ check.expect_almost_eq_const(%res, dense<[
+ [[3.04814, 8.18679], [-1.74598, 3.39266]],
+ [[-6.91349, -1.77484], [8.29239, -6.56897]]]> : tensor<2x2x2xf32>) : tensor<2x2x2xf32>
+ return
+}
