Revert "[inductor][dynamo] Include operator name in size/stride/alignment assertion (pytorch#152353)"

This reverts commit 725bbb6.

Reverted pytorch#152353 on behalf of https://github.com/jeanschmidt due to seems to have broken a few internal tests, @jansel may you help the author get his PR merged? ([comment](pytorch#152353 (comment)))

Author: pytorchmergebot

benchmarks/dynamo/pr_time_benchmarks/expected_results.csv

@@ -10,7 +10,7 @@ add_loop_inductor,compile_time_instruction_count,29370000000,0.015
 
 
 
-add_loop_inductor_dynamic_gpu,compile_time_instruction_count,44130000000,0.025
+add_loop_inductor_dynamic_gpu,compile_time_instruction_count,44480000000,0.025
 
 
 

test/distributed/test_functional_api.py

@@ -715,13 +715,6 @@ def run_with_backward():
 
         _, codes = run_and_get_code(run_with_backward)
         for code in codes:
-            assert_keywords = ["assert_size_stride", "assert_alignment"]
-            filtered_lines = [
-                line
-                for line in code.splitlines()
-                if not any(assert_key in line for assert_key in assert_keywords)
-            ]
-            code = "\n".join(filtered_lines)
             FileCheck().check_count(
                 "_c10d_functional.all_to_all_single.default", 1, exactly=True
             ).check_count("_c10d_functional.wait_tensor.default", 1, exactly=True).run(

test/inductor/test_mkldnn_pattern_matcher.py

@@ -231,14 +231,6 @@ def _test_code_common(
                 torch.compile(mod, fullgraph=True, dynamic=check_dynamic),
                 *clone_inputs,
             )
-            assert_keywords = ["assert_size_stride", "assert_alignment"]
-            filtered_lines = [
-                line
-                for line in source_code.splitlines()
-                if not any(assert_key in line for assert_key in assert_keywords)
-            ]
-            source_code = "\n".join(filtered_lines)
-
             for op in include_ops:
                 self.assertIn(op, source_code)
             if num_include_ops is not None:

test/inductor/test_torchinductor.py

@@ -30,7 +30,6 @@
 import torch._dynamo.config as dynamo_config
 import torch._inductor.aoti_eager
 import torch.nn as nn
-from torch._C._dynamo.guards import assert_alignment, assert_size_stride
 from torch._dispatch.python import enable_python_dispatcher
 from torch._dynamo.debug_utils import aot_graph_input_parser
 from torch._dynamo.device_interface import get_interface_for_device
@@ -1411,10 +1410,9 @@ def fn(a, b):
             )
             _, code = run_and_get_code(fn, x, y)
             code = " ".join(code)
-            if config.cpp_wrapper:
-                self.assertEqual(code.count("view_dtype"), 3)
-            else:
-                self.assertEqual(code.count("aten.view"), 9)
+            self.assertEqual(
+                code.count("view_dtype" if config.cpp_wrapper else "aten.view"), 3
+            )
 
     def test_add_complex5(self):
         def fn(a, b, alpha):
@@ -11884,82 +11882,6 @@ def fn(x):
             check_lowp=False,
         )
 
-    @requires_gpu()
-    @skip_if_not_triton
-    @skip_if_cpp_wrapper("skip cpp_wrapper tests")
-    @config.patch(implicit_fallbacks=True)
-    def test_generated_code_has_size_stride_assert(self):
-        def foo(x):
-            return 3 * x
-
-        def foo_meta(x):
-            return torch.empty_like(x)
-
-        define_custom_op_for_test("foo", foo, foo_meta)
-
-        def fn(x):
-            a = torch.nn.functional.relu(x)
-            b = torch.ops.test.foo(a)
-            return b
-
-        a = torch.randn((16, 32), device=self.device)
-
-        _, code = run_and_get_code(
-            torch.compile(fn),
-            a,
-        )
-        if not is_dynamic_shape_enabled():
-            FileCheck().check(
-                "assert_size_stride(buf2, (16, 32), (32, 1), 'torch.ops.test.foo.default')"
-            ).run(code[0])
-
-    @requires_gpu()
-    @skip_if_not_triton
-    @skip_if_cpp_wrapper("skip cpp_wrapper tests")
-    @config.patch(implicit_fallbacks=True)
-    def test_generated_code_has_alignment_assert(self):
-        def foo(x):
-            return 3 * x
-
-        def foo_meta(x):
-            return torch.empty_like(x)
-
-        define_custom_op_for_test("foo", foo, foo_meta)
-
-        def fn(x):
-            a = torch.nn.functional.relu(x)
-            b = torch.ops.test.foo(a)
-            return b
-
-        a = torch.randn((16, 32), device=self.device)
-
-        _, code = run_and_get_code(
-            torch.compile(fn),
-            a,
-        )
-        if not is_dynamic_shape_enabled():
-            FileCheck().check(
-                "assert_alignment(buf2, 16, 'torch.ops.test.foo.default')"
-            ).run(code[0])
-
-    def test_assert_size_stride_op_name_pass(self):
-        tensor = torch.empty((16, 32))
-        assert_size_stride(tensor, (16, 32), (32, 1), "torch.ops.dummy.op_name")
-
-    def test_assert_size_stride_op_name_fail(self):
-        tensor = torch.empty((16, 32))
-        with self.assertRaisesRegex(AssertionError, "torch.ops.dummy.op_name"):
-            assert_size_stride(tensor, (32, 64), (32, 1), "torch.ops.dummy.op_name")
-
-    def test_assert_alignment_op_name_pass(self):
-        tensor = torch.empty((16, 32))
-        assert_alignment(tensor, 16, "torch.ops.dummy.op_name")
-
-    def test_assert_alignment_op_name_fail(self):
-        tensor = torch.empty((16, 32))
-        with self.assertRaisesRegex(AssertionError, "torch.ops.dummy.op_name"):
-            assert_alignment(tensor, 0, "torch.ops.dummy.op_name")
-
     @torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True)
     @torch._inductor.config.patch(implicit_fallbacks=True)
     def test_custom_op_unbacked_symints(self):
