Repository: pytorch/torchdynamo

Issue:
ValueError: Cannot broadcast
Repro:
```
benchmarks/huggingface.py -d cuda --inductor --training --float32 --use-eval-mode -k CamemBert
```

Error:
```
ERROR:root:unhandled error
Traceback (most recent call last):
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1309, in _compile
    generator.visit(self.parse())
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 754, in visit
    return super().visit(node)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 245, in visit_Module
    ast.NodeVisitor.generic_visit(self, node)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/ast.py", line 368, in generic_visit
    self.visit(item)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 754, in visit
    return super().visit(node)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 306, in visit_FunctionDef
    has_ret = self.visit_compound_statement(node.body)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 239, in visit_compound_statement
    self.last_ret = self.visit(stmt)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 754, in visit
    return super().visit(node)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 631, in visit_For
    self.visit_compound_statement(node.body)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 239, in visit_compound_statement
    self.last_ret = self.visit(stmt)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 754, in visit
    return super().visit(node)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 741, in visit_Expr
    ast.NodeVisitor.generic_visit(self, node)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/ast.py", line 370, in generic_visit
    self.visit(value)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 754, in visit
    return super().visit(node)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/ast.py", line 360, in visit
    return visitor(node)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 693, in visit_Call
    ret = fn(*args, _builder=self.builder, **kws)
  File "/fsx/users/binbao/triton/python/triton/language/core.py", line 46, in wrapper
    return fn(*args, **kwargs)
  File "/fsx/users/binbao/triton/python/triton/language/core.py", line 878, in atomic_add
    return semantic.atomic_add(pointer, val, mask, _builder)
  File "/fsx/users/binbao/triton/python/triton/language/semantic.py", line 890, in atomic_add
    ptr, val, mask = atom_red_typechecking_impl(ptr, val, mask, 'add', builder)
  File "/fsx/users/binbao/triton/python/triton/language/semantic.py", line 802, in atom_red_typechecking_impl
    val = broadcast_impl_shape(val, ptr.type.get_block_shapes(), builder)
  File "/fsx/users/binbao/triton/python/triton/language/semantic.py", line 498, in broadcast_impl_shape
    raise ValueError(f"Cannot broadcast, the expanded size of the tensor ({shape[i]})"
ValueError: Cannot broadcast, the expanded size of the tensor (1) must match the existing size (64) at non-singleton dimension 0: [64, 64], [1, 64]

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/fsx/users/binbao/torchdynamo-tip/benchmarks/common.py", line 974, in run_one_model
    new_result = accuracy_model_iter_fn(model, example_inputs)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/eval_frame.py", line 135, in _fn
    return fn(*args, **kwargs)
  File "benchmarks/huggingface.py", line 417, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "benchmarks/huggingface.py", line 417, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "benchmarks/huggingface.py", line 421, in forward_and_backward_pass
    pred = mod(**cloned_inputs)
  File "/fsx/users/binbao/pytorch-release/torch/nn/modules/module.py", line 1185, in _call_impl
    return forward_call(*input, **kwargs)
  File "/fsx/users/binbao/conda/envs/release/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 1060, in forward
    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/eval_frame.py", line 135, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 790, in forward
    return compiled_f(
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 776, in new_func
    compiled_fn = create_aot_dispatcher_function(
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 355, in create_aot_dispatcher_function
    aot_dispatch_autograd(flat_fn, fake_flat_tensor_args, aot_config)
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 274, in aot_dispatch_autograd
    compiled_bw = aot_config.bw_compiler(bw_module, bw_args)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/optimizations/backends.py", line 542, in _wrapped_bw_compiler
    return torchdynamo.disable(bw_compiler(*args, **kwargs))
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 246, in bw_compiler
    return compile_fx_inner(
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 125, in compile_fx_inner
    compiled_fn = cudagraphify(
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 166, in cudagraphify
    model(*inputs)
  File "/tmp/torchinductor_binbao/ky/ckyjxnaywoz5k2s5yl3w5lqpfalqlslyirv3lipwugijpvdsqyel.py", line 1602, in call
    kernel23[grid(512)](buf608, mul_tensor_2, div_tensor_62, view_default_492, view_default_494, view_default_496, buf616, buf618, buf620, 512, 768)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 999, in __call__
    return self.kernel(*wargs, **kwargs, grid=self.grid)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1073, in __call__
    timings = {config: self._bench(*args, config=config, **kwargs)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1073, in <dictcomp>
    timings = {config: self._bench(*args, config=config, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/triton_ops/autotune.py", line 28, in _bench
    return super()._bench(*args, config=config, **kwargs)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1054, in _bench
    return triton.testing.do_bench(kernel_call)
  File "/fsx/users/binbao/triton/python/triton/testing.py", line 137, in do_bench
    fn()
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1053, in kernel_call
    self.kernel(*args, num_warps=config.num_warps, num_stages=config.num_stages, **current)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 988, in __call__
    return _triton.runtime.launch(wargs, self.fn.do_not_specialize, cache_key, self.fn.arg_names,
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 956, in add_to_cache
    return self.fn._warmup(key, arg_types=arg_types, device=device_idx, attributes=attributes, constants=constants, num_warps=num_warps, num_stages=num_stages,
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1285, in _warmup
    binary = self._compile(**compile)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1314, in _compile
    raise CompilationError(self.src, node) from e
triton.code_gen.CompilationError: at 52:44:
def kernel23(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, in_ptr5, out_ptr3, out_ptr4, out_ptr5, xnumel: tl.constexpr, rnumel: tl.constexpr, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK, 1])
    xmask = xindex < xnumel
    rbase = tl.reshape(tl.arange(0, RBLOCK), [1, RBLOCK])
    x0 = xindex
    _tmp1 = tl.zeros([XBLOCK, RBLOCK], tl.float32) + 0
    _tmp4 = tl.zeros([XBLOCK, RBLOCK], tl.float32) + 0
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r1 = rindex
        tmp0 = tl.load(in_ptr0 + r1 + (768*x0), xmask & rmask, eviction_policy='evict_last')
        tmp2 = tl.load(in_ptr1 + r1 + (768*x0), xmask & rmask, eviction_policy='evict_last')
        _tmp1 = tl.where(xmask & rmask, _tmp1 + tmp0, _tmp1)
        tmp3 = tmp0 * tmp2
        _tmp4 = tl.where(xmask & rmask, _tmp4 + tmp3, _tmp4)
    tmp1 = tl.reshape(tl.sum(_tmp1, 1), [XBLOCK, 1])
    tmp4 = tl.reshape(tl.sum(_tmp4, 1), [XBLOCK, 1])
    tmp5 = tl.load(in_ptr2 + x0, xmask)
    tmp14 = tl.load(in_ptr3 + x0, xmask)
    tmp21 = tl.load(in_ptr4 + x0, xmask)
    tmp27 = tl.load(in_ptr5 + x0, xmask)
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r1 = rindex
        tmp6 = tl.load(in_ptr0 + r1 + (768*x0), xmask & rmask, eviction_policy='evict_last')
        tmp10 = tl.load(in_ptr1 + r1 + (768*x0), xmask & rmask, eviction_policy='evict_last')
        tmp7 = 768
        tmp8 = tmp6 * tmp7
        tmp9 = tmp8 - tmp1
        tmp11 = tmp10 * tmp4
        tmp12 = tmp9 - tmp11
        tmp13 = tmp5 * tmp12
        tmp15 = 1
        tmp16 = tmp14 != tmp15
        tmp17 = 0
        tmp18 = tmp16 | tl.zeros(tmp13.shape, tmp16.dtype) if tmp13.numel > 1 else tmp16
        tmp19 = tmp18 | tl.zeros(tmp17.shape, tmp18.dtype) if tmp17.numel > 1 else tmp18
        tmp20 = tl.where(tmp19, tmp13, tmp17)
        tmp22 = -1
        tmp23 = tmp21 != tmp22
        tmp24 = tmp23 | tl.zeros(tmp13.shape, tmp23.dtype) if tmp13.numel > 1 else tmp23
        tmp25 = tmp24 | tl.zeros(tmp17.shape, tmp24.dtype) if tmp17.numel > 1 else tmp24
        tmp26 = tl.where(tmp25, tmp13, tmp17)
        tmp28 = tmp27 != tmp15
        tmp29 = tmp28 | tl.zeros(tmp13.shape, tmp28.dtype) if tmp13.numel > 1 else tmp28
        tmp30 = tmp29 | tl.zeros(tmp17.shape, tmp29.dtype) if tmp17.numel > 1 else tmp29
        tmp31 = tl.where(tmp30, tmp13, tmp17)
        tl.atomic_add(out_ptr3 + r1 + (768*tmp14) + tl.zeros([XBLOCK, RBLOCK], tl.int32), tmp20, xmask & rmask)
        tl.atomic_add(out_ptr4 + r1, tmp26, rmask)
                                            ^
```


Discussion:
@voznesenskym , please go ahead

> @voznesenskym , please go ahead
> 
> 

Awesome ty 
You should run this with the minifier :)

`INDUCTOR_REPRO_LEVEL=2 ./benchmarks/huggingface.py -d cuda --inductor --training --float32 --use-eval-mode -k CamemBert`

and then `python minimizer_repro.py`,
I'm not seeing this anymore. 

```
 benchmarks/huggingface.py -d cuda --inductor --training --float32 --use-eval-mode -k CamemBert

[W OperatorEntry.cpp:133] Warning: Overriding a previously registered kernel for the same operator and the same dispatch key
  operator: aten::inverse(Tensor self) -> Tensor
    registered at aten/src/ATen/RegisterSchema.cpp:6
  dispatch key: FuncTorchBatched
  previous kernel: registered at ../aten/src/ATen/autocast_mode.cpp:487
       new kernel: registered at /tmp/pip-req-build-uava6yj_/functorch/functorch/csrc/BatchRulesDecompositions.cpp:28 (function registerKernel)
[W OperatorEntry.cpp:133] Warning: Overriding a previously registered kernel for the same operator and the same dispatch key
  operator: aten::inverse(Tensor self) -> Tensor
    registered at aten/src/ATen/RegisterSchema.cpp:6
  dispatch key: FuncTorchBatched
  previous kernel: registered at ../aten/src/ATen/autocast_mode.cpp:487
       new kernel: registered at /tmp/pip-req-build-uava6yj_/functorch/functorch/csrc/BatchRulesDecompositions.cpp:28 (function registerKernel)
cuda train CamemBert                          WARNING torchinductor.ir: Using FallbackKernel: aten.cumsum
Similarity score=-0.001350881066173315
Accuracy failed for key name roberta.embeddings.word_embeddings.weight.grad
INCORRECT
1.876x p=0.00
speedup    gmean=1.88x mean=1.88x
```

Can someone confirm or deny my non-repro?
I still see this error.

Do you perhaps have an old version of PyTorch? I remember seeing that inverse error a couple weeks ago.
> I still see this error.
> 
> Do you perhaps have an old version of PyTorch? I remember seeing that inverse error a couple weeks ago.

Yea, used the makefile version. Discussed offline, small QOL PR coming. Will also build latest torch and try. 
> > I still see this error.
> > Do you perhaps have an old version of PyTorch? I remember seeing that inverse error a couple weeks ago.
> 
> Yea, used the makefile version. Discussed offline, small QOL PR coming. Will also build latest torch and try.

https://github.com/pytorch/torchdynamo/pull/908 relevant - will try to repro this on latest. Thanks for the ack @Chillee. 
Alright repro confirmed on latest torch. Looking now. 
Weird bug with sum, I think:

This is a repro generated + some of my own changes to trim things down further

```

import torch
from torch import tensor, device
import torch.fx as fx
from torchdynamo.testing import rand_strided
from math import inf
from torchinductor.compile_fx import compile_fx_inner
from torch.fx.experimental.proxy_tensor import make_fx

# torch version: 1.13.0a0+git1e98c87
# torch cuda version: 11.6
# torch git version: 1e98c87d38384cf8de79d0365c4feb9a9befd2b5


# 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 

class Repro(torch.nn.Module):
    def forward(self, _to_copy_default_228, view_default_270, expand_default_12, full_like_default):
        sum_dim_int_list_190 = torch.ops.aten.sum(_to_copy_default_228, [0, 1]); 
        where_self = torch.ops.aten.where.self(expand_default_12, full_like_default, full_like_default); 
        return (where_self,)
    
args = [((1, 512, 768), (393216, 768, 1), torch.float32, 'cuda'), ((512, 768), (768, 1), torch.float16, 'cuda'), ((512, 768), (1, 0), torch.bool, 'cuda'), ((512, 768), (768, 1), torch.float16, 'cuda')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
mod = make_fx(Repro())(*args)

compiled = compile_fx_inner(mod, args)
compiled(*args)
```

This has the error.

And this does not:

```

import torch
from torch import tensor, device
import torch.fx as fx
from torchdynamo.testing import rand_strided
from math import inf
from torchinductor.compile_fx import compile_fx_inner
from torch.fx.experimental.proxy_tensor import make_fx

# torch version: 1.13.0a0+git1e98c87
# torch cuda version: 11.6
# torch git version: 1e98c87d38384cf8de79d0365c4feb9a9befd2b5


# 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 

class Repro(torch.nn.Module):

    def forward(self, _to_copy_default_228, view_default_270, expand_default_12, full_like_default):
        where_self = torch.ops.aten.where.self(expand_default_12, full_like_default, full_like_default); 
        return (where_self,)
    
args = [((1, 512, 768), (393216, 768, 1), torch.float32, 'cuda'), ((512, 768), (768, 1), torch.float16, 'cuda'), ((512, 768), (1, 0), torch.bool, 'cuda'), ((512, 768), (768, 1), torch.float16, 'cuda')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
mod = make_fx(Repro())(*args)

compiled = compile_fx_inner(mod, args)
compiled(*args)
```

which means that sum has a weird side effect on `where` / operations downstream of it. 
Removing the lowering for sum fixes it :( 

Going to dig into that guy. 
So it looks like the ranges block produced during `codegen_header` is wrong - it makes one of `[XBLOCK, 1]` when we actually need `[XBLOCK, RBLOCK]`. The mask correctly uses `[XBLOCK, RBLOCK]` but the condition for where does not have the right shape, so broadcast fails. However, reshaping into `XBLOCK, RBLOCK` is bound to cause some other problems... Alternatively, it could be the the condition is fine but the in shape we load is not, because removing the rindex changes (both the rindex addition and rmask) makes it pass.
Yeah, I think we need some more rules about how rindex and rmask are used here, I dont understand that well enough so going to look into it. 
I wonder if we have any examples of atomic add in a reduction loop working? That might help here. I do not understand how, as its written, atomic add broadcast ever matches shapes in a reduction unless its by accident. 

Resolving patch:
diff --git a/torchinductor/codegen/triton.py b/torchinductor/codegen/triton.py
--- a/torchinductor/codegen/triton.py
+++ b/torchinductor/codegen/triton.py
@@ -717,7 +717,7 @@ def load(self, name: str, index: sympy.Expr, upcast: bool = False):
 
     def store(self, name, index, value, mode=None):
         var = self.args.output(name)
-        index, mask = self.indexing(index, value)
+        index, mask = self.indexing(index, value, dense_indexing=True)
         if mode is None:
             line = f"tl.store({var} + {index}, {value}, {mask})"
         elif mode == "atomic_add":


Test patch: