Skip to content

[TEST][Gluon] Fix test_lowerings.py::test_histogram #5148

@whitneywhtsang

Description

@whitneywhtsang

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions