[Inductor] Kill MutationLayout from ir.py (#112925)

oulgen · pytorchmergebot · commit 611a7457cad5 · 2023-11-07T17:03:52.000Z
Pull Request resolved: #112925 Approved by: https://github.com/jansel
diff --git a/test/distributed/test_inductor_collectives.py b/test/distributed/test_inductor_collectives.py
@@ -528,9 +528,9 @@ def func(inp, *, tag, ranks, group_size):
             .check("buf2_work = dist.all_reduce(buf2") \
             .check("fun_col_impl._register_tensor_work(buf2, buf2_work)") \
             .check("buf1 = _wait_tensor(buf1)") \
-            .check("buf3 = buf1") \
-            .check("buf4 = empty") \
-            .check("return (buf1, buf4") \
+            .check("buf4 = buf1") \
+            .check("buf5 = empty") \
+            .check("return (buf1, buf5") \
             .run(code)
         out = compiled(inputs, **self.get_world_trs())
         correct = func(inputs, **self.get_world_trs())
@@ -561,16 +561,16 @@ def func(inp, *, tag, ranks, group_size):
         # wait_tensors before they are returned from the graph.
         FileCheck() \
             .check("buf0 = empty(") \
-            .check("buf4 = empty(") \
-            .check("triton_poi__0.run(arg0_1, buf0, buf4") \
+            .check("buf5 = empty(") \
+            .check("triton_poi__0.run(arg0_1, buf0, buf5") \
             .check_not("copy_(") \
             .check("buf1 = buf0; del buf0  # reuse") \
             .check("buf2 = buf1") \
             .check("buf2_work = dist.all_reduce(buf2") \
             .check("fun_col_impl._register_tensor_work(buf2, buf2_work)") \
             .check("buf1 = _wait_tensor(buf1)") \
-            .check("buf3 = buf1") \
-            .check("return (buf1, buf4, buf5") \
+            .check("buf4 = buf1") \
+            .check("return (buf1, buf5, buf6") \
             .run(code)
         out = compiled(inputs, **self.get_world_trs())
         correct = func(inputs, **self.get_world_trs())
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -2471,6 +2471,7 @@ class NoneLayout(IRNode):
     def __init__(self, device):
         self.device = device
         self.size = [0]
+        self.stride = [0]
 
     def storage_size(self):
         return 0
@@ -3843,7 +3844,6 @@ def __init__(self, *, kernel_idx, grid, kernel_args):
                 t = InputsKernel.unwrap_storage_for_input(self.realize_input(v))
                 inputs.append(t)
                 kwargs[k] = t
-                V.graph.mark_buffer_mutated(t.get_name())
             else:
                 constant_args.append(v)
                 kwargs[k] = v
@@ -3861,20 +3861,25 @@ def __init__(self, *, kernel_idx, grid, kernel_args):
         self.name = V.graph.register_buffer(self)
         self.kernel_idx = kernel_idx
         self.grid = grid
-
-    @classmethod
-    def create(cls, *, kernel_idx, grid, kernel_args):
-        packed = UserDefinedTritonKernel(
-            kernel_idx=kernel_idx, grid=grid, kernel_args=kernel_args
+        mark_node_as_mutating(
+            self, *[a for a in kernel_args.values() if isinstance(a, TensorBox)]
         )
-        for arg in kernel_args.values():
-            if isinstance(arg, TensorBox):
-                MutationOutput(arg.layout, arg, packed)
 
     def get_alias_names(self):
         return [i.get_name() for i in self.inputs]
 
 
+def mark_node_as_mutating(cur_buffer, *mutated_ops):
+    """
+    Allows ops in mutated_ops to be marked as being mutated as well as
+    indicates to the scheduler that these ops depend on cur_buffer.
+    """
+    for op in mutated_ops:
+        assert isinstance(op, TensorBox)
+        V.graph.mark_buffer_mutated(op.get_name())
+        MutationOutput(op.layout, op, cur_buffer)
+
+
 class MutationOutput(ExternKernel):
     def get_mutation_names(self):
         return [self.inputs[0].get_name()]
@@ -3913,17 +3918,20 @@ def should_allocate(self):
         return False
 
     def get_mutation_names(self):
-        assert isinstance(self.layout, MutationLayout)
-        return (self.layout.target.get_name(),)
+        return [self.inputs[0].get_name()]
+
+    def get_unbacked_symbol_defs(self):
+        return {}
 
     def __init__(self, x, *constant_args):
         super().__init__(
             None,
-            MutationLayout(x),
+            NoneLayout(x.get_device()),  # type: ignore[arg-type]
             self.unwrap_storage([x]),
             constant_args,
         )
         self.name = V.graph.register_buffer(self)
+        mark_node_as_mutating(self, x)
 
 
 class AccumulateGrad(ExternKernel):
@@ -3941,16 +3949,19 @@ def should_allocate(self):
         return False
 
     def get_mutation_names(self):
-        assert isinstance(self.layout, MutationLayout)
-        return (self.layout.target.get_name(),)
+        return [self.inputs[0].get_name()]
+
+    def get_unbacked_symbol_defs(self):
+        return {}
 
     def __init__(self, variable, new_grad):
         super().__init__(
             None,
-            MutationLayout(variable),
+            NoneLayout(variable.get_device()),  # type: ignore[arg-type]
             self.unwrap_storage([variable, new_grad]),
         )
         self.name = V.graph.register_buffer(self)
+        mark_node_as_mutating(self, variable)
 
 
 class ScatterFallback(ExternKernel):
@@ -3997,6 +4008,12 @@ def get_cpp_kernel(self, fn, reduce):
             kernel = "at::scatter_reduce_out"
         return kernel
 
+    def get_mutation_names(self):
+        return [self.inputs[0].get_name()]
+
+    def get_unbacked_symbol_defs(self):
+        return {}
+
     def __init__(
         self,
         fn,
@@ -4028,15 +4045,17 @@ def __init__(
         else:
             tensors = [self.realize_input(t) for t in [x, index]]
             constant_args = (dim, src)
+
         super().__init__(
             None,
-            MutationLayout(x),
+            NoneLayout(x.get_device()),  # type: ignore[arg-type]
             self.unwrap_storage(tensors),
             constant_args,
             {"reduce": reduce, "include_self": include_self},
         )
         self.ordered_kwargs_for_cpp_kernel = ["reduce", "include_self"]
         self.name = V.graph.register_buffer(self)
+        mark_node_as_mutating(self, x)
 
 
 class IndexPutFallback(ExternKernel):
@@ -4061,18 +4080,25 @@ def codegen(self, wrapper):
     def should_allocate(self):
         return False
 
+    def get_mutation_names(self):
+        return [self.inputs[0].get_name()]
+
+    def get_unbacked_symbol_defs(self):
+        return {}
+
     def __init__(self, x, indices, values, accumulate):
         self.indices = indices
         valid_indices = [i for i in indices if i is not None]
         tensors = [self.realize_input(x) for x in [x, values, *valid_indices]]
         super().__init__(
             None,
-            MutationLayout(x),
+            NoneLayout(x.get_device()),  # type: ignore[arg-type]
             self.unwrap_storage(tensors),
             (accumulate,),
         )
         self.name = V.graph.register_buffer(self)
         self.kernel = "at::index_put_" if V.graph.cpp_wrapper else "aten.index_put_"
+        mark_node_as_mutating(self, x)
 
 
 class DeviceCopy(ExternKernelOut):
@@ -5019,8 +5045,10 @@ def codegen(self, wrapper):
         )
 
     def get_mutation_names(self):
-        assert isinstance(self.layout, MutationLayout)
-        return (self.layout.target.get_name(),)
+        return [self.inputs[1].get_name()]
+
+    def get_unbacked_symbol_defs(self):
+        return {}
 
     @classmethod
     def create(
@@ -5056,11 +5084,13 @@ def create(
             may_convert_to_optional(unary_scalars),
             unary_algorithm,
         ]
-        return ConvolutionBinaryInplace(
-            kernel_layout=MutationLayout(inputs[1]),
+        packed = ConvolutionBinaryInplace(
+            kernel_layout=NoneLayout(inputs[1].get_device()),  # type: ignore[arg-type]
             inputs=inputs,
             constant_args=constant_args,
         )
+        mark_node_as_mutating(packed, inputs[1])
+        return packed
 
 
 class MKLPackedLinear(ExternKernelAlloc):
@@ -6638,6 +6668,9 @@ def should_allocate(self):
     def get_mutation_names(self):
         return [self.inputs[0].get_name()]
 
+    def get_unbacked_symbol_defs(self):
+        return {}
+
     @classmethod
     def create(
         cls,
@@ -6648,14 +6681,13 @@ def create(
         group_size: int,
     ):
         inplace_inputs = cls.wrap_inputs_as_inplace(inputs)
-        layout = MutationLayout(inplace_inputs[0])
-
-        _ = AllReduceCoalesced(
-            layout=layout,
+        packed = AllReduceCoalesced(
+            layout=NoneLayout(inplace_inputs[0].get_device()),  # type: ignore[arg-type]
             inputs=inplace_inputs,
             constant_args=[tag, ranks, group_size],
             reduce_op=reduce_op,
         )
+        mark_node_as_mutating(packed, inplace_inputs[0])
         return inplace_inputs
 
     def codegen_collective(self, wrapper, output_name, input_names):
@@ -6676,19 +6708,22 @@ def __init__(self, layout, inputs, constant_args, reduce_op):
     def get_mutation_names(self):
         return [self.inputs[0].get_name()]
 
+    def get_unbacked_symbol_defs(self):
+        return {}
+
     @classmethod
     def create(
         cls, x: "TensorBox", reduce_op: str, tag: str, ranks: List[int], group_size: int
     ):
         inplace_inputs = cls.wrap_inputs_as_inplace([x])
-        layout = MutationLayout(inplace_inputs[0])
 
-        _ = AllReduce(
-            layout=layout,
+        packed = AllReduce(
+            layout=NoneLayout(inplace_inputs[0].get_device()),  # type: ignore[arg-type]
             inputs=inplace_inputs,
             constant_args=[tag, ranks, group_size],
             reduce_op=reduce_op,
         )
+        mark_node_as_mutating(packed, inplace_inputs[0])
         return inplace_inputs[0]
 
     def codegen_collective(self, wrapper, output_name, input_names):
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -5067,9 +5067,7 @@ def accumulate_grad_(variable, new_grad):
 
 @register_lowering(triton_kernel_wrapper_mutation)
 def triton_kernel_wrap_(*, kernel_idx, grid, kwargs):
-    ir.UserDefinedTritonKernel.create(
-        kernel_idx=kernel_idx, grid=grid, kernel_args=kwargs
-    )
+    ir.UserDefinedTritonKernel(kernel_idx=kernel_idx, grid=grid, kernel_args=kwargs)
     return {key: val for key, val in kwargs.items() if isinstance(val, TensorBox)}
 
 
