Repository: pytorch/torchdynamo

Issue:
[inductor] Cannot make_shape_compatible: blocks must have the same rank
Repro


~~~
import torch
from torch import tensor, device
import torch.fx as fx
from torchdynamo.testing import rand_strided
from math import inf
from torch.fx.experimental.proxy_tensor import make_fx

# torch version: 1.14.0a0+git25725fd
# torch cuda version: 11.6
# torch git version: 25725fd62448165b91647304c26d676db22b6955


# CUDA Info:
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2022 NVIDIA Corporation
# Built on Thu_Feb_10_18:23:41_PST_2022
# Cuda compilation tools, release 11.6, V11.6.112
# Build cuda_11.6.r11.6/compiler.30978841_0

# GPU Hardware Info:
# NVIDIA A100-SXM4-40GB : 8


from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()



    def forward(self, view_1, sub, exp, alias_4, alias_2):
        sum_1 = torch.ops.aten.sum.dim_IntList(exp, [1], True);  exp = None
        log = torch.ops.aten.log.default(sum_1);  sum_1 = None
        sub_1 = torch.ops.aten.sub.Tensor(sub, log);  sub = log = None
        _to_copy_1 = torch.ops.aten._to_copy.default(sub_1, dtype = torch.float16);  sub_1 = None
        _to_copy_2 = torch.ops.aten._to_copy.default(_to_copy_1, dtype = torch.float32);  _to_copy_1 = None
        unsqueeze = torch.ops.aten.unsqueeze.default(view_1, 1);  view_1 = None
        gather = torch.ops.aten.gather.default(_to_copy_2, 1, unsqueeze);  _to_copy_2 = None
        squeeze = torch.ops.aten.squeeze.dim(gather, 1);  gather = None
        neg = torch.ops.aten.neg.default(squeeze);  squeeze = None
        mean = torch.ops.aten.mean.default(neg);  neg = None
        scatter = torch.ops.aten.scatter.value(alias_2, 1, unsqueeze, -1.0);  alias_2 = unsqueeze = None
        return [mean, scatter, alias_4]

args = [((1,), (1,), torch.int64, 'cuda'), ((1, 2), (2, 1), torch.float32, 'cuda'), ((1, 2), (2, 1), torch.float32, 'cuda'), ((1, 2), (2, 1), torch.float16, 'cuda'), ((1, 2), (2, 1), torch.float32, 'cuda')]
args = [rand_strided(sh, st, dt, dev) for (sh, st, dt, dev) in args]
mod = make_fx(Repro().to(device="cuda"))(*args)

from torchinductor.compile_fx import compile_fx_inner
from torchdynamo.debug_utils import same_two_models

compiled = compile_fx_inner(mod, args)
compiled(*args)
~~~




Discussion:

Resolving patch:
diff --git a/torchinductor/codegen/triton.py b/torchinductor/codegen/triton.py
--- a/torchinductor/codegen/triton.py
+++ b/torchinductor/codegen/triton.py
@@ -649,6 +649,7 @@ def combine_contiguous_dims(self, index: sympy.Expr, tree: IterationRangesRoot):
     def indexing(
         self,
         index: sympy.Expr,
+        *,
         copy_shape=None,
         dense_indexing=False,
     ):
@@ -774,7 +775,7 @@ def load(self, name: str, index: sympy.Expr):
 
     def store(self, name, index, value, mode=None):
         var = self.args.output(name)
-        index, mask = self.indexing(index, value, dense_indexing=True)
+        index, mask = self.indexing(index, dense_indexing=True)
         if mode is None:
             line = f"tl.store({var} + ({index}), {value}, {mask})"
         elif mode == "atomic_add":
@@ -856,7 +857,7 @@ def reduction(self, name, dtype, src_dtype, reduction_type, index, value):
             var_name = self.cse.reduction_cache[(src_dtype, reduction_type, value)]
             self.suffix.writeline(f"{result_var} = {var_name}")
         self.inside_reduction = False
-        index, mask = self.indexing(index, result_var)
+        index, mask = self.indexing(index)
         assert "rmask" not in index
         self.inside_reduction = True
         self.outside_loop_vars.add(result_var)
diff --git a/torchinductor/codegen/triton_template.py b/torchinductor/codegen/triton_template.py
--- a/torchinductor/codegen/triton_template.py
+++ b/torchinductor/codegen/triton_template.py
@@ -79,7 +79,9 @@ def assign_block_numel(self):
 
     def indexing(self, index: sympy.Expr, copy_shape=None, dense_indexing=True):
         # use dense_indexing for TritonTemplateKernel to avoid map::at error
-        return super().indexing(index, copy_shape, dense_indexing)
+        return super().indexing(
+            index, copy_shape=copy_shape, dense_indexing=dense_indexing
+        )
 
     def codegen_body(
         self, name, fuse, could_remove_kernel_buf, kernel_buf_replace_name


Test patch: