commit 301091ac3cf4d2de7bda116547dde04c0571a70e
author Lei Zhang <antiagainst@google.com> Thu Mar 17 09:08:06 2022 -0400
committer GitHub <noreply@github.com> Thu Mar 17 06:08:06 2022 -0700
tree 61a1e8887d5d1c933393dae1041c8c104c9bd916
parent 3fe8427e613659a3140162acdc4523e2becee4ad

Re-enable large_dot_general2 for Vulkan (#8561)

This was hitting a GPU driver bug in Pixel 4.

Fixes https://github.com/google/iree/issues/6070

iree/test/e2e/vulkan_specific/dot_general.mlir

diff --git a/iree/test/e2e/vulkan_specific/dot_general.mlir b/iree/test/e2e/vulkan_specific/dot_general.mlir
index a5098f4..af2b39c 100644
--- a/iree/test/e2e/vulkan_specific/dot_general.mlir
+++ b/iree/test/e2e/vulkan_specific/dot_general.mlir
@@ -48,19 +48,18 @@
   return
 }
 
-// TODO(#6070): Re-enable this after fixing wrong answer on Pixel 4.
-// func @large_dot_general2() {
-//   %lhs = util.unfoldable_constant dense<1.0> : tensor<4x32x1024xf32>
-//   %rhs = util.unfoldable_constant dense<0.4> : tensor<4x1024x64xf32>
-//   %res = "mhlo.dot_general"(%lhs, %rhs) {
-//            dot_dimension_numbers = {
-//              lhs_batching_dimensions = dense<0> : tensor<1xi64>,
-//              lhs_contracting_dimensions = dense<2> : tensor<1xi64>,
-//              rhs_batching_dimensions = dense<0> : tensor<1xi64>,
-//              rhs_contracting_dimensions = dense<1> : tensor<1xi64>
-//            },
-//            precision_config = ["DEFAULT", "DEFAULT"]
-//          } : (tensor<4x32x1024xf32>, tensor<4x1024x64xf32>) -> tensor<4x32x64xf32>
-//   check.expect_almost_eq_const(%res, dense<409.596> : tensor<4x32x64xf32>) : tensor<4x32x64xf32>
-//   return
-// }
+func @large_dot_general2() {
+  %lhs = util.unfoldable_constant dense<1.0> : tensor<4x32x1024xf32>
+  %rhs = util.unfoldable_constant dense<0.4> : tensor<4x1024x64xf32>
+  %res = "mhlo.dot_general"(%lhs, %rhs) {
+           dot_dimension_numbers = #mhlo.dot<
+             lhs_batching_dimensions = [0],
+             lhs_contracting_dimensions = [2],
+             rhs_batching_dimensions = [0],
+             rhs_contracting_dimensions = [1],
+           >,
+           precision_config = ["DEFAULT", "DEFAULT"]
+         } : (tensor<4x32x1024xf32>, tensor<4x1024x64xf32>) -> tensor<4x32x64xf32>
+  check.expect_almost_eq_const(%res, dense<409.596> : tensor<4x32x64xf32>) : tensor<4x32x64xf32>
+  return
+}