Some pointwise unary ops return 2-D results from 1-D inputs

[jdh8]

As we found out in tenstorrent/pytorch2.0_ttnn#198, several ops produce (1, N) tensors when (N,) tensors are expected.

Affected ops:

• ceil
• floor
• gelu
• rsqrt
• sqrt

Spared ops:

• cos
• erf
• exp

The lists above are non-exhaustive (yet). I fired this ticket before keeping trying on other ops because this issue is probably widespread.

[ayerofieiev-tt]

@jdh8 is the input tilized?

[jdh8]

I’m trying to compare implementations of unary ops, for example, ttnn::cos and ttnn::sqrt. However, I searched the C++ codebase, but I only found call sites instead of definitions. Where do I find these implementations?

(python_env) jdh8@tt-metal-skymizer-n150-3:~/tt-metal/ttnn/cpp$ git grep -i 'cos *('
ttnn/operations/eltwise/complex_unary/device/complex_unary_op.cpp:    Tensor c = ttnn::cos(input_b,output_mem_config);
ttnn/operations/eltwise/unary_backward/device/unary_backward_op.cpp:// name: cos(Tensor self) -> Tensor
ttnn/operations/eltwise/unary_backward/device/unary_backward_op.cpp:// # - name: acos(Tensor self) -> Tensor
ttnn/operations/eltwise/unary_backward/device/unary_backward_op.cpp:// self: grad * self.cos()
ttnn/operations/eltwise/unary_backward/device/unary_backward_op.cpp:    Tensor grad_input = ttnn::multiply(grad, ttnn::cos(input_tensor, output_mem_config), std::nullopt, output_mem_config);

[jdh8]

I've inspected op implementations in both groups, but I can't find a clear difference in treating 1-D tensors.

https://github.com/tenstorrent/tt-metal/blob/32ad2318f05b904ed51952b6097ceae8d60d479e/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/ckernel_sfpu_exp.h

https://github.com/tenstorrent/tt-metal/blob/32ad2318f05b904ed51952b6097ceae8d60d479e/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/ckernel_sfpu_sqrt.h

[jdh8]

After #13339, cos becomes affected. Now I believe that returning a 2-D tensor is intended. Squeezing 2-D tensors in the compiler looks like a true solution now.

[jdh8]

I made a workaround in the compiler at tenstorrent/pytorch2.0_ttnn#198 because I can't find the root cause. That PR also enables conversion for all rounding ops: ceil, floor, round, trunc.

[eyonland]

At the moment we believe that all unary eltwise ops are impacted by this. It sounds like our to_layout is converting to [1[32], 1024] for tiled output instead of the expected shape of [1024, 1[32]] .
To align with torch we have to do ttnn.transpose(input_tensor, -2, -1). See example below.

import ttnn
import torch

w = 1024
torch_input_tensor = torch.rand((w,), dtype=torch.bfloat16)
device_id = 0
device = ttnn.open_device(device_id=device_id)
input_tensor = ttnn.from_torch(torch_input_tensor, layout=ttnn.TILE_LAYOUT, device=device)

# required to align with torch
input_tensor = ttnn.transpose(input_tensor, -2, -1)
print(input_tensor.shape)
scalar = 3
output_tensor = ttnn.erf(input_tensor)
print(output_tensor.shape)

output_tensor = ttnn.from_device(output_tensor)
output_tensor = ttnn.to_torch(output_tensor)
print(output_tensor.shape)

pytorch_output_tensor = torch.erf(torch_input_tensor)
print("what torch does")
print(pytorch_output_tensor.shape)

ttnn.close_device(device=device)

@jdh8 can you confirm?

@ayerofieiev-tt , if our understanding is correct, this sounds like it is not an eltwise issue.

[ayerofieiev-tt]

@jdh8 , I will keep track of this one. We might have to update ops later, but at this moment this is pretty much blocked by Tensor Layout work

[prajaramanTT]

@jdh8 @ayerofieiev-tt Is this still an open issue ? If not, can you please close this ? Thanks