Revert "[Inductor] Constrain the shape of other tensor for Conv/Linear + broadcast add fusion. (pytorch#141759)"

This reverts commit 35752cb.

Reverted pytorch#141759 on behalf of https://github.com/atalman due to Failing internally ([comment](pytorch#141759 (comment)))

Author: pytorchmergebot

aten/src/ATen/native/mkldnn/Linear.cpp

@@ -295,15 +295,12 @@ Tensor mkldnn_linear_pointwise_binary(
         input_reshaped.size(0), weight_t.size(0)};
     output = output.reshape(output_size_reshaped);
     other_reshaped = other_reshaped.reshape(output_size_reshaped);
-    TORCH_CHECK(
-        output.sizes() == other_reshaped.sizes(),
-        "linear_binary_run expects the size of output and other tensor to be the same");
-  } else {
-    TORCH_CHECK(
-        output.dim() == other_reshaped.dim(),
-        "linear_binary_run expects the dimension of output and other tensor to be the same");
   }
 
+  TORCH_CHECK(
+      output.dim() == other_reshaped.dim(),
+      "linear_binary_run expects the dimension of output and other tensor to be the same");
+
   c10::impl::ExcludeDispatchKeyGuard edkg(c10::autograd_dispatch_keyset);
   ideep::tensor mkldnn_output = itensor_from_tensor(output);
   const ideep::tensor mkldnn_other = itensor_from_tensor(other_reshaped);

test/inductor/test_mkldnn_pattern_matcher.py

@@ -709,34 +709,26 @@ def forward(self, x, x2):
             dtypes.append(torch.float16)
         cl_format = torch.channels_last if dim == 4 else torch.channels_last_3d
         test_memory_format = [torch.contiguous_format, cl_format]
-        if dim == 4:
-            input_shapes = [
-                [2, 3, 56, 56],
-            ]
-            other_shapes = [[2, 16, 1, 1], [1, 16, 1, 1], [1, 1, 1, 1]]
-        else:
-            input_shapes = [
-                [2, 3, 20, 56, 56],
-            ]
-            other_shapes = [[2, 16, 1, 1, 1], [1, 16, 1, 1, 1], [1, 1, 1, 1, 1]]
         options = itertools.product(
             binary_list,
-            input_shapes,
-            other_shapes,
             [True, False],
             test_memory_format,
             dtypes,
         )
 
         for (
             binary_fn,
-            x_shape,
-            other_shape,
             has_relu,
             memory_format,
             dtype,
         ) in options:
             metrics.reset()
+            if dim == 4:
+                x_shape = (1, 3, 56, 56)
+                other_shape = (1, 16, 1, 1)
+            else:
+                x_shape = (1, 3, 20, 56, 56)
+                other_shape = (1, 16, 1, 1, 1)
             mod = M(binary_fn, has_relu).eval()
             x = (
                 torch.randn(x_shape, dtype=torch.float32, requires_grad=True)
@@ -855,23 +847,15 @@ def forward(self, x, y):
         if torch.ops.mkldnn._is_mkldnn_fp16_supported():
             dtypes.append(torch.float16)
         options = itertools.product(
-            binary_list,
-            (
-                ([2, 3, 10], [1, 1, 30]),
-                ([2, 3, 10], [1, 1, 1]),
-                ([2, 10], [1, 30]),
-                ([2, 10], [1, 1]),
-            ),
-            (True, False),
-            dtypes,
+            binary_list, [[2, 3, 10], [2, 10]], [True, False], dtypes
         )
         out_feature = 30
 
-        for binary_fn, (input_shape, other_shape), bias, dtype in options:
+        for binary_fn, input_shape, bias, dtype in options:
             metrics.reset()
             mod = M(binary_fn, input_shape[-1], out_feature, bias).eval()
             v = torch.randn(input_shape)
-            other = torch.randn(other_shape).to(dtype)
+            other = torch.randn(input_shape[:-1] + [1]).to(dtype)
 
             def matcher_check_fn():
                 self.assertEqual(

torch/_inductor/fx_passes/mkldnn_fusion.py

@@ -356,8 +356,8 @@ def fn(match, *args, **kwargs):
         ops.sub: "sub",
     }
 
-    def _is_valid_binary(match, computation_op, binary_op):
-        binary_nodes = filter_nodes(match.nodes, binary_op)
+    def _is_valid_binary(match, fn):
+        binary_nodes = filter_nodes(match.nodes, fn)
         if len(binary_nodes) < 1:
             return False
 
@@ -381,50 +381,13 @@ def get_meta_value(argument: torch.fx.node.Argument):
         ):
             return False
 
-        def _check_input_sizes(n, computation_op):
-            # Check if the tensor shape of the 'other' node is the same as or
-            # can be broadcasted to the tensor shape of the computation node.
-            computation_node = (
-                n.args[0] if n.args[1] is match.kwargs["other"] else n.args[1]
-            )
-            assert computation_node.target == computation_op
-            computation_node_size = get_meta_value(computation_node).size()
-            if computation_op is mkldnn._linear_pointwise.default:
-                if len(computation_node_size) >= 2:
-                    broadcast_sizes = [
-                        torch.Size(
-                            [1 for _ in range(len(computation_node_size) - 1)]
-                            + [computation_node_size[-1]]
-                        ),
-                        torch.Size([1 for _ in range(len(computation_node_size))]),
-                    ]
-                else:
-                    broadcast_sizes = [
-                        torch.Size([1 for _ in range(len(computation_node_size))]),
-                    ]
-            else:
-                assert len(computation_node_size) > 2
-                broadcast_sizes = [
-                    torch.Size(
-                        [computation_node_size[0], computation_node_size[1]]
-                        + [1 for _ in range(len(computation_node_size) - 2)]
-                    ),
-                    torch.Size(
-                        [1, computation_node_size[1]]
-                        + [1 for _ in range(len(computation_node_size) - 2)]
-                    ),
-                    torch.Size([1 for _ in range(len(computation_node_size))]),
-                ]
-            return (
-                get_meta_value(match.kwargs["other"]).size()
-                in [
-                    computation_node_size,
-                ]
-                + broadcast_sizes
-            )
-
         if any(
-            not _check_input_sizes(n, computation_op)
+            get_meta_value(n.args[0]).dim() != get_meta_value(n.args[1]).dim()
+            or not all(
+                get_meta_value(n.args[0]).size(i) == get_meta_value(n.args[1]).size(i)
+                or get_meta_value(match.kwargs["other"]).size(i) == 1
+                for i in range(get_meta_value(n.args[0]).dim())
+            )
             or get_meta_value(n.args[0]).device != get_meta_value(n.args[1]).device
             or get_meta_value(n.args[0]).dtype != get_meta_value(n.args[1]).dtype
             for n in binary_nodes
@@ -439,7 +402,7 @@ def _is_valid_computation_binary(computation_op, binary_op, other_index=None):
         def fn(match):
             if not _is_single_computation_op(computation_op)(match):
                 return False
-            if not _is_valid_binary(match, computation_op, binary_op):
+            if not _is_valid_binary(match, binary_op):
                 return False
             return True
 
@@ -562,7 +525,7 @@ def _can_be_inplace(_other):
         else:
             return not (
                 isinstance(_other.data, ir.ReinterpretView)
-                or len(_other.data.get_inputs_that_alias_output()) > 0
+                or len(_other.get_inputs_that_alias_output()) > 0
             )
 
     def _register_binary_unary_maybe_inplace_fusion_lowering(