diff --git a/ascend/examples/generalization_cases/test_flip.py b/ascend/examples/generalization_cases/test_flip.py index 74507a6e0c74b32186ccdb2255726f50d49475e9..dc321a3a61889b46cc0cb3012fd534c83775ef63 100644 --- a/ascend/examples/generalization_cases/test_flip.py +++ b/ascend/examples/generalization_cases/test_flip.py @@ -13,7 +13,7 @@ from test_common import TestUtils, check_ub_mem_overflow import triton.language.extra.ascend.libdevice as libdevice @triton.jit -def fn_npu_1d(output_ptr, x_ptr, XB: tl.constexpr): +def fn_npu_1d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): xidx = tl.arange(0, XB) idx = xidx X = tl.load(x_ptr + idx) @@ -22,10 +22,9 @@ def fn_npu_1d(output_ptr, x_ptr, XB: tl.constexpr): tl.store(output_ptr + oidx, ret) @triton.jit -def fn_npu_2d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr): - xoffs = tl.program_id(0) * XB - xidx = tl.arange(0, XB) + xoffs - yidx = tl.arange(0, YB) +def fn_npu_2d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + tl.program_id(0) * XB + yidx = tl.arange(0, YB) + tl.program_id(1) * YB idx = xidx[:, None] * YB + yidx[None, :] X = tl.load(x_ptr + idx) ret = libdevice.flip(X, 1) @@ -34,9 +33,9 @@ def fn_npu_2d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr): @triton.jit def fn_npu_3d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): - xidx = tl.arange(0, XB) - yidx = tl.arange(0, YB) - zidx = tl.arange(0, ZB) + xidx = tl.arange(0, XB) + tl.program_id(0) * XB + yidx = tl.arange(0, YB) + tl.program_id(1) * YB + zidx = tl.arange(0, ZB) + tl.program_id(2) * ZB idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] X = tl.load(x_ptr + idx) ret = libdevice.flip(X, 2) @@ -69,23 +68,34 @@ def test_flip(shape, dtype): else: x = torch.randint(low=0, high=128, size=shape, dtype=data_dtype).npu() - # torch_npu䷾M弾T¯弾L~Auint32漾Z~Dflip torch_input = x if x.dtype != torch.uint32 else x.to(torch.float32) torch_res = torch.flip(torch_input, dims=(-1,)) triton_res = torch.empty(shape, dtype=data_dtype).npu() + + grid = [1, 1, 1] + shape0 = shape[0] + shape1 = 1 + shape2 = 1 + if len(shape) == 2: + shape1 = shape[1] + if len(shape) == 3: + shape1 = shape[1] + shape2 = shape[2] + values = [shape0, shape1, shape2] + if dtype == 'bool': + if x.numel() * x.element_size() >= 512 and len(shape) > 1: + grid = [shape0, 1, 1] + values[0] = 1 + else: + if x.numel() * x.element_size() >= 8192 and len(shape) > 1: + grid = [shape0, 1, 1] + values[0] = 1 if len(shape) == 1: - fn_npu_1d[1, 1, 1](triton_res, x, shape[0]) + fn_npu_1d[grid](triton_res, x, values[0], values[1], values[2]) elif len(shape) == 2: - shape0 = shape[0] - shape1 = shape[1] - if x.numel() * x.element_size() >= 8192: - grid = (shape0, 1, 1) - shape0 = 1 - else: - grid = (1, 1, 1) - fn_npu_2d[grid](triton_res, x, shape0, shape1) + fn_npu_2d[grid](triton_res, x, values[0], values[1], values[2]) elif len(shape) == 3: - fn_npu_3d[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + fn_npu_3d[grid](triton_res, x, values[0], values[1], values[2]) triton_res = triton_res if triton_res.dtype != torch.uint32 else triton_res.to(torch.float32) cmp_dtype = dtype if dtype != 'uint32' else 'float32' diff --git a/ascend/examples/generalization_cases/test_general_add.py b/ascend/examples/generalization_cases/test_general_add.py index e1a1b79213445624e18b1bd220ef3a842a39bd68..7b58cfd24dfc4cf6d126d70dfb8557553fa25480 100644 --- a/ascend/examples/generalization_cases/test_general_add.py +++ b/ascend/examples/generalization_cases/test_general_add.py @@ -71,10 +71,15 @@ def test_add(shape, dtype): znumel = shape[2] grid = (1, 1, 1) - if x.numel() * x.element_size() >= 8192: - grid = (1, 1, ZB) - ZB = 1 + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 triton_add[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) - test_common.validate_cmp(dtype, ans, output) \ No newline at end of file + test_common.validate_cmp(dtype, ans, output) diff --git a/ascend/examples/generalization_cases/test_general_floordiv.py b/ascend/examples/generalization_cases/test_general_floordiv.py index 87826a8755e2f8c83cff2dff3a73b7471c76cf91..f7753415991db8f2ef36680537507babd02015ef 100644 --- a/ascend/examples/generalization_cases/test_general_floordiv.py +++ b/ascend/examples/generalization_cases/test_general_floordiv.py @@ -70,9 +70,14 @@ def test_add(shape, dtype): znumel = shape[2] grid = (1, 1, 1) - if x.numel() * x.element_size() >= 8192: - grid = (1, 1, ZB) - ZB = 1 + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 triton_floordiv[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) diff --git a/ascend/examples/generalization_cases/test_general_mul.py b/ascend/examples/generalization_cases/test_general_mul.py index 71035213b3923700ecb4e855af818bc51b332d4f..fea21df8b0036bf498c22b2abea700e58355f0b0 100644 --- a/ascend/examples/generalization_cases/test_general_mul.py +++ b/ascend/examples/generalization_cases/test_general_mul.py @@ -72,9 +72,14 @@ def test_mul(shape, dtype): znumel = shape[2] grid = (1, 1, 1) - if x.numel() * x.element_size() >= 8192: - grid = (1, 1, ZB) - ZB = 1 + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 triton_mul[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) diff --git a/ascend/examples/generalization_cases/test_general_sub.py b/ascend/examples/generalization_cases/test_general_sub.py index 074ed2c72bc18462fc122d15aaa1bb921c91830f..46e61d07c04b42558b7d96956eda3039607d3e3c 100644 --- a/ascend/examples/generalization_cases/test_general_sub.py +++ b/ascend/examples/generalization_cases/test_general_sub.py @@ -71,9 +71,14 @@ def test_sub(shape, dtype): znumel = shape[2] grid = (1, 1, 1) - if x.numel() * x.element_size() >= 8192: - grid = (1, 1, ZB) - ZB = 1 + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 triton_sub[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) diff --git a/ascend/examples/generalization_cases/test_mod.py b/ascend/examples/generalization_cases/test_mod.py index c3ccc6e2f573e76c8fd497993dd34ce754347852..f0ec652b1e146277fca1ead5288a2e45bb45df7e 100644 --- a/ascend/examples/generalization_cases/test_mod.py +++ b/ascend/examples/generalization_cases/test_mod.py @@ -35,10 +35,8 @@ def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, tl.store(output_ptr + idx, ret) @pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) -#@pytest.mark.parametrize('dtype', ['int8']) -@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32']) def test_case2(dtype, shape): - # 漾T~_弾H~P弾U°弾M® x = test_common.generate_tensor(shape, dtype).npu() x[x <= 0] = 1 y = test_common.generate_tensor(shape, dtype).npu() @@ -79,14 +77,12 @@ def test_case2(dtype, shape): znumel = shape[2] grid = (1, 1, 1) - if x.numel() * x.element_size() >= 8192: - if max(XB, YB, ZB) == XB: - grid = (XB, 1, 1) - XB = 1 - elif max(XB, YB, ZB) == YB: - grid = (1, YB, 1) - YB = 1 - else: + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: grid = (1, 1, ZB) ZB = 1 diff --git a/ascend/examples/generalization_cases/test_neg.py b/ascend/examples/generalization_cases/test_neg.py index 1880b1e20dd50a366c93f3bd3a7215306f26624b..77128667a610cba0c28f863753b73ad613fd652a 100644 --- a/ascend/examples/generalization_cases/test_neg.py +++ b/ascend/examples/generalization_cases/test_neg.py @@ -39,7 +39,6 @@ def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, @pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64']) def test_case2(dtype, shape): - # 生成数据 x = test_common.generate_tensor(shape, dtype).npu() y = test_common.generate_tensor(shape, dtype).npu() z = test_common.generate_tensor(shape, dtype).npu() @@ -76,9 +75,14 @@ def test_case2(dtype, shape): znumel = shape[2] grid = (1, 1, 1) - if x.numel() * x.element_size() >= 8192: - grid = (1, 1, ZB) - ZB = 1 + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel)