Reductions reducing with padding on TILE_LAYOUT inputs.

[LPanosTT]

The following test fails:

def test_broken_reduction(device):

    shape = (1, 10)

    torch_input = torch.zeros(shape, dtype=torch.bfloat16)

    golden = torch.exp(torch_input).sum(dim=1, keepdim=True)

    tt_input = ttnn.from_torch(torch_input, device=device)
    tt_input = ttnn.to_layout(tt_input, ttnn.TILE_LAYOUT)

    tt_exp = ttnn.exp(tt_input) # Padded values become 1
    tt_sum = ttnn.sum(tt_exp, dim=1, keepdim=True) # reduction adds padded values
    tt_out = ttnn.to_layout(tt_sum, ttnn.ROW_MAJOR_LAYOUT)
    tt_out = ttnn.from_device(tt_out)
    tt_out = ttnn.to_torch(tt_out) # result is [[32.0]] rather than [[10.0]]

    assert torch.allclose(tt_out, golden)

I have a 1x10 tensor of values (0 in this case) I need to exponentiate and then take their sum. Because ttnn.exp is an element-wise operation the input must be in TILE_LAYOUT. After executing the op the padded 0's introduced by ttnn.to_layout are assigned the value of 1.0. I then perform the reduction and those 1.0 values are included in the sum. Thus giving an incorrect result.

[cmaryanTT]

@bbradelTT - reduce issue with non-zero padding value. Is this solved by any of your recent fixes?

[bbradelTT]

This issue is not solved by the recent fixes.

I had a PR out but had to remove the padding due to underlying issues.

This problem is really prevalent for min/max because they can't handle the default 0.

Looks like I'll need to pad the values for all ops, not just min/max.

A future PR will fix this issue.

[bbradelTT]

The PR I mentioned is #16157

Earlier commits have the padding for min/max

#16218 is blocking the changes that were there for min/max padding.

For debug mode, there's also an assert, although that can be removed for testing purposes (see #16157 )

[bbradelTT]

This has the potential to affect all ops, including matmul.

@ntarafdar what is the status of the proof of concept that you said you'd look at for dealing with garbage padding of all ops?

[LPanosTT]

This has the potential to affect all ops, including matmul.

@ntarafdar what is the status of the proof of concept that you said you'd look at for dealing with garbage padding of all ops?

Here's an example of where it does affect matmul.

def test_broken_mm(device):

    shape = (1, 10)
    torch_lhs = torch.zeros(shape, dtype=torch.bfloat16)
    torch_rhs = torch.zeros(10, 10, dtype=torch.bfloat16)

    golden = torch.exp(torch_lhs) @ torch.exp(torch_rhs)

    tt_lhs = ttnn.from_torch(torch_lhs, device=device)
    tt_rhs = ttnn.from_torch(torch_rhs, device=device)
    tt_lhs = ttnn.to_layout(tt_lhs, ttnn.TILE_LAYOUT)
    tt_rhs = ttnn.to_layout(tt_rhs, ttnn.TILE_LAYOUT)

    tt_exp_lhs = ttnn.exp(tt_lhs) # Padded values become 1
    tt_exp_rhs = ttnn.exp(tt_rhs) # Padded values become 1
    tt_sum = ttnn.matmul(tt_exp_lhs, tt_exp_rhs) # padded 1's are added....
    tt_out = ttnn.to_layout(tt_sum, ttnn.ROW_MAJOR_LAYOUT)
    tt_out = ttnn.from_device(tt_out)
    tt_out = ttnn.to_torch(tt_out) # result is [[32.0, 32.0, ..., 32.0]] (shape (1, 10)) rather than [[10.0, 10.0, ..., 10.0]] (shape (1, 10))

    assert torch.allclose(tt_out, golden)

[LPanosTT]

I've thought of two ways we could handle this (possibly):

• For reduction at least, we pass a pad value to run_with_autoformat and that value is set correctly when it's passed. However, at ttnn/cpp/ttnn/run_operation.cpp:356 the variable pad_input is determined to be false, which I believe is because the input tensor is already padded. As a result, the pad value is not applied to the input tensor.

  • We could remedy this by adding a bool force_update_pad_value argument to run_with_autoformat (or overload the function) that would apply the pad_value regardless of whether the shape is already padded.

• To me it seems like we generally want the padded value to be 0 for any TILE_LAYOUT tensor which has a logical shape that is not tile aligned. We could force a pad value update on the input to exp and set the pad value to -inf. If you go this route I supposed you'd want to do the same for any eltwise unary where f(0) != 0 (i.e. cos).

[eyonland]

No vision for dealing with pad value has been settled as far as I know.

I tend to favor the idea that we choose to allow padded values to always be undefined and explicitly handle the padding like we do softmax where we set the -inf in padded sections as part of the operation. Basically we allow garbage to exist in the padded sections after ttnn::exp and and ttnn::cos and instead you handle all reduction ops like we do softmax.

[bbradelTT]

ttnn/cpp/ttnn/operations/normalization/softmax/device/softmax_op.cpp
has

            ttnn::operations::experimental::auto_format::FormatParams input_format_params = {
                .pad_shape = input_pad_shape,
                .pad_value = -std::numeric_limits<float>::infinity(),
                .target_layout = Layout::TILE};

[ntarafdar]

We met today to discuss this:

Short term Goals:

• @bbradelTT will look in short term if autoformat in reduce can fix it
• @bbradelTT will also look at the overhead of padding relative to the rest of the op
• @ntarafdar will look at immediately staffing explicit fill op (this is needed in general, and can also be an alternative if autoformat in reduce does not fix it)
• @ntarafdar will look at overhead of explicit fill op.

Longer term Goals: from next week look at adding this padding value to tensor spec with @TT-BrianLiu

• Ops should be on an opt-in basis
• Default behaviour should be safe (maybe assume garbage padding if not specified)