From 0bb358ee70c9fd737f169cfff8d4fca21434f0a7 Mon Sep 17 00:00:00 2001
From: Alexander Belyaev <pifon@google.com>
Date: Mon, 25 Jan 2021 11:03:00 -0800
Subject: [PATCH] [KERNEL_GEN] Disable loop unrolling for binary kernels.

The unrolling is performed, but it does not lead to any vector instructions.
We disable it for now.

PiperOrigin-RevId: 353684863
Change-Id: I394d86e04e48edaecde8f38fd0208d0d4d1e9f26
---
 tensorflow/core/kernels/mlir_generated/BUILD | 55 ++++++++++++--------
 1 file changed, 33 insertions(+), 22 deletions(-)

diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD
index d0344bcad1e..6ed9f174587 100644
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@@ -508,14 +508,15 @@ gen_kernel_library(
 [
     gen_kernel_library(
         name = name,
-        tile_size = "256,1,1",
+        tile_size = "1024",
         types = [
             "f16",
             "f32",
             "f64",
             "i64",
         ],
-        unroll_factors = "4",
+        # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+        # unroll_factors = "4",
     )
     for name in [
         "square",
@@ -526,12 +527,13 @@ gen_kernel_library(
 
 gen_kernel_library(
     name = "complex",
-    tile_size = "256,1,1",
+    tile_size = "1024",
     types = [
         "f32",
         "f64",
     ],
-    unroll_factors = "2",
+    # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+    # unroll_factors = "2",
 )
 
 gen_kernel_library(
@@ -545,7 +547,7 @@ gen_kernel_library(
 
 gen_kernel_library(
     name = "div",
-    tile_size = "256,1,1",
+    tile_size = "1024",
     types = [
         "f16",
         "f32",
@@ -553,12 +555,13 @@ gen_kernel_library(
         "i16",
         "i64",
     ],
-    unroll_factors = "4",
+    # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+    # unroll_factors = "4",
 )
 
 gen_kernel_library(
     name = "mul",
-    tile_size = "256,1,1",
+    tile_size = "1024",
     types = [
         "f16",
         "f32",
@@ -567,14 +570,15 @@ gen_kernel_library(
         "i16",
         "i64",
     ],
-    unroll_factors = "4",
+    # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+    # unroll_factors = "4",
 )
 
 # Bitwise operations.
 [
     gen_kernel_library(
         name = name,
-        tile_size = "256,1,1",
+        tile_size = "1024",
         types = [
             "i8",
             "i16",
@@ -586,7 +590,8 @@ gen_kernel_library(
             # "ui32",
             # "ui64",
         ],
-        unroll_factors = "4",
+        # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+        # unroll_factors = "4",
     )
     for name in [
         "bitwise_and",
@@ -611,11 +616,12 @@ gen_kernel_library(
 [
     gen_kernel_library(
         name = name,
-        tile_size = "256,1,1",
+        tile_size = "1024",
         types = [
             "i1",
         ],
-        unroll_factors = "4",
+        # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+        # unroll_factors = "4",
     )
     for name in [
         "logical_and",
@@ -626,7 +632,7 @@ gen_kernel_library(
 [
     gen_kernel_library(
         name = name,
-        tile_size = "256,1,1",
+        tile_size = "1024",
         types = [
             "f16",
             "f32",
@@ -637,7 +643,8 @@ gen_kernel_library(
             "i32",
             "i64",
         ],
-        unroll_factors = "4",
+        # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+        # unroll_factors = "4",
     )
     for name in [
         "equal",
@@ -648,7 +655,7 @@ gen_kernel_library(
 [
     gen_kernel_library(
         name = name,
-        tile_size = "256,1,1",
+        tile_size = "1024",
         types = [
             "f16",
             "f32",
@@ -658,7 +665,8 @@ gen_kernel_library(
             "i32",
             "i64",
         ],
-        unroll_factors = "4",
+        # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+        # unroll_factors = "4",
     )
     for name in [
         "less",
@@ -671,7 +679,7 @@ gen_kernel_library(
 [
     gen_kernel_library(
         name = name,
-        tile_size = "256,1,1",
+        tile_size = "1024",
         types = [
             "f16",
             "f32",
@@ -680,7 +688,8 @@ gen_kernel_library(
             "i32",
             "i64",
         ],
-        unroll_factors = "4",
+        # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+        # unroll_factors = "4",
     )
     for name in [
         "maximum",
@@ -711,7 +720,7 @@ gen_kernel_library(
 
 gen_kernel_library(
     name = "floor_div",
-    tile_size = "256",
+    tile_size = "1024",
     # TODO(172804967): Enable for integer types also once unsigned integers are
     # supported.
     types = [
@@ -719,7 +728,8 @@ gen_kernel_library(
         "f32",
         "f64",
     ],
-    unroll_factors = "4",
+    # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+    # unroll_factors = "4",
 )
 
 # Kernels that support all floating-point types.
@@ -770,12 +780,13 @@ gen_kernel_library(
 
 gen_kernel_library(
     name = "pow",
-    tile_size = "256,1,1",
+    tile_size = "1024",
     types = [
         "f16",
         "f32",
         "f64",
         "i64",
     ],
-    unroll_factors = "4",
+    # TODO(b/178388085): Enable unrolling after vectorization is fixed.
+    # unroll_factors = "4",
 )