-
Notifications
You must be signed in to change notification settings - Fork 73
Open
Labels
Milestone
Description
Fix the variants below:
python/test/gluon/test_lowerings.py::test_histogram[2048-2-src_layout3-dst_layout3]
python/test/gluon/test_lowerings.py::test_histogram[32-32-src_layout4-dst_layout4]
M = 2048, bins = 2
src_layout = DistributedLinearLayout(reg_bases=[[32], [64], [128], [256], [512], [1024]], lane_bases=[[0], [16], [4], [2], [1]], warp_bases=[[0], [8]], block_bases=[], shape=[2048])
dst_layout = BlockedLayout(size_per_thread=[1], threads_per_warp=[32], warps_per_cta=[4], order=[0], ctas_per_cga=[1], cta_split_num=[1], cta_order=[0])
device = 'xpu'
@pytest.mark.parametrize("M, bins, src_layout, dst_layout", _histogram_cases())
def test_histogram(M, bins, src_layout, dst_layout, device):
@gluon.jit
def kernel(x_ptr, z_ptr, M: ttgl.constexpr, B: ttgl.constexpr, src_layout: ttgl.constexpr,
dst_layout: ttgl.constexpr):
offs = ttgl.arange(0, M, layout=src_layout)
x = ttgl.load(x_ptr + offs)
h = ttgl.histogram(x, B, layout=dst_layout)
z_offs = ttgl.arange(0, B, layout=dst_layout)
ttgl.store(z_ptr + z_offs, h)
torch.manual_seed(0)
x = torch.randint(0, bins, (M, ), dtype=torch.int32, device=device)
z = torch.zeros((bins, ), dtype=torch.int32, device=device)
z_torch = torch.histc(x.float(), bins=bins, min=0, max=bins - 1).to(torch.int32)
kernel[(1, )](x, z, M, bins, src_layout, dst_layout, num_warps=4)
> torch.testing.assert_close(z, z_torch, atol=0, rtol=0)
E AssertionError: Tensor-likes are not equal!
E
E Mismatched elements: 2 / 2 (100.0%)
E Greatest absolute difference: 35 at index (0,)
E Greatest relative difference: 0.03496503457427025 at index (0,)
test_lowerings.py:316: AssertionError