# test.py
import torch
import ctypes
import triton
import triton.language as tl


@triton.jit
def histogram_kernel(x_ptr, z_ptr):
    offset = tl.arange(0, 1)
    x = tl.load(x_ptr + offset)
    z = tl.histogram(x, 1)
    buf = (ctypes.c_int32 * 2).from_address(int(z_ptr))

    print(f'before store: {list(buf)}')
    tl.store(z_ptr + offset, z) # tl.store treats z values as int64 while they're int32
    print(f'after store: {list(buf)}')


device = 'cpu'
torch.manual_seed(17)
x = torch.ones(1, device=device, dtype=torch.int32)
z = torch.ones(2, dtype=torch.int32, device=device)
histogram_kernel[(1, )](x, z)

# Output:
# TRITON_INTERPRET=1 TRITON_TEST_SUITE=interpreter python test.py 
# before store: [1, 1]
# after store: [1, 0] <- second element shouldn't be cleared

import numpy as np

data = np.array([1], dtype=np.int32)
bins = 1

print(f'Data dtype before: {data.dtype}')
histogram = np.histogram(data, bins=bins, range=(0, bins))[0]
print(f'Data dtype after: {histogram.dtype}')

# Data dtype before: int32                                                                                                                                           
# Data dtype after: int64