Tensor and Add Operation

ttnn.Tensor is the central type of ttnn.

It is similar to torch.Tensor in the sense that it represents multi-dimensional matrix containing elements of a single data type.

The are a few key differences:

• ttnn.Tensor can be stored in the SRAM or DRAM of TensTorrent devices

• ttnn.Tensor doesn’t have a concept of the strides, however it has a concept of row-major and tile layout

• ttnn.Tensor has support for data types not supported by torch such as bfp8 for example

• ttnn.Tensor’s shape stores the padding added to the tensor due to TILE_LAYOUT

Creating a tensor

The recommended way to create a tensor is by using torch create function and then simply calling ttnn.from_torch. So, let’s import both torch and ttnn

import torch
import ttnn

And now let’s create a torch Tensor and convert it to ttnn Tensor

torch_tensor = torch.rand(3, 4)
ttnn_tensor = ttnn.from_torch(torch_tensor)

print(f"shape: {ttnn_tensor.shape}")
print(f"layout: {ttnn_tensor.layout}")
print(f"dtype: {ttnn_tensor.dtype}")

shape: ttnn.Shape([3, 4])
layout: Layout.ROW_MAJOR
dtype: DataType.FLOAT32

As expected we get a tensor of shape [3, 4] in row-major layout with a data type of float32.

Host Storage: Borrowed vs Owned

In this particular case, ttnn Tensor will borrow the data of the torch Tensor because ttnn Tensor is in row-major layout, torch tensor is contiguous and their data type matches.

Let’s print the current ttnn tensor, set element of torch tensor to 1234 and print the ttnn Tensor again to see borrowed storage in action

print(f"Original values:\n{ttnn_tensor}")
torch_tensor[:] = 1234
print(f"New values are all going to be 1234:\n{ttnn_tensor}")

Original values:
Tensor([ [0.868396, 0.199809, 0.505658, 0.0919966],
    [0.441207, 0.465399, 0.225584, 0.497159],
    [0.205919, 0.219386, 0.0836022, 0.761129]], dtype=float32 )

New values are all going to be 1234:
Tensor([ [1234, 1234, 1234, 1234],
    [1234, 1234, 1234, 1234],
    [1234, 1234, 1234, 1234]], dtype=float32 )

We try our best to use borrowed storage but if the torch data type is not supported in ttnn, then we don’t have a choice but to automatically pick a different data type and copy data

torch_tensor = torch.rand(3, 4).to(torch.float16)
ttnn_tensor = ttnn.from_torch(torch_tensor)
print("torch_tensor.dtype:", torch_tensor.dtype)
print("ttnn_tensor.dtype:", ttnn_tensor.dtype)

print(f"Original values:\n{ttnn_tensor}")
torch_tensor[0, 0] = 1234
#print(f"Original values again because the tensor doesn't use borrowed storage:\n{ttnn_tensor}")

torch_tensor.dtype: torch.float16
ttnn_tensor.dtype: DataType.BFLOAT16
Original values:
Tensor([ [0.9375, 0.0683594, 0.765625, 0.894531],
    [0.100098, 0.285156, 0.597656, 0.21582],
    [0.203125, 0.730469, 0.310547, 0.453125]], dtype=bfloat16 )

Data Type

The data type of the ttnn tensor can be controlled explicitly when conversion from torch.

torch_tensor = torch.rand(3, 4).to(torch.float32)
ttnn_tensor = ttnn.from_torch(torch_tensor, dtype=ttnn.bfloat16)
print(f"torch_tensor.dtype: {torch_tensor.dtype}")
print(f"ttnn_tensor.dtype: {ttnn_tensor.dtype}")

torch_tensor.dtype: torch.float32
ttnn_tensor.dtype: DataType.BFLOAT16

Layout

TensTorrent hardware is most efficiently utilized when running tensors using tile layout. The current tile size is hard-coded to [32, 32]. It was determined to be the optimal size for a tile given the compute, memory and data transfer constraints.

ttnn provides easy and intuitive way to convert from row-major layout to tile layout and back.

torch_tensor = torch.rand(3, 4).to(torch.float16)
ttnn_tensor = ttnn.from_torch(torch_tensor)
print(f"Tensor in row-major layout:\nShape {ttnn_tensor.shape}\nLayout: {ttnn_tensor.layout}\n{ttnn_tensor}")
ttnn_tensor = ttnn.to_layout(ttnn_tensor, ttnn.TILE_LAYOUT)
print(f"Tensor in tile layout:\nShape {ttnn_tensor.shape}\nLayout: {ttnn_tensor.layout}\n{ttnn_tensor}")
ttnn_tensor = ttnn.to_layout(ttnn_tensor, ttnn.ROW_MAJOR_LAYOUT)
print(f"Tensor back in row-major layout:\nShape {ttnn_tensor.shape}\nLayout: {ttnn_tensor.layout}\n{ttnn_tensor}")

Tensor in row-major layout:
Shape ttnn.Shape([3, 4])
Layout: Layout.ROW_MAJOR
Tensor([ [0.020752, 0.0820312, 0.664062, 0.0742188],
    [0.0463867, 0.785156, 0.664062, 0.0195312],
    [0.304688, 0.287109, 0.671875, 0.808594]], dtype=bfloat16 )

Tensor in tile layout:
Shape ttnn.Shape([3[32], 4[32]])
Layout: Layout.TILE
Tensor([ [0.020752, 0.0820312, 0.664062, 0.0742188, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0.0463867, 0.785156, 0.664062, 0.0195312, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0.304688, 0.287109, 0.671875, 0.808594, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=bfloat16 )

Tensor back in row-major layout:
Shape ttnn.Shape([3, 4])
Layout: Layout.ROW_MAJOR
Tensor([ [0.020752, 0.0820312, 0.664062, 0.0742188],
    [0.0463867, 0.785156, 0.664062, 0.0195312],
    [0.304688, 0.287109, 0.671875, 0.808594]], dtype=bfloat16 )

Note that padding is automatically inserted to put the tensor into tile layout and it automatically removed after the tensor is converted back to row-major layout

The conversion to tile layout can be done when caling ttnn.from_torch

torch_tensor = torch.rand(3, 4).to(torch.float16)
ttnn_tensor = ttnn.from_torch(torch_tensor)
print(f"Tensor in row-major layout:\nShape {ttnn_tensor.shape}; Layout: {ttnn_tensor.layout}")

Tensor in row-major layout:
Shape ttnn.Shape([3, 4]); Layout: Layout.ROW_MAJOR

Note that ttnn.to_torch will always convert to row-major layout

Device storage

Finally, in order to actually utilize the tensor, we need to put it on the device. So, that we can run ttnn operations on it

Open the device

Use ttnn.open to get a handle to the device

device_id = 0
device = ttnn.open_device(device_id=device_id)

                  Metal | INFO     | Initializing device 0
                 Device | INFO     | Opening user mode device driver
2024-02-16 19:46:19.597 | INFO     | SiliconDriver   - Detected 1 PCI device : {0}
2024-02-16 19:46:19.610 | WARNING  | SiliconDriver   - init_detect_tt_device_numanodes(): Could not determine NumaNodeSet for TT device (physical_device_id: 0 pci_bus_id: 0000:00:08.0)
2024-02-16 19:46:19.610 | WARNING  | SiliconDriver   - Could not find NumaNodeSet for TT Device (physical_device_id: 0 pci_bus_id: 0000:00:08.0)
2024-02-16 19:46:19.612 | WARNING  | SiliconDriver   - bind_area_memory_nodeset(): Unable to determine TT Device to NumaNode mapping for physical_device_id: 0. Skipping membind.
---- ttSiliconDevice::init_hugepage: bind_area_to_memory_nodeset() failed (physical_device_id: 0 ch: 0). Hugepage allocation is not on NumaNode matching TT Device. Side-Effect is decreased Device->Host perf (Issue #893).
                  Metal | INFO     | AI CLK for device 0 is:   1202 MHz

Initialize tensors a and b with random values using torch

To create a tensor that can be used by a ttnn operation: 1. Create a tensor using torch 2. Use ttnn.from_torch to convert the tensor from torch.Tensor to ttnn.Tensor, change the layout to ttnn.TILE_LAYOUT and put the tensor on the device

torch.manual_seed(0)

torch_input_tensor_a = torch.rand((32, 32), dtype=torch.bfloat16)
torch_input_tensor_b = torch.rand((32, 32), dtype=torch.bfloat16)

input_tensor_a = ttnn.from_torch(torch_input_tensor_a, layout=ttnn.TILE_LAYOUT, device=device)
input_tensor_b = ttnn.from_torch(torch_input_tensor_b, layout=ttnn.TILE_LAYOUT, device=device)

Add tensor a and b

ttnn supports operator overloading, therefore operator + can be used instead of torch.add

output_tensor = input_tensor_a + input_tensor_b

Inspect the output tensor of the add in ttnn

As can be seen the tensor of the same shape, layout and dtype is produced

print(f"shape: {output_tensor.shape}")
print(f"dtype: {output_tensor.dtype}")
print(f"layout: {output_tensor.layout}")

shape: ttnn.Shape([32, 32])
dtype: DataType.BFLOAT16
layout: Layout.TILE

In general we expect layout and dtype to stay the same when running most operations unless explicit arguments to modify them are passed in. However, there are obvious exceptions like an embedding operation that takes in ttnn.uint32 and produces ttnn.bfloat16

Convert to torch and inspect the attributes of the torch tensor

When converting the tensor to torch, ttnn.to_torch will move the tensor from the device, convert to tile layout and figure out the best data type to use on the torch side

output_tensor = ttnn.to_torch(output_tensor)
print(f"shape: {output_tensor.shape}")
print(f"dtype: {output_tensor.dtype}")

shape: torch.Size([32, 32])
dtype: torch.bfloat16

Close the device

Close the handle the device. This is a very important step as the device can hang currently if not closed properly

ttnn.close_device(device)

                  Metal | INFO     | Closing device 0