Fix inplacing with multiple, fused uses (pytorch#150892)

Fix inplacing with multiple, fused uses (pytorch#150845)

We had `can_inplace` defined on a single use. When that buffer has multiple uses inside a fused node, we need to check if the other accesses have the same index. Otherwise we may read memory that has already been written to from inplacing.

Pull Request resolved: pytorch#150845
Approved by: https://github.com/zou3519, https://github.com/exclamaforte, https://github.com/atalman, https://github.com/jansel

(cherry picked from commit 27ded35)

Co-authored-by: eellison <[email protected]>

Author: pytorchbot

test/inductor/test_cuda_repro.py

@@ -1292,6 +1292,185 @@ def fn(x, y, z):
 
         self.assertEqual(ref, res)
 
+    @torch._inductor.config.patch(emulate_precision_casts=True)
+    def test_dont_inplace_disjoint_accesses(self):
+        # TODO - would not need mms if we could annotate donated buffer..
+        def forward(  # noqa: F821, F722
+            arg0_1: "bf16[2048, 2048][2048, 1]cuda:0",  # noqa: F821, F722
+            arg1_1: "bf16[8, 4096, 2048][8388608, 2048, 1]cuda:0",  # noqa: F821, F722
+            arg2_1: "bf16[2048, 2048][2048, 1]cuda:0",  # noqa: F821, F722
+            arg3_1: "bf16[2048, 2048][2048, 1]cuda:0",  # noqa: F821, F722
+            arg4_1: "bf16[2048][1]cuda:0",  # noqa: F821, F722
+            arg5_1: "bf16[2048][1]cuda:0",  # noqa: F821, F722
+            arg6_1: "f32[4096, 128][128, 1]cuda:0",  # noqa: F821, F722
+            arg7_1: "f32[4096, 128][128, 1]cuda:0",  # noqa: F821, F722
+        ):
+            permute = torch.ops.aten.permute.default(arg0_1, [1, 0])
+            arg0_1 = None
+            view = torch.ops.aten.view.default(arg1_1, [32768, 2048])
+            mm = torch.ops.aten.mm.default(view, permute)
+            view = permute = None
+            view_1 = torch.ops.aten.view.default(mm, [8, 4096, 2048])
+            mm = None
+            permute_1 = torch.ops.aten.permute.default(arg2_1, [1, 0])
+            arg2_1 = None
+            view_2 = torch.ops.aten.view.default(arg1_1, [32768, 2048])
+            mm_1 = torch.ops.aten.mm.default(view_2, permute_1)
+            view_2 = permute_1 = None
+            view_3 = torch.ops.aten.view.default(mm_1, [8, 4096, 2048])
+            mm_1 = None
+            permute_2 = torch.ops.aten.permute.default(arg3_1, [1, 0])
+            arg3_1 = None
+            view_4 = torch.ops.aten.view.default(arg1_1, [32768, 2048])
+            arg1_1 = None
+            mm_2 = torch.ops.aten.mm.default(view_4, permute_2)
+            view_4 = permute_2 = None
+            view_5 = torch.ops.aten.view.default(mm_2, [8, 4096, 2048])
+            mm_2 = None
+            convert_element_type_6 = torch.ops.prims.convert_element_type.default(
+                view_1, torch.float32
+            )
+            view_1 = None
+            pow_1 = torch.ops.aten.pow.Tensor_Scalar(convert_element_type_6, 2)
+            mean = torch.ops.aten.mean.dim(pow_1, [-1], True)
+            pow_1 = None
+            add = torch.ops.aten.add.Tensor(mean, 1e-06)
+            mean = None
+            rsqrt = torch.ops.aten.rsqrt.default(add)
+            add = None
+            mul = torch.ops.aten.mul.Tensor(convert_element_type_6, rsqrt)
+            convert_element_type_6 = rsqrt = None
+            convert_element_type_7 = torch.ops.prims.convert_element_type.default(
+                arg4_1, torch.float32
+            )
+            arg4_1 = None
+            mul_1 = torch.ops.aten.mul.Tensor(convert_element_type_7, mul)
+            convert_element_type_7 = mul = None
+            convert_element_type_8 = torch.ops.prims.convert_element_type.default(
+                mul_1, torch.bfloat16
+            )
+            mul_1 = None
+            convert_element_type_9 = torch.ops.prims.convert_element_type.default(
+                view_3, torch.float32
+            )
+            view_3 = None
+            pow_2 = torch.ops.aten.pow.Tensor_Scalar(convert_element_type_9, 2)
+            mean_1 = torch.ops.aten.mean.dim(pow_2, [-1], True)
+            pow_2 = None
+            add_1 = torch.ops.aten.add.Tensor(mean_1, 1e-06)
+            mean_1 = None
+            rsqrt_1 = torch.ops.aten.rsqrt.default(add_1)
+            add_1 = None
+            mul_2 = torch.ops.aten.mul.Tensor(convert_element_type_9, rsqrt_1)
+            convert_element_type_9 = rsqrt_1 = None
+            convert_element_type_10 = torch.ops.prims.convert_element_type.default(
+                arg5_1, torch.float32
+            )
+            arg5_1 = None
+            mul_3 = torch.ops.aten.mul.Tensor(convert_element_type_10, mul_2)
+            convert_element_type_10 = mul_2 = None
+            convert_element_type_11 = torch.ops.prims.convert_element_type.default(
+                mul_3, torch.bfloat16
+            )
+            mul_3 = None
+            view_6 = torch.ops.aten.view.default(
+                convert_element_type_8, [8, 4096, -1, 128]
+            )
+            convert_element_type_8 = None
+            view_7 = torch.ops.aten.view.default(
+                convert_element_type_11, [8, 4096, -1, 128]
+            )
+            convert_element_type_11 = None
+            view_8 = torch.ops.aten.view.default(view_5, [8, 4096, -1, 128])
+            view_5 = None
+            convert_element_type_12 = torch.ops.prims.convert_element_type.default(
+                view_6, torch.float32
+            )
+            view_6 = None
+            convert_element_type_13 = torch.ops.prims.convert_element_type.default(
+                view_7, torch.float32
+            )
+            view_7 = None
+            unsqueeze = torch.ops.aten.unsqueeze.default(arg6_1, 0)
+            unsqueeze_1 = torch.ops.aten.unsqueeze.default(unsqueeze, 2)
+            unsqueeze = None
+            unsqueeze_2 = torch.ops.aten.unsqueeze.default(arg7_1, 0)
+            unsqueeze_3 = torch.ops.aten.unsqueeze.default(unsqueeze_2, 2)
+            unsqueeze_2 = None
+            mul_4 = torch.ops.aten.mul.Tensor(convert_element_type_12, unsqueeze_3)
+            unsqueeze_3 = None
+            view_9 = torch.ops.aten.view.default(
+                convert_element_type_12, [8, 4096, 16, 2, 64]
+            )
+            convert_element_type_12 = None
+            unbind = torch.ops.aten.unbind.int(view_9, -2)
+            view_9 = None
+            getitem = unbind[0]
+            getitem_1 = unbind[1]
+            unbind = None
+            neg = torch.ops.aten.neg.default(getitem_1)
+            getitem_1 = None
+            cat = torch.ops.aten.cat.default([neg, getitem], -1)
+            neg = getitem = None
+            mul_5 = torch.ops.aten.mul.Tensor(cat, unsqueeze_1)
+            cat = unsqueeze_1 = None
+            add_2 = torch.ops.aten.add.Tensor(mul_4, mul_5)
+            mul_4 = mul_5 = None
+            unsqueeze_4 = torch.ops.aten.unsqueeze.default(arg6_1, 0)
+            arg6_1 = None
+            unsqueeze_5 = torch.ops.aten.unsqueeze.default(unsqueeze_4, 2)
+            unsqueeze_4 = None
+            unsqueeze_6 = torch.ops.aten.unsqueeze.default(arg7_1, 0)
+            arg7_1 = None
+            unsqueeze_7 = torch.ops.aten.unsqueeze.default(unsqueeze_6, 2)
+            unsqueeze_6 = None
+            mul_6 = torch.ops.aten.mul.Tensor(convert_element_type_13, unsqueeze_7)
+            unsqueeze_7 = None
+            view_10 = torch.ops.aten.view.default(
+                convert_element_type_13, [8, 4096, 16, 2, 64]
+            )
+            convert_element_type_13 = None
+            unbind_1 = torch.ops.aten.unbind.int(view_10, -2)
+            view_10 = None
+            getitem_2 = unbind_1[0]
+            getitem_3 = unbind_1[1]
+            unbind_1 = None
+            neg_1 = torch.ops.aten.neg.default(getitem_3)
+            getitem_3 = None
+            cat_1 = torch.ops.aten.cat.default([neg_1, getitem_2], -1)
+            neg_1 = getitem_2 = None
+            mul_7 = torch.ops.aten.mul.Tensor(cat_1, unsqueeze_5)
+            cat_1 = unsqueeze_5 = None
+            add_3 = torch.ops.aten.add.Tensor(mul_6, mul_7)
+            mul_6 = mul_7 = None
+            convert_element_type_14 = torch.ops.prims.convert_element_type.default(
+                add_2, torch.bfloat16
+            )
+            add_2 = None
+            convert_element_type_15 = torch.ops.prims.convert_element_type.default(
+                add_3, torch.bfloat16
+            )
+            add_3 = None
+            permute_3 = torch.ops.aten.permute.default(
+                convert_element_type_14, [0, 2, 1, 3]
+            )
+            convert_element_type_14 = None
+            permute_4 = torch.ops.aten.permute.default(
+                convert_element_type_15, [0, 2, 1, 3]
+            )
+            convert_element_type_15 = None
+            permute_5 = torch.ops.aten.permute.default(view_8, [0, 2, 1, 3])
+            view_8 = None
+            return (permute_3, permute_4, permute_5)
+
+        from torch._dynamo.debug_utils import aot_graph_input_parser
+
+        kwargs = aot_graph_input_parser(forward)
+        out, code = run_and_get_code(torch.compile(forward), **kwargs)
+        # ignore tiny values.. prior to this fix absolute error was ~28
+        self.assertEqual(forward(**kwargs), out, atol=0.01, rtol=2)
+        FileCheck().check_not("in_out").run(code[0])
+
     # https://github.com/pytorch/pytorch/issues/104937
     def test_linear_with_zero_infeature_size(self):
         m = nn.Linear(in_features=0, out_features=0, bias=True).to("cuda")

torch/_inductor/scheduler.py

@@ -462,6 +462,38 @@ def decide_inplace_update(self) -> None:
             | self.scheduler.completed_operations
         )
 
+        def single_index_in_fused_node(buf_to_be_inplaced: SchedulerBuffer) -> bool:
+            # Inside of NodeUser, we track that the read and write are equivalent
+            # before deciding if the use can be inplace.
+            # But if that use is fused into a larger kernel, we need to check equivalence
+            # of other accesses in fused scheduler node as well.
+            fused_node = buf_to_be_inplaced.scheduler.get_fused_node(self)
+            buf_name = buf_to_be_inplaced.get_name()
+            # Dedup read/writes with equivalent indices
+            # TODO - would be nice if we could just cache accesses on ReadWrites,
+            # and inforce variant that this class & members are functional..
+            deps: OrderedSet[Dep] = OrderedSet()
+            for user in buf_to_be_inplaced.users:
+                user_node = user.node
+                if not isinstance(user_node, BaseSchedulerNode):
+                    continue
+
+                if (
+                    buf_to_be_inplaced.scheduler.get_fused_node(user_node)
+                    is not fused_node
+                ):
+                    continue
+
+                deps |= (
+                    o
+                    for o in user_node.read_writes.reads_and_writes()
+                    if o.name == buf_name
+                )
+                if len(deps) > 1:
+                    return False
+
+            return True
+
         for buf in self.get_outputs():
             buf_node = buf.node
             assert buf_node is not None
@@ -513,6 +545,7 @@ def decide_inplace_update(self) -> None:
                             and len(input_buf.node.get_inputs_that_alias_output()) > 0
                         )
                         and can_match_buffer_size(input_buf.node, buf.node)
+                        and single_index_in_fused_node(input_buf)
                     ):
                         # if there isn't a triton kernel, then we don't need to call triton-specific things.
                         # but TODO this might be a convenient place to signal to the Collective kernels to inplace