Revert "[inductor] Fix issue with scalar arg handling" (pytorch#163737)

This reverts commit a8cd437.

See pytorch#163481 (comment)

This PR might also cause issues with cudagraphs.

Pull Request resolved: pytorch#163737
Approved by: https://github.com/ezyang
ghstack dependencies: pytorch#163386, pytorch#163398, pytorch#163387, pytorch#163414, pytorch#163415, pytorch#163419, pytorch#163434, pytorch#163393, pytorch#163412, pytorch#163422, pytorch#163481, pytorch#163520, pytorch#163482

Author: jansel

test/inductor/test_cpu_repro.py

@@ -485,17 +485,6 @@ def forward(self, x):
             example_inputs = (torch.rand(1, 10),)
             self.common(Model(), example_inputs)
 
-    @torch._dynamo.config.patch(capture_scalar_outputs=True)
-    def test_fill_diagonal_item_scalar_cpu(self):
-        def fn():
-            x = torch.ones(3, 3)
-            x.fill_diagonal_(0)
-            return x.sum().item()
-
-        compiled = torch.compile(fn, backend="inductor", fullgraph=True)
-        eager = fn()
-        self.assertEqual(compiled(), eager)
-
     @unittest.skipIf(not torch.backends.mkldnn.is_available(), "MKLDNN is not enabled")
     @patch("torch.cuda.is_available", lambda: False)
     def test_linear_packed(self):

torch/_inductor/codecache.py

@@ -2864,8 +2864,6 @@ def _worker_compile_cpp(
 # Customized Python binding for cpp kernels
 @clear_on_fresh_cache
 class CppPythonBindingsCodeCache(CppCodeCache):
-    """Compile and cache CPU C++ kernels together with lightweight Python bindings."""
-
     cache: dict[str, Callable[[], Union[CDLL, ModuleType]]] = {}
     cache_clear = staticmethod(cache.clear)
     cpp_compile_command_flags = {
@@ -2875,28 +2873,7 @@ class CppPythonBindingsCodeCache(CppCodeCache):
     }
     entry_function = "kernel"
     call_entry_function = "kernel({}); Py_RETURN_NONE;"
-    extra_parse_arg = textwrap.dedent(
-        """
-        template <> inline double parse_arg<double>(PyObject* args, size_t n) {{
-            auto result = PyFloat_AsDouble(PyTuple_GET_ITEM(args, n));
-            if(unlikely(result == -1.0 && PyErr_Occurred()))
-                throw std::runtime_error("expected float arg");
-            return result;
-        }}
-        template <> inline float parse_arg<float>(PyObject* args, size_t n) {{
-            auto result = PyFloat_AsDouble(PyTuple_GET_ITEM(args, n));
-            if(unlikely(result == -1.0 && PyErr_Occurred()))
-                throw std::runtime_error("expected float arg");
-            return static_cast<float>(result);
-        }}
-        template <> inline bool parse_arg<bool>(PyObject* args, size_t n) {{
-            int result = PyObject_IsTrue(PyTuple_GET_ITEM(args, n));
-            if(unlikely(result == -1 && PyErr_Occurred()))
-                throw std::runtime_error("expected bool arg");
-            return result;
-        }}
-        """
-    )
+    extra_parse_arg = ""
     suffix_template = textwrap.dedent(
         """
         // Python bindings to call {entry_func}():

torch/_inductor/codegen/common.py

@@ -290,12 +290,6 @@ def alias_of(self) -> Optional[str]:
         return None
 
 
-@dataclasses.dataclass
-class ScalarArg:
-    name: str
-    dtype: torch.dtype
-
-
 @dataclasses.dataclass
 class ConstexprArg:
     name: str
@@ -317,14 +311,7 @@ class DeviceCodegen:
     fx_wrapper_codegen: Optional[WrapperConstructor] = None
 
 
-KernelArgType = Union[
-    WorkspaceArg,
-    TensorArg,
-    SizeArg,
-    ScalarArg,
-    TMADescriptorArg,
-    ConstexprArg,
-]
+KernelArgType = Union[WorkspaceArg, TensorArg, SizeArg, TMADescriptorArg, ConstexprArg]
 
 device_codegens: dict[str, DeviceCodegen] = {}
 
@@ -1480,7 +1467,6 @@ def __init__(self) -> None:
         self.output_buffers: dict[str, Union[str, RemovedArg]] = {}
         self.inplace_buffers: dict[str, Union[InplacedBuffer, RemovedArg]] = {}
         self.sizevars: dict[sympy.Expr, str] = {}
-        self.scalar_vars: dict[sympy.Symbol, tuple[str, torch.dtype]] = {}
         self.workspace_args: list[WorkspaceArg] = []
 
     def __repr__(self) -> str:
@@ -1493,7 +1479,6 @@ def __repr__(self) -> str:
                         self.output_buffers,
                         self.inplace_buffers,
                         self.sizevars,
-                        self.scalar_vars,
                     ],
                 )
             )
@@ -1641,27 +1626,9 @@ def size(self, name: sympy.Symbol) -> str:
             return "seed"
         return self._lookup("ks", self.sizevars, name)
 
-    def scalar(self, name: sympy.Symbol, dtype: torch.dtype) -> str:
-        assert isinstance(name, sympy.Symbol), (type(name), name)
-        if name in self.scalar_vars:
-            inner, existing_dtype = self.scalar_vars[name]
-            if existing_dtype != dtype:
-                try:
-                    promoted = torch.promote_types(existing_dtype, dtype)
-                except TypeError:
-                    promoted = dtype
-                self.scalar_vars[name] = (inner, promoted)
-            return self.scalar_vars[name][0]
-        inner = f"kscalar{len(self.scalar_vars)}"
-        self.scalar_vars[name] = (inner, dtype)
-        return inner
-
     def call_names(self) -> Iterator[str]:
         return chain(
-            self.input_buffers.keys(),
-            self.output_buffers.keys(),
-            self.sizevars.keys(),
-            self.scalar_vars.keys(),
+            self.input_buffers.keys(), self.output_buffers.keys(), self.sizevars.keys()
         )
 
     def arg_name(self, name: str) -> Optional[str]:
@@ -1682,9 +1649,6 @@ def wrap_ptr_arg(self, buf: str, dtype: torch.dtype) -> str:
     def wrap_size_arg(self, size: SymbolLike) -> str:
         return str(size)
 
-    def wrap_scalar_arg(self, scalar: sympy.Symbol) -> str:
-        return str(scalar)
-
     def cpp_argdefs(
         self, dtype_to_cpp_type: Optional[dict[torch.dtype, str]] = None
     ) -> tuple[list[str], list[str], list[str]]:
@@ -1730,11 +1694,6 @@ def cpp_argdefs(
             arg_types.append(f"const {INDEX_TYPE}")
             if V.graph.wrapper_code:
                 V.graph.wrapper_code.ensure_size_computed(outer)
-        for outer, (inner, dtype) in self.scalar_vars.items():
-            cpp_dtype = dtype_to_cpp_type[dtype]
-            arg_defs.append(f"const {cpp_dtype} {inner}")
-            call_args.append(self.wrap_scalar_arg(outer))
-            arg_types.append(f"const {cpp_dtype}")
         assert not self.workspace_args, "Workspace not supported on CPU "
         return arg_defs, call_args, arg_types
 
@@ -1780,11 +1739,6 @@ def python_argdefs(
             precompile_args.append(SizeArg(inner, outer))
             if V.graph.wrapper_code:
                 V.graph.wrapper_code.ensure_size_computed(outer)
-        for outer, (inner, dtype) in self.scalar_vars.items():
-            arg_defs.append(ArgName(inner))
-            call_args.append(self.wrap_scalar_arg(outer))
-            arg_types.append(dtype)
-            precompile_args.append(ScalarArg(inner, dtype))
         for arg in self.workspace_args:
             arg_defs.append(ArgName(arg.inner_name))
             call_args.append(arg.outer_name)
@@ -2339,10 +2293,6 @@ def rename_indexing(
                 ),
             )
         }
-        for x in sorted_symbols:
-            if symbol_is_type(x, (SymT.FLOAT, SymT.UNBACKED_FLOAT)):
-                dtype = V.graph.get_dynamic_scalar_dtype(x)
-                replacements[x] = self.args.scalar(x, dtype)
         return sympy_subs(index, replacements)
 
     def create_cse_var(self, *args: Any, **kwargs: Any) -> CSEVariable:

torch/_inductor/codegen/triton_utils.py

@@ -13,7 +13,6 @@
     ArgName,
     ConstexprArg,
     KernelArgType,
-    ScalarArg,
     SizeArg,
     TensorArg,
     TMADescriptorArg,
@@ -92,9 +91,6 @@ def signature_of(arg: KernelArgType, *, size_dtype: Optional[str]) -> str:
             raise NotImplementedError(f"unhandled size_dtype {size_dtype}")
     if isinstance(arg, WorkspaceArg):
         return _type_of(arg.dtype)
-    if isinstance(arg, ScalarArg):
-        typ = _type_of(arg.dtype)
-        return typ.removeprefix("*")
     if isinstance(arg, TMADescriptorArg):
         if arg.api_type == "experimental":
             return "nvTmaDesc"

torch/_inductor/graph.py

@@ -385,7 +385,6 @@ def __init__(
             const_module.device_idxs if const_module else OrderedSet()
         )
         self.device_type = "cpu"
-        self.dynamic_scalar_dtypes: dict[sympy.Symbol, torch.dtype] = {}
 
         # Inplace padding may require Inductor to allocate slightly larger
         # tensor for padding.
@@ -950,17 +949,6 @@ def get_dtype(self, buffer_name: str) -> torch.dtype:
             return self.get_dtype(m.group(1))
         raise KeyError(f"could not find {buffer_name}")
 
-    def register_dynamic_scalar_dtype(
-        self, sym: sympy.Symbol, dtype: torch.dtype
-    ) -> None:
-        existing = self.dynamic_scalar_dtypes.get(sym)
-        if existing is not None and existing != dtype:
-            dtype = torch.promote_types(existing, dtype)
-        self.dynamic_scalar_dtypes[sym] = dtype
-
-    def get_dynamic_scalar_dtype(self, sym: sympy.Symbol) -> torch.dtype:
-        return self.dynamic_scalar_dtypes.get(sym, torch.float64)
-
     def get_numel(self, buffer_name: str) -> Union[int, Expr]:
         if buffer_name in self.constants:
             return self.constants[buffer_name].numel()

torch/_inductor/lowering.py

@@ -3236,7 +3236,6 @@ def _local_scalar_dense(data):
     buffer = ir.DynamicScalar(binding_sym, keypath, data)
     buffer.name = V.graph.register_buffer(buffer)
     V.graph.register_operation(buffer)
-    V.graph.register_dynamic_scalar_dtype(binding_sym, data.get_dtype())
     # NB: the replaced expr is OK to use directly downstream, we want
     # simplifications in this case!
     val = V.graph.current_node.meta["val"]