@@ -13092,12 +13014,12 @@ def f(x):
         code = run_and_get_triton_code(f, x)
 
         if is_dynamic_shape_enabled():
-            FileCheck().check("assert_size_stride(buf1, (s77, s27), (s27, 1)").check(
-                "assert_size_stride(buf2, (s77, s27), (s27, 1)"
+            FileCheck().check("assert_size_stride(buf1, (s77, s27), (s27, 1))").check(
+                "assert_size_stride(buf2, (s77, s27), (s27, 1))"
             ).run(code)
         else:
-            FileCheck().check("assert_size_stride(buf1, (16, 32), (32, 1)").check(
-                "assert_size_stride(buf2, (16, 32), (32, 1)"
+            FileCheck().check("assert_size_stride(buf1, (16, 32), (32, 1))").check(
+                "assert_size_stride(buf2, (16, 32), (32, 1))"
             ).run(code)
 
     @requires_cuda

torch/_C/_dynamo/guards.pyi

@@ -176,12 +176,6 @@ def assert_size_stride(
     item: torch.Tensor,
     size: torch.types._size,
     stride: torch.types._size,
-    op_name: str | None = None,
-): ...
-def assert_alignment(
-    item: torch.Tensor,
-    alignment: int,
-    op_name: str | None = None,
 ): ...
 def check_obj_id(obj: object, expected: int) -> bool: ...
 def check_type_id(obj: object, expected: int) -> bool: ...

torch/_inductor/ir.py

@@ -5772,42 +5772,26 @@ def codegen_kwargs(self, skip_out=False):  # type: ignore[no-untyped-def]
             ]
         return kwargs
 
-    def get_op_name(self) -> str:
-        if self.fx_node is not None:
-            target = self.fx_node.target
-            op_namespace = getattr(target, "__module__", "unknown_namespace")
-            op_namespace = op_namespace.replace("._ops.", ".ops.")
-            op_namespace = op_namespace.rsplit(".", 1)[0]
-            op_name = f"{op_namespace}.{target}"
-        else:
-            op_name = "unknown_op"
-        return op_name
-
     def codegen_size_asserts(self, wrapper) -> None:  # type: ignore[no-untyped-def]
         if config.size_asserts and not V.graph.cpp_wrapper:
             # comparing strides for 0 size tensor is tricky. Ignore them for now.
             if sympy_product(self.get_size()) == 0:
                 return
             size = V.graph.wrapper_code.codegen_shape_tuple(self.get_size())
             stride = V.graph.wrapper_code.codegen_shape_tuple(self.get_stride())
-            op_name = self.get_op_name()
+
             wrapper.writeline(
-                f"assert_size_stride({self.get_name()}, {size}, {stride}, {op_name!r})"
+                f"assert_size_stride({self.get_name()}, {size}, {stride})"
             )
 
     def codegen_alignment_asserts(self, wrapper) -> None:  # type: ignore[no-untyped-def]
         if config.alignment_asserts and not V.graph.cpp_wrapper:
             name = self.get_name()
             aligned = name not in V.graph.unaligned_buffers
-            op_name = self.get_op_name()
             if aligned:
-                wrapper.writeline(
-                    f"assert_alignment({name}, {GPU_ALIGN_BYTES}, {op_name!r})"
-                )
+                wrapper.writeline(f"assert_alignment({name}, {GPU_ALIGN_BYTES})")
             else:
-                wrapper.writeline(
-                    f"# buffer {name} (op: {op_name}) is assumed to be not aligned"
-                )
+                wrapper.writeline(f"# buffer {name} is assumed to be not aligned")
 
     def get_group_stride(self):  # type: ignore[no-untyped-def]
         """

torch/csrc/dynamo/guards.cpp

@@ -844,38 +844,21 @@ static PyObject* assert_size_stride(PyObject* dummy, PyObject* args) {
   PyObject* item = nullptr;
   PyObject* size = nullptr;
   PyObject* stride = nullptr;
-  const char* op_name = nullptr;
-
-  if (!PyArg_ParseTuple(args, "OOO|s", &item, &size, &stride, &op_name)) {
+  if (!PyArg_ParseTuple(args, "OOO", &item, &size, &stride)) {
     return nullptr;
   }
   if (!THPVariable_CheckExact(item) && !THPVariable_Check(item)) {
-    std::stringstream msg;
-    msg << "expected Tensor()";
-    if (op_name) {
-      msg << " for op: " << op_name;
-    }
-    PyErr_SetString(PyExc_TypeError, msg.str().c_str());
+    PyErr_SetString(PyExc_TypeError, "expected Tensor()");
     return nullptr;
   }
   if (!PyTuple_CheckExact(size) || !PyTuple_CheckExact(stride)) {
-    std::stringstream msg;
-    msg << "expected tuple()";
-    if (op_name) {
-      msg << " for op: " << op_name;
-    }
-    PyErr_SetString(PyExc_TypeError, msg.str().c_str());
+    PyErr_SetString(PyExc_TypeError, "expected tuple()");
     return nullptr;
   }
   at::Tensor tensor = THPVariable_Unpack(item);
   int64_t ndim = tensor.ndimension();
   if (PyTuple_GET_SIZE(size) != ndim || PyTuple_GET_SIZE(stride) != ndim) {
-    std::stringstream msg;
-    msg << "wrong number of dimensions" << ndim;
-    if (op_name) {
-      msg << " for op: " << op_name;
-    }
-    PyErr_SetString(PyExc_AssertionError, msg.str().c_str());
+    PyErr_SetString(PyExc_AssertionError, "wrong number of dimensions");
     return nullptr;
   }
 
@@ -904,9 +887,6 @@ static PyObject* assert_size_stride(PyObject* dummy, PyObject* args) {
   }
 
   if (num_errors) {
-    if (op_name) {
-      msg << "\nError in op: " << op_name;
-    }
     msg << "\nThis error most often comes from a incorrect fake (aka meta) kernel for a custom op.";
     msg << "\nUse torch.library.opcheck to test your custom op.";
     msg << "\nSee https://pytorch.org/docs/stable/library.html#torch.library.opcheck";
@@ -924,27 +904,15 @@ static PyObject* assert_alignment(PyObject* dummy, PyObject* args) {
    */
   PyObject* item = nullptr;
   unsigned long alignment = 0;
-  const char* op_name = nullptr;
-
-  if (!PyArg_ParseTuple(args, "Ok|s", &item, &alignment, &op_name)) {
+  if (!PyArg_ParseTuple(args, "Ok", &item, &alignment)) {
     return nullptr;
   }
   if (!THPVariable_CheckExact(item) && !THPVariable_Check(item)) {
-    std::stringstream msg;
-    msg << "expected Tensor()";
-    if (op_name) {
-      msg << " for op: " << op_name;
-    }
-    PyErr_SetString(PyExc_TypeError, msg.str().c_str());
+    PyErr_SetString(PyExc_TypeError, "expected Tensor()");
     return nullptr;
   }
   if (alignment == 0) {
-    std::stringstream msg;
-    msg << "alignment cannot be 0";
-    if (op_name) {
-      msg << " in op: " << op_name;
-    }
-    PyErr_SetString(PyExc_AssertionError, msg.str().c_str());
+    PyErr_SetString(PyExc_AssertionError, "alignment can not be 0");
     return nullptr;
   }
 
@@ -954,10 +922,7 @@ static PyObject* assert_alignment(PyObject* dummy, PyObject* args) {
   size_t itemsize = tensor.itemsize();
   if (storage_offset * itemsize % alignment != 0) {
     std::stringstream msg;
-    if (op_name) {
-      msg << "\nError in op: " << op_name;
-    }
-    msg << "\nExpect the tensor to be " << alignment
+    msg << "Expect the tensor to be " << alignment
         << " bytes aligned. Fail due to storage_offset=" << storage_offset
         << " itemsize=" << itemsize;
     PyErr_SetString(PyExc_AssertionError, msg.str().c_str());