[FEA]: Identity matrix

[Da1sypetals]

Is this a new feature, an improvement, or a change to existing functionality?

New Feature

How would you describe the priority of this feature request?

Critical (currently preventing usage)

Please provide a clear description of problem this feature solves

Have trouble solving small linear systems using iterative methods without the support of creating identity matrix(tile).

Feature Description

Need a way to create identity tile (1s on the diagonal, 0s elsewhere).
It will be great if diag is also supported.

Describe your ideal solution

Proposed API would be:

ct.identity(4, dtype=ct.float32)
# returns Tile[float32, (4, 4)]

# for any tile `values` with type Tile[float32, (8)]
ct.diag(values)
# returns Tile[float32, (8, 8)]

Describe any alternatives you have considered

No response

Additional context

No response

Contributing Guidelines

• I agree to follow cuTile Python's contributing guidelines
• I have searched the open feature requests and have found no duplicates for this feature request

[haijieg]

@Da1sypetals this is a fair request that we will take on in the near future. Here's a workaround to unblock you.

import torch
import cuda.tile as ct

def identity(rank: int, dtype: ct.DType) -> ct.Tile:
    tmp = ct.arange(rank, dtype=ct.int32)
    tmp = (tmp[:, None] == tmp[None, :]).astype(dtype)
    return tmp

@ct.kernel
def identity_example(x):
    tile = identity(4, ct.int32)
    ct.store(x, (0, 0), tile)

x = torch.zeros((4, 4), dtype=torch.int32, device='cuda')
ct.launch(torch.cuda.current_stream(), (1,), identity_example, (x,))
print(x)