From 6ca79f2b8efea9dfa99bb8eb3307f88d7f17e52c Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 19 Oct 2026 11:54:53 +0800 Subject: [PATCH] Fix Ruff CI --- src/ntops/kernels/__init__.py | 40 +++--- src/ntops/kernels/acosh.py | 36 +++-- src/ntops/kernels/adaptive_avg_pool2d.py | 47 ++++--- src/ntops/kernels/adaptive_max_pool2d.py | 41 +++--- src/ntops/kernels/addmv.py | 6 +- src/ntops/kernels/argsort.py | 15 ++- src/ntops/kernels/atan.py | 58 +++++---- src/ntops/kernels/batch_norm.py | 21 +-- src/ntops/kernels/bincount.py | 60 +++++---- src/ntops/kernels/fmax.py | 14 +- src/ntops/kernels/logsumexp.py | 2 + src/ntops/kernels/lp_pool1d.py | 31 +++-- src/ntops/kernels/lp_pool2d.py | 95 +++++++++----- src/ntops/kernels/lp_pool3d.py | 81 ++++++++---- src/ntops/kernels/max.py | 10 +- src/ntops/kernels/max_pool1d.py | 26 ++-- src/ntops/kernels/max_pool3d.py | 38 +++--- src/ntops/kernels/maximum.py | 18 ++- src/ntops/kernels/mean.py | 18 ++- src/ntops/kernels/median.py | 30 +++-- src/ntops/kernels/sort.py | 4 +- src/ntops/kernels/stack.py | 15 ++- src/ntops/torch/__init__.py | 42 +++--- src/ntops/torch/acosh.py | 38 +++--- src/ntops/torch/adaptive_avg_pool2d.py | 30 +++-- src/ntops/torch/adaptive_max_pool2d.py | 32 +++-- src/ntops/torch/addmv.py | 4 +- src/ntops/torch/argsort.py | 21 +-- src/ntops/torch/atan.py | 37 +++--- src/ntops/torch/batch_norm.py | 12 +- src/ntops/torch/bincount.py | 32 +++-- src/ntops/torch/fmax.py | 26 ++-- src/ntops/torch/lp_pool1d.py | 3 +- src/ntops/torch/lp_pool2d.py | 34 +++-- src/ntops/torch/lp_pool3d.py | 34 +++-- src/ntops/torch/max.py | 23 ++-- src/ntops/torch/max_pool1d.py | 35 +++-- src/ntops/torch/max_pool3d.py | 45 ++++--- src/ntops/torch/maximum.py | 31 ++--- src/ntops/torch/mean.py | 48 ++++--- src/ntops/torch/median.py | 26 ++-- src/ntops/torch/stack.py | 24 ++-- tests/test_acosh.py | 159 ++++++++++++----------- tests/test_adaptive_avg_pool2d.py | 21 +-- tests/test_adaptive_max_pool2d.py | 17 +-- tests/test_addmv.py | 12 +- tests/test_argsort.py | 14 +- tests/test_atan.py | 57 ++++---- tests/test_batch_norm.py | 28 ++-- tests/test_bincount.py | 25 ++-- tests/test_fmax.py | 10 +- tests/test_logsumexp.py | 135 ++++++++++--------- tests/test_lp_pool1d.py | 4 +- tests/test_lp_pool2d.py | 15 ++- tests/test_lp_pool3d.py | 8 +- tests/test_max.py | 8 +- tests/test_max_pool1d.py | 13 +- tests/test_max_pool3d.py | 15 ++- tests/test_maximum.py | 16 ++- tests/test_mean.py | 19 ++- tests/test_median.py | 20 ++- tests/test_stack.py | 14 +- 62 files changed, 1113 insertions(+), 780 deletions(-) diff --git a/src/ntops/kernels/__init__.py b/src/ntops/kernels/__init__.py index 12d337b..3167b4d 100644 --- a/src/ntops/kernels/__init__.py +++ b/src/ntops/kernels/__init__.py @@ -1,9 +1,17 @@ from ntops.kernels import ( abs, + acosh, + adaptive_avg_pool2d, + adaptive_max_pool2d, add, addmm, + addmv, alpha_dropout, + argsort, + atan, avg_pool2d, + batch_norm, + bincount, bitwise_and, bitwise_not, bitwise_or, @@ -18,6 +26,7 @@ dropout, eq, exp, + fmax, ge, gelu, gt, @@ -26,8 +35,18 @@ isnan, layer_norm, le, + logsumexp, + lp_pool1d, + lp_pool2d, + lp_pool3d, lt, + max, + max_pool1d, max_pool2d, + max_pool3d, + maximum, + mean, + median, mm, msort, mul, @@ -51,29 +70,10 @@ sin, softmax, sort, + stack, sub, tanh, threshold, - max_pool1d, - max_pool3d, - stack, - mean, - median, - maximum, - atan, - batch_norm, - bincount, - adaptive_max_pool2d, - acosh, - adaptive_avg_pool2d, - addmv, - argsort, - fmax, - logsumexp, - lp_pool1d, - lp_pool2d, - lp_pool3d, - max, ) __all__ = [ diff --git a/src/ntops/kernels/acosh.py b/src/ntops/kernels/acosh.py index 51d1a16..4acf023 100644 --- a/src/ntops/kernels/acosh.py +++ b/src/ntops/kernels/acosh.py @@ -21,52 +21,50 @@ def _log(x, dtype): def application(input, output): """ 计算反双曲余弦函数 acosh(x) = ln(x + sqrt(x² - 1)) - + 参数: input: 输入张量，形状为 (C // block_size, block_size) output: 输出张量，形状为 (C // block_size, block_size) - + 数值稳定性考虑: 1. 当x接近1时，x² - 1接近0，使用(x-1)(x+1)形式避免精度损失 2. 当x很大时，避免x²溢出，使用代数变换 3. float16特殊处理，提升到float32计算 """ dtype = output.dtype.dtype - + for i in range(input.shape[0]): # 获取当前块的数据 input_block = ntl.cast(input[i], dtype) - + # 数值稳定的acosh计算 # acosh(x) = ln(x + sqrt(x² - 1)) - + # 处理x接近1的情况：x² - 1 = (x-1)(x+1) # 这样可以避免当x接近1时的精度损失 x_minus_one = input_block - ntl.cast(1.0, dtype) x_plus_one = input_block + ntl.cast(1.0, dtype) - + # 计算 sqrt(x² - 1) = sqrt((x-1)(x+1)) sqrt_term = _sqrt(x_minus_one * x_plus_one, dtype) - + # 计算 x + sqrt(x² - 1) # 当x很大时，这可能会导致数值问题，但acosh的定义域x≥1 sum_term = input_block + sqrt_term - + # 最终计算 ln(x + sqrt(x² - 1)) result = _log(sum_term, dtype) - + # 处理边界情况：当x < 1时，acosh未定义，返回NaN # 当x == 1时，acosh(1) = 0 result = ntl.where( input_block < ntl.cast(1.0, dtype), ntl.cast(float("nan"), dtype), ntl.where( - input_block == ntl.cast(1.0, dtype), - ntl.cast(0.0, dtype), - result - ) + input_block == ntl.cast(1.0, dtype), ntl.cast(0.0, dtype), result + ), ) - + # 将结果存入输出 output[i] = result @@ -74,13 +72,13 @@ def application(input, output): def premake(ndim, dim, dtype=None, block_size=None): """ 准备acosh内核 - + 参数: ndim: 输入张量的维度 dim: 要计算acosh的维度 dtype: 数据类型 block_size: 分块大小，用于优化内存访问 - + 返回: arrangement_: 张量排列函数 application: 计算函数 @@ -89,8 +87,8 @@ def premake(ndim, dim, dtype=None, block_size=None): arrangement_ = functools.partial(arrangement, dim=dim, block_size=block_size) tensors = ( - Tensor(ndim, dtype=dtype), # 输入张量 - Tensor(ndim, dtype=dtype), # 输出张量 + Tensor(ndim, dtype=dtype), # 输入张量 + Tensor(ndim, dtype=dtype), # 输出张量 ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/adaptive_avg_pool2d.py b/src/ntops/kernels/adaptive_avg_pool2d.py index e506c48..6c4888d 100644 --- a/src/ntops/kernels/adaptive_avg_pool2d.py +++ b/src/ntops/kernels/adaptive_avg_pool2d.py @@ -1,32 +1,43 @@ import functools + import ninetoothed import ninetoothed.language as ntl from ninetoothed import Tensor + def _cast_to_f32(x, dtype): """ 为了保证累加精度，如果是 float16 则转为 float32 计算 """ return ntl.cast(x, ntl.float32) if dtype == ntl.float16 else x -def arrangement(input, output, kernel_size_flatted, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size): + +def arrangement( + input, + output, + kernel_size_flatted, + kernel_size_h, + kernel_size_w, + stride_h, + stride_w, + block_size, +): if block_size is None: block_size = ninetoothed.block_size() - # input: (N, C, H_in, W_in) + # input: (N, C, H_in, W_in) # output: (N, C, H_out, W_out) # 1. 窗口切分 input_arranged = input.tile( - (1, 1, kernel_size_h, kernel_size_w), - (1, 1, stride_h, stride_w) + (1, 1, kernel_size_h, kernel_size_w), (1, 1, stride_h, stride_w) ) # => (N, C, H_out, W_out), dtype=(1, 1, k_h, k_w) # 2. 展平与重排 input_arranged = input_arranged.ravel() # => (N, C, H_out, W_out, 1, 1, k_h, k_w) - + input_arranged = input_arranged.flatten(end_dim=4).flatten(start_dim=1) # => (N*C*H_out*W_out, k_h*k_w) @@ -35,7 +46,7 @@ def arrangement(input, output, kernel_size_flatted, kernel_size_h, kernel_size_w nearest_pow2 = 1 << (kernel_size_h * kernel_size_w - 1).bit_length() input_arranged = input_arranged.tile((1, nearest_pow2)) # => (..., 1), dtype=(1, nearest_pow2) - + input_arranged.dtype = input_arranged.dtype.squeeze(0) input_arranged = input_arranged.tile((block_size, -1)) input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) @@ -50,28 +61,32 @@ def arrangement(input, output, kernel_size_flatted, kernel_size_h, kernel_size_w return input_arranged, output_arranged, kernel_size_flatted + def application(input, output, kernel_size_flatted): # input: (block_size, nearest_pow2) # output: (block_size, ) # kernel_size_flatted: scalar tensor (k_h * k_w) dtype = input.dtype - + # 转为高精度进行 Sum val = _cast_to_f32(input, dtype) - + # 求和 (Axis 1 对应 nearest_pow2 维度) # 这里的 0 填充不会影响 Sum 结果 acc = ntl.sum(val, axis=1) - + # 求平均： Sum / Area # 注意：kernel_size_flatted 是实际的窗口大小，不是 nearest_pow2 res = acc / ntl.cast(kernel_size_flatted, acc.dtype) - + # 转回原类型 - output = ntl.cast(res, dtype) + output = ntl.cast(res, dtype) # noqa: F841 + -def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None): +def premake( + ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None +): arrangement_ = functools.partial( arrangement, kernel_size_h=kernel_size_h, @@ -83,9 +98,9 @@ def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=N tensors = ( # input: 设置 other=0，保证 tile 补齐的值不影响 sum - Tensor(ndim, dtype=dtype, other=0), - Tensor(ndim, dtype=dtype), # output - Tensor(0, dtype=dtype), # kernel_size_flatted (scalar) + Tensor(ndim, dtype=dtype, other=0), + Tensor(ndim, dtype=dtype), # output + Tensor(0, dtype=dtype), # kernel_size_flatted (scalar) ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/adaptive_max_pool2d.py b/src/ntops/kernels/adaptive_max_pool2d.py index 5caaac1..20aaa4c 100644 --- a/src/ntops/kernels/adaptive_max_pool2d.py +++ b/src/ntops/kernels/adaptive_max_pool2d.py @@ -1,26 +1,29 @@ import functools + import ninetoothed import ninetoothed.language as ntl from ninetoothed import Tensor -def arrangement(input, output, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size): + +def arrangement( + input, output, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size +): if block_size is None: block_size = ninetoothed.block_size() - # input: (N, C, H_in, W_in) + # input: (N, C, H_in, W_in) # output: (N, C, H_out, W_out) - + # 使用 tile 将输入切分为窗口 # floor_mode=True 对应默认行为，对于 Adaptive Pool，我们通常通过计算好的 stride/kernel 确保覆盖 input_arranged = input.tile( - (1, 1, kernel_size_h, kernel_size_w), - (1, 1, stride_h, stride_w) + (1, 1, kernel_size_h, kernel_size_w), (1, 1, stride_h, stride_w) ) # => (N, C, H_out, W_out), dtype=(1, 1, k_h, k_w) - + input_arranged = input_arranged.ravel() # => (N, C, H_out, W_out, 1, 1, k_h, k_w) - + input_arranged = input_arranged.flatten(end_dim=4).flatten(start_dim=1) # => (N*C*H_out*W_out, k_h*k_w) @@ -28,10 +31,10 @@ def arrangement(input, output, kernel_size_h, kernel_size_w, stride_h, stride_w, nearest_pow2 = 1 << (kernel_size_h * kernel_size_w - 1).bit_length() input_arranged = input_arranged.tile((1, nearest_pow2)) # => (..., 1), dtype=(1, nearest_pow2) - + input_arranged.dtype = input_arranged.dtype.squeeze(0) # => (..., 1), dtype=(nearest_pow2, ) - + input_arranged = input_arranged.tile((block_size, -1)) # => (..., 1), dtype=(block_size, nearest_pow2) input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) @@ -45,19 +48,23 @@ def arrangement(input, output, kernel_size_h, kernel_size_w, stride_h, stride_w, return input_arranged, output_arranged + def application(input, output): - # input: (block_size, nearest_pow2) + # input: (block_size, nearest_pow2) # output: (block_size, ) - + # 简单的 max reduction # 因为在 premake 中设置了 other=float("-inf")，padding 部分的值为负无穷， # 或者是 nearest_pow2 补齐产生的部分，通常默认为 0 或 padding 值， # 这里为了安全，可以显式处理 padding，但如果 arrange padding 正确，直接 max 即可。 # 假设 DSL 的 tile 填充行为遵循 Tensor 的 other 属性。 - - output = ntl.max(input, axis=1) -def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None): + output = ntl.max(input, axis=1) # noqa: F841 + + +def premake( + ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None +): arrangement_ = functools.partial( arrangement, kernel_size_h=kernel_size_h, @@ -69,8 +76,8 @@ def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=N tensors = ( # input: 设置 other 为负无穷，这样 tile 越界填充的值不会影响 max - Tensor(ndim, dtype=dtype, other=float("-inf")), - Tensor(ndim, dtype=dtype), # output + Tensor(ndim, dtype=dtype, other=float("-inf")), + Tensor(ndim, dtype=dtype), # output ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/addmv.py b/src/ntops/kernels/addmv.py index d10488c..2cf3e47 100644 --- a/src/ntops/kernels/addmv.py +++ b/src/ntops/kernels/addmv.py @@ -23,7 +23,7 @@ def arrangement( # 关键：强制 block_size_n 为 1，因为这是 MV 乘法 if block_size_n is None: - block_size_n = 1 + block_size_n = 1 if block_size_k is None: block_size_k = mm.BLOCK_SIZE_K @@ -68,7 +68,7 @@ def application(input, mat, vec, beta, alpha, output, input_precision): # 这里 output 是 (M, 1)，zeros 也创建 (M, 1) mm_output = ntl.zeros(output.shape, dtype=ntl.float32) mm.application(mat, vec, mm_output, input_precision) - + # 逐元素操作，形状兼容 (M, 1) output = beta * input + alpha * mm_output @@ -102,4 +102,4 @@ def premake( Tensor(0, dtype=dtype, constexpr=True, value=input_precision), ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/argsort.py b/src/ntops/kernels/argsort.py index a9a9c6d..9ec7dcb 100644 --- a/src/ntops/kernels/argsort.py +++ b/src/ntops/kernels/argsort.py @@ -5,13 +5,14 @@ from ntops.kernels.reduction import arrangement + def application(input, values, indices, dim_size, descending): val_block = input[0] # 初始化结果 buffer res_vals = ntl.zeros(val_block.shape, dtype=val_block.dtype) res_idxs = ntl.zeros(val_block.shape, dtype=indices.dtype.dtype) - + # 用于记录当前写入位置的索引 output_range = ntl.arange(0, val_block.shape[0]) # 原始数据的索引，用于比较 @@ -40,7 +41,7 @@ def application(input, values, indices, dim_size, descending): # 确定当前写入的位置 (target_mask 只有一个位置是 True) target_mask = output_range == i - + # 将找到的最值和索引写入结果 Tensor res_vals = ntl.where(target_mask, real_val, res_vals) res_idxs = ntl.where(target_mask, current_max_idx, res_idxs) @@ -65,11 +66,11 @@ def premake( pad_val = float("-inf") if descending else float("inf") tensors = ( - Tensor(ndim, dtype=dtype, other=pad_val), # Input - Tensor(ndim, dtype=dtype), # Output Values (辅助，虽然 argsort 主要要 indices) - Tensor(ndim, dtype=indices_dtype), # Output Indices - Tensor(0, constexpr=True, value=dim_size),# Loop bound (k=dim_size) + Tensor(ndim, dtype=dtype, other=pad_val), # Input + Tensor(ndim, dtype=dtype), # Output Values (辅助，虽然 argsort 主要要 indices) + Tensor(ndim, dtype=indices_dtype), # Output Indices + Tensor(0, constexpr=True, value=dim_size), # Loop bound (k=dim_size) Tensor(0, constexpr=True, value=descending), ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/atan.py b/src/ntops/kernels/atan.py index 00d84e2..1fe4552 100644 --- a/src/ntops/kernels/atan.py +++ b/src/ntops/kernels/atan.py @@ -13,7 +13,7 @@ def _atan_taylor_poly(x, dtype): 在此范围内，15阶多项式足以提供 float32/double 级别的精度。 """ x2 = x * x - + # 泰勒展开系数: 1, -1/3, 1/5, -1/7, ... # 为了精度，我们保留足够多的小数位 c3 = -0.333333333333 @@ -32,7 +32,7 @@ def _atan_taylor_poly(x, dtype): p = p * x2 + ntl.cast(c7, dtype) p = p * x2 + ntl.cast(c5, dtype) p = p * x2 + ntl.cast(c3, dtype) - + # result = x + x^3 * p = x * (1 + x^2 * p) # 提公因式 x 以减少一次乘法并提高小数值的稳定性 return x + x * x2 * p @@ -44,19 +44,19 @@ def _atan(x, dtype): 使用两级范围归约策略将输入映射到小区间，以保证多项式精度。 """ calc_dtype = dtype if dtype != ntl.float16 else ntl.float32 - + # === 常量定义 (局部定义以避开作用域问题) === PI_OVER_2 = 1.5707963267948966 PI_OVER_4 = 0.7853981633974483 - TAN_PI_8 = 0.4142135623730950 # tan(pi/8) - + TAN_PI_8 = 0.4142135623730950 # tan(pi/8) + x_arg = ntl.cast(x, calc_dtype) - + # 0. 提取符号并取绝对值 # atan(-x) = -atan(x) sign = ntl.where(x_arg < 0.0, -1.0, 1.0) abs_x = ntl.abs(x_arg) - + # 1. 第一级归约：处理 |x| > 1 的情况 # 使用恒等式：atan(x) = pi/2 - atan(1/x) (当 x > 0) # 如果 x > 1: @@ -69,67 +69,69 @@ def _atan(x, dtype): # coef_1 = 1 # 当前结果表达式: offset_1 + coef_1 * atan(val_1) mask_gt_1 = abs_x > 1.0 - + # 安全除法：防止 abs_x 为 0 时的除零错误（虽然此时分支不被选择） safe_abs_x = ntl.where(mask_gt_1, abs_x, ntl.cast(1.0, calc_dtype)) - + val_1 = ntl.where(mask_gt_1, ntl.cast(1.0, calc_dtype) / safe_abs_x, abs_x) offset_1 = ntl.where(mask_gt_1, PI_OVER_2, 0.0) coef_1 = ntl.where(mask_gt_1, -1.0, 1.0) - + # 此时 val_1 在 [0, 1] 范围内 - + # 2. 第二级归约：处理 x 接近 1 的情况 # 使用恒等式：atan(x) = pi/4 + atan((x-1)/(x+1)) # 阈值取 tan(pi/8) ≈ 0.414，这样归约后的值域在 [-0.29, 0.414] 之间 # 这对泰勒级数收敛非常有利 mask_gt_tan_pi_8 = val_1 > TAN_PI_8 - + # 计算 (x-1)/(x+1) # 注意：val_1 均为非负数，分母 val_1 + 1 永远 >= 1，无除零风险 reduced_val = (val_1 - 1.0) / (val_1 + 1.0) - + val_2 = ntl.where(mask_gt_tan_pi_8, reduced_val, val_1) offset_2 = ntl.where(mask_gt_tan_pi_8, PI_OVER_4, 0.0) - + # 当前结果表达式: atan(val_1) = offset_2 + atan(val_2) # 此时 |val_2| <= 0.4142... - + # 3. 多项式计算 poly_res = _atan_taylor_poly(val_2, calc_dtype) - + # 4. 结果组合 # result = sign * (offset_1 + coef_1 * (offset_2 + poly_res)) - + # 先计算内层: atan_val_1 atan_val_1 = ntl.cast(offset_2, calc_dtype) + poly_res - + # 再计算外层: abs_result - abs_result = ntl.cast(offset_1, calc_dtype) + ntl.cast(coef_1, calc_dtype) * atan_val_1 - + abs_result = ( + ntl.cast(offset_1, calc_dtype) + ntl.cast(coef_1, calc_dtype) * atan_val_1 + ) + # 最后恢复符号 final_result = ntl.cast(sign, calc_dtype) * abs_result - + return ntl.cast(final_result, dtype) def application(input, output): """ 计算反正切函数 atan(x) - + 参数: input: 输入张量，形状为 (C // block_size, block_size) output: 输出张量，形状为 (C // block_size, block_size) """ dtype = output.dtype.dtype - + for i in range(input.shape[0]): # 获取当前块的数据 input_block = ntl.cast(input[i], dtype) - + # 计算 atan(x) result = _atan(input_block, dtype) - + # 将结果存入输出 output[i] = result @@ -141,8 +143,8 @@ def premake(ndim, dim, dtype=None, block_size=None): arrangement_ = functools.partial(arrangement, dim=dim, block_size=block_size) tensors = ( - Tensor(ndim, dtype=dtype), # 输入张量 - Tensor(ndim, dtype=dtype), # 输出张量 + Tensor(ndim, dtype=dtype), # 输入张量 + Tensor(ndim, dtype=dtype), # 输出张量 ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/batch_norm.py b/src/ntops/kernels/batch_norm.py index 0df582a..a13f3c4 100644 --- a/src/ntops/kernels/batch_norm.py +++ b/src/ntops/kernels/batch_norm.py @@ -1,5 +1,4 @@ import functools -import math import ninetoothed.language as ntl from ninetoothed import Tensor @@ -10,7 +9,7 @@ def application(input, weight, bias, eps, output, num_normalized_elements): # 使用 E[x^2] - E[x]^2 公式计算方差，避免显式处理 Padding Mask # 因为 Padding 处 input 为 0，0 的平方也是 0，不会污染 sum 和 sum_sq - + _sum = ntl.zeros(input.dtype.shape, dtype=ntl.float32) _sum_sq = ntl.zeros(input.dtype.shape, dtype=ntl.float32) @@ -22,7 +21,7 @@ def application(input, weight, bias, eps, output, num_normalized_elements): mean = ntl.sum(_sum, 0) / num_normalized_elements mean_sq = ntl.sum(_sum_sq, 0) / num_normalized_elements - + # Var = E[x^2] - (E[x])^2 var = mean_sq - mean * mean # 确保方差非负 (处理数值误差) @@ -38,15 +37,17 @@ def application(input, weight, bias, eps, output, num_normalized_elements): def premake(ndim, reduction_dims, num_elements, dtype=None, block_size=None): # reduction_dims 指定了需要在哪些维度上进行规约 - arrangement_ = functools.partial(arrangement, dim=reduction_dims, block_size=block_size) + arrangement_ = functools.partial( + arrangement, dim=reduction_dims, block_size=block_size + ) tensors = ( - Tensor(ndim, other=0, dtype=dtype), # Input (other=0 确保 padding 读入 0) - Tensor(ndim, dtype=dtype), # Weight - Tensor(ndim, dtype=dtype), # Bias - Tensor(0, dtype=dtype), # eps - Tensor(ndim, dtype=dtype), # Output + Tensor(ndim, other=0, dtype=dtype), # Input (other=0 确保 padding 读入 0) + Tensor(ndim, dtype=dtype), # Weight + Tensor(ndim, dtype=dtype), # Bias + Tensor(0, dtype=dtype), # eps + Tensor(ndim, dtype=dtype), # Output Tensor(0, dtype=dtype, constexpr=True, value=num_elements), ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/bincount.py b/src/ntops/kernels/bincount.py index daff19c..8ec525e 100644 --- a/src/ntops/kernels/bincount.py +++ b/src/ntops/kernels/bincount.py @@ -1,6 +1,8 @@ import functools -from ninetoothed import Tensor + import ninetoothed.language as ntl +from ninetoothed import Tensor + def arrangement(input, weights, output, bin_ids, T, S, T_pow2, S_pow2, block_size=None): # input: (T,) @@ -11,14 +13,14 @@ def arrangement(input, weights, output, bin_ids, T, S, T_pow2, S_pow2, block_siz # 1. 对 Output 进行分块并行 (Grid 维度) # output_tiled: (GridSize, block_size) output_tiled = output.tile((block_size,)) - + # bin_ids 也随 output 一起分块，以便我们在 Kernel 中知道当前处理的是哪些 bin bin_ids_tiled = bin_ids.tile((block_size,)) # 2. Input 和 Weights 需要被所有 Block 访问 # 我们先将它们扩展到 GridSize，然后调整维度以符合 T_pow2 的要求 grid_size = output_tiled.shape[0] - + # input: (T,) -> (GridSize, T) -> (GridSize, 1, T_pow2) (通过 tile dim 1) # 注意：这里的 tile((1, T_pow2)) 是为了让 ninetoothed 框架正确处理内部维度 input_expand = input.unsqueeze(0).expand((grid_size, -1)) @@ -29,6 +31,7 @@ def arrangement(input, weights, output, bin_ids, T, S, T_pow2, S_pow2, block_siz return input_tiled, weights_tiled, output_tiled, bin_ids_tiled, T, S + def application(input, weights, output, bin_ids, T, S): # input: (1, T_pow2) <-- 来自 arrangement 的广播 # weights: (1, T_pow2) @@ -39,48 +42,55 @@ def application(input, weights, output, bin_ids, T, S): # 我们要计算矩阵: Match[i, j] = (bin_ids[i] == input[j]) # bin_ids: (block_size, 1) # input: (1, T_pow2) - - bin_ids_col = ntl.expand_dims(bin_ids, 1) # (block_size, 1) - + + bin_ids_col = ntl.expand_dims(bin_ids, 1) # (block_size, 1) + # 确保 input 和 weights 在 dim 0 上广播以匹配 block_size - input_b = ntl.broadcast_to(input, (bin_ids.shape[0], input.shape[1])) # (block_size, T_pow2) - weights_b = ntl.broadcast_to(weights, (bin_ids.shape[0], weights.shape[1])) # (block_size, T_pow2) + input_b = ntl.broadcast_to( + input, (bin_ids.shape[0], input.shape[1]) + ) # (block_size, T_pow2) + weights_b = ntl.broadcast_to( + weights, (bin_ids.shape[0], weights.shape[1]) + ) # (block_size, T_pow2) # 2. 生成有效性掩码 (处理 Padding) # input 的真实长度是 T，T_pow2 之外的是 padding - col_indices = ntl.arange(0, input.shape[1]) # (T_pow2,) - col_valid = col_indices < T # (T_pow2,) - + col_indices = ntl.arange(0, input.shape[1]) # (T_pow2,) + col_valid = col_indices < T # (T_pow2,) + # 广播 mask col_valid_b = ntl.expand_dims(col_valid, 0) col_valid_b = ntl.broadcast_to(col_valid_b, (bin_ids.shape[0], input.shape[1])) # 3. 核心计算：Masking + Sum # 找出哪些 input 值落入了当前 block 负责的 bin 中 - match_mask = (input_b == bin_ids_col) - + match_mask = input_b == bin_ids_col + # 结合有效性检查 final_mask = ntl.where(col_valid_b, match_mask, False) - + # 选择权重 (如果 weights 是全 1，则相当于计数) selected = ntl.where(final_mask, weights_b, 0.0) - + # 沿着 T 维度求和，得到每个 bin 的总值 - result = ntl.sum(selected, axis=1) # (block_size,) + result = ntl.sum(selected, axis=1) # (block_size,) # 4. 写回 - output = result + output = result # noqa: F841 + def premake(T_pow2, S_pow2, dtype=None, block_size=None): - arrangement_ = functools.partial(arrangement, T_pow2=T_pow2, S_pow2=S_pow2, block_size=block_size) + arrangement_ = functools.partial( + arrangement, T_pow2=T_pow2, S_pow2=S_pow2, block_size=block_size + ) tensors = ( - Tensor(1, dtype=int, shape_options={'constexpr': True}), # input - Tensor(1, dtype=dtype, shape_options={'constexpr': True}), # weights - Tensor(1, dtype=dtype, shape_options={'constexpr': True}), # output - Tensor(1, dtype=int, shape_options={'constexpr': True}), # bin_ids - Tensor(0, dtype=int, constexpr=True), # T - Tensor(0, dtype=int, constexpr=True), # S + Tensor(1, dtype=int, shape_options={"constexpr": True}), # input + Tensor(1, dtype=dtype, shape_options={"constexpr": True}), # weights + Tensor(1, dtype=dtype, shape_options={"constexpr": True}), # output + Tensor(1, dtype=int, shape_options={"constexpr": True}), # bin_ids + Tensor(0, dtype=int, constexpr=True), # T + Tensor(0, dtype=int, constexpr=True), # S ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/fmax.py b/src/ntops/kernels/fmax.py index 82de996..7640285 100644 --- a/src/ntops/kernels/fmax.py +++ b/src/ntops/kernels/fmax.py @@ -1,29 +1,33 @@ import functools + import ninetoothed import ninetoothed.language as ntl from ninetoothed import Tensor + def arrangement_elementwise(input, other, output, block_size=None): if block_size is None: block_size = ninetoothed.block_size() - + input = input.flatten().tile((block_size,)) other = other.flatten().tile((block_size,)) output = output.flatten().tile((block_size,)) - + return input, other, output + def application(input, other, output): # Compute element-wise maximum val = ntl.maximum(input, other) - + # Generate index range [0, block_size) to match the shape of 'val' # input.shape[0] retrieves the block_size of the tile indices = ntl.arange(0, input.shape[0]) - + # Write the vector 'val' to the memory locations defined by output + indices output[indices] = val + def premake(ndim, dtype=None, block_size=None): arrangement_ = functools.partial(arrangement_elementwise, block_size=block_size) tensors = ( @@ -31,4 +35,4 @@ def premake(ndim, dtype=None, block_size=None): Tensor(ndim, dtype=dtype), Tensor(ndim, dtype=dtype), ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/logsumexp.py b/src/ntops/kernels/logsumexp.py index 44f6903..d34bfe0 100644 --- a/src/ntops/kernels/logsumexp.py +++ b/src/ntops/kernels/logsumexp.py @@ -10,10 +10,12 @@ def _exp(x, dtype): exp_dtype = dtype if dtype != ntl.float16 else ntl.float32 return ntl.cast(ntl.exp(ntl.cast(x, exp_dtype)), dtype) + def _log(x, dtype): log_dtype = dtype if dtype != ntl.float16 else ntl.float32 return ntl.cast(ntl.log(ntl.cast(x, log_dtype)), dtype) + def application(input, output): # input&output: (C // block_size, ) # input.dtype: (block_size, ) diff --git a/src/ntops/kernels/lp_pool1d.py b/src/ntops/kernels/lp_pool1d.py index bec4beb..a15cbb7 100644 --- a/src/ntops/kernels/lp_pool1d.py +++ b/src/ntops/kernels/lp_pool1d.py @@ -2,18 +2,28 @@ import ninetoothed import ninetoothed.language as ntl -from ninetoothed.language import libdevice from ninetoothed import Tensor -from ninetoothed import Symbol +from ninetoothed.language import libdevice -def arrangement(input, output, norm_type, kernel_size_val, kernel_size, stride, block_size, ceil_mode): +def arrangement( + input, + output, + norm_type, + kernel_size_val, + kernel_size, + stride, + block_size, + ceil_mode, +): if block_size is None: block_size = ninetoothed.block_size() # input: (N, C, L_in) output: (N, C, L_out) - input_arranged = input.tile((1, 1, kernel_size), (1, 1, stride), floor_mode=not ceil_mode) + input_arranged = input.tile( + (1, 1, kernel_size), (1, 1, stride), floor_mode=not ceil_mode + ) # => (N, C, L_out), dtype=(1, 1, k) input_arranged = input_arranged.ravel() # => (N, C, L_out, 1, 1, k) @@ -30,15 +40,15 @@ def arrangement(input, output, norm_type, kernel_size_val, kernel_size, stride, input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) # => (..., 1), dtype=(block_size, nearest_pow2) - output_arranged = output.tile((1, 1, 1)) + output_arranged = output.tile((1, 1, 1)) # => (N, C, L_out), dtype=(1, 1, 1) output_arranged = output_arranged.ravel() # => (N, C, L_out, 1, 1, 1) output_arranged = output_arranged.flatten(end_dim=3).flatten(start_dim=1) # => (N*C*L_out, 1) - output_arranged = output_arranged.tile((block_size, -1)) + output_arranged = output_arranged.tile((block_size, -1)) # => (..., 1), dtype=(block_size, 1) - output_arranged.dtype = output_arranged.dtype.squeeze(1) + output_arranged.dtype = output_arranged.dtype.squeeze(1) # => (..., 1), dtype=(block_size, ) return input_arranged, output_arranged, norm_type, kernel_size_val @@ -48,6 +58,7 @@ def _pow(x, norm, dtype): pow_dtype = dtype if dtype != ntl.float16 else ntl.float32 return ntl.cast(libdevice.pow(ntl.cast(x, pow_dtype), norm), dtype) + def application(input, output, norm_type, kernel_size): # input: (block_size, nearest_pow2) # output: (block_size) @@ -57,7 +68,7 @@ def application(input, output, norm_type, kernel_size): input_masked = ntl.where(~mask, 0, input) x_pow = _pow(input_masked, norm_type, dtype) acc_sim = ntl.sum(x_pow, 1) / cnt * kernel_size - output = _pow(acc_sim, 1.0 / norm_type, dtype) + output = _pow(acc_sim, 1.0 / norm_type, dtype) # noqa: F841 def premake(ndim, kernel_size, stride, ceil_mode=False, dtype=None, block_size=None): @@ -72,8 +83,8 @@ def premake(ndim, kernel_size, stride, ceil_mode=False, dtype=None, block_size=N tensors = ( Tensor(ndim, dtype=dtype, other=float("inf")), # input Tensor(ndim, dtype=dtype), # output - Tensor(0, dtype=dtype), # norm_type - Tensor(0, dtype=dtype, constexpr=True), # kernel_size + Tensor(0, dtype=dtype), # norm_type + Tensor(0, dtype=dtype, constexpr=True), # kernel_size ) return arrangement_, application, tensors diff --git a/src/ntops/kernels/lp_pool2d.py b/src/ntops/kernels/lp_pool2d.py index 7bdbbf6..085ac90 100644 --- a/src/ntops/kernels/lp_pool2d.py +++ b/src/ntops/kernels/lp_pool2d.py @@ -2,16 +2,18 @@ import ninetoothed import ninetoothed.language as ntl -from ninetoothed.language import libdevice from ninetoothed import Tensor -from ninetoothed import Symbol +from ninetoothed.language import libdevice def _pow(x, norm, dtype): pow_dtype = dtype if dtype != ntl.float16 else ntl.float32 return ntl.cast(libdevice.pow(ntl.cast(x, pow_dtype), norm), dtype) -def arrangement_ceil_mode(*tensors, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size, ceil_mode): + +def arrangement_ceil_mode( + *tensors, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size, ceil_mode +): input, output, norm_type, kernel_size_flatted = tensors if block_size is None: block_size = ninetoothed.block_size() @@ -19,7 +21,11 @@ def arrangement_ceil_mode(*tensors, kernel_size_h, kernel_size_w, stride_h, stri # input: (N, C, H_in, W_in) output: (N, C, H_out, W_out) # ref. example 里的 max_pool2d arrangement - input_arranged = input.tile((1, 1, kernel_size_h, kernel_size_w), (1, 1, stride_h, stride_w), floor_mode=not ceil_mode) + input_arranged = input.tile( + (1, 1, kernel_size_h, kernel_size_w), + (1, 1, stride_h, stride_w), + floor_mode=not ceil_mode, + ) # => (N, C, H_out, W_out), dtype=(1, 1, k_h, k_w) input_arranged = input_arranged.ravel() # => (N, C, H_out, W_out, 1, 1, k_h, k_w) @@ -36,21 +42,20 @@ def arrangement_ceil_mode(*tensors, kernel_size_h, kernel_size_w, stride_h, stri input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) # => (..., 1), dtype=(block_size, nearest_pow2) - output_arranged = output.tile((1, 1, 1, 1)) + output_arranged = output.tile((1, 1, 1, 1)) # => (N, C, H_out, W_out), dtype=(1, 1, 1, 1) output_arranged = output_arranged.ravel() # => (N, C, H_out, W_out, 1, 1, 1, 1) output_arranged = output_arranged.flatten(end_dim=4).flatten(start_dim=1) # => (N*C*H_out*W_out, 1) - output_arranged = output_arranged.tile((block_size, -1)) + output_arranged = output_arranged.tile((block_size, -1)) # => (..., 1), dtype=(block_size, 1) - output_arranged.dtype = output_arranged.dtype.squeeze(1) + output_arranged.dtype = output_arranged.dtype.squeeze(1) # => (..., 1), dtype=(block_size, ) return input_arranged, output_arranged, norm_type, kernel_size_flatted - def application_ceil_mode(input, output, norm_type, kernel_size_flatted): # input: (block_size, nearest_pow2) arrangement 之后最外层被用于并行计算 # output: (block_size, ) @@ -72,10 +77,19 @@ def application_ceil_mode(input, output, norm_type, kernel_size_flatted): input_masked = ntl.where(~mask, 0, input) x_pow = _pow(input_masked, norm_type, dtype) acc_sim = ntl.sum(x_pow, 1) / cnt * kernel_size_flatted - output = _pow(acc_sim, 1.0 / norm_type, dtype) - - -def premake_ceil_mode(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, ceil_mode=False, dtype=None): + output = _pow(acc_sim, 1.0 / norm_type, dtype) # noqa: F841 + + +def premake_ceil_mode( + ndim, + kernel_size_h, + kernel_size_w, + stride_h, + stride_w, + block_size=None, + ceil_mode=False, + dtype=None, +): arrangement_ = functools.partial( arrangement_ceil_mode, kernel_size_h=kernel_size_h, @@ -87,24 +101,37 @@ def premake_ceil_mode(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, bl ) tensors = ( - Tensor(ndim, dtype=dtype, other=float("inf")), # input - Tensor(ndim, dtype=dtype), # output - Tensor(0, dtype=dtype), # norm_type - Tensor(0, dtype=dtype), # kernel_size_flatted + Tensor(ndim, dtype=dtype, other=float("inf")), # input + Tensor(ndim, dtype=dtype), # output + Tensor(0, dtype=dtype), # norm_type + Tensor(0, dtype=dtype), # kernel_size_flatted ) return arrangement_, application_ceil_mode, tensors - -def arrangement(input, output, norm_type, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size, ceil_mode): +def arrangement( + input, + output, + norm_type, + kernel_size_h, + kernel_size_w, + stride_h, + stride_w, + block_size, + ceil_mode, +): if block_size is None: block_size = ninetoothed.block_size() # input: (N, C, H_in, W_in) output: (N, C, H_out, W_out) # ref. example 里的 max_pool2d arrangement - input_arranged = input.tile((1, 1, kernel_size_h, kernel_size_w), (1, 1, stride_h, stride_w), floor_mode=not ceil_mode) + input_arranged = input.tile( + (1, 1, kernel_size_h, kernel_size_w), + (1, 1, stride_h, stride_w), + floor_mode=not ceil_mode, + ) # => (N, C, H_out, W_out), dtype=(1, 1, k_h, k_w) input_arranged = input_arranged.ravel() # => (N, C, H_out, W_out, 1, 1, k_h, k_w) @@ -121,29 +148,39 @@ def arrangement(input, output, norm_type, kernel_size_h, kernel_size_w, stride_h input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) # => (..., 1), dtype=(block_size, nearest_pow2) - output_arranged = output.tile((1, 1, 1, 1)) + output_arranged = output.tile((1, 1, 1, 1)) # => (N, C, H_out, W_out), dtype=(1, 1, 1, 1) output_arranged = output_arranged.ravel() # => (N, C, H_out, W_out, 1, 1, 1, 1) output_arranged = output_arranged.flatten(end_dim=4).flatten(start_dim=1) # => (N*C*H_out*W_out, 1) - output_arranged = output_arranged.tile((block_size, -1)) + output_arranged = output_arranged.tile((block_size, -1)) # => (..., 1), dtype=(block_size, 1) - output_arranged.dtype = output_arranged.dtype.squeeze(1) + output_arranged.dtype = output_arranged.dtype.squeeze(1) # => (..., 1), dtype=(block_size, ) return input_arranged, output_arranged, norm_type + def application(input, output, norm_type): # input: (block_size, nearest_pow2) # output: (block_size, ) dtype = input.dtype x_pow = _pow(input, norm_type, dtype) acc = ntl.sum(x_pow, axis=1) - output = _pow(acc, 1.0 / norm_type, dtype) - - -def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, ceil_mode=False, dtype=None): + output = _pow(acc, 1.0 / norm_type, dtype) # noqa: F841 + + +def premake( + ndim, + kernel_size_h, + kernel_size_w, + stride_h, + stride_w, + block_size=None, + ceil_mode=False, + dtype=None, +): arrangement_ = functools.partial( arrangement, kernel_size_h=kernel_size_h, @@ -155,9 +192,9 @@ def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=N ) tensors = ( - Tensor(ndim, dtype=dtype, other=0), # input - Tensor(ndim, dtype=dtype), # output - Tensor(0, dtype=dtype), # norm_type + Tensor(ndim, dtype=dtype, other=0), # input + Tensor(ndim, dtype=dtype), # output + Tensor(0, dtype=dtype), # norm_type ) return arrangement_, application, tensors diff --git a/src/ntops/kernels/lp_pool3d.py b/src/ntops/kernels/lp_pool3d.py index d872c01..a96f28c 100644 --- a/src/ntops/kernels/lp_pool3d.py +++ b/src/ntops/kernels/lp_pool3d.py @@ -2,15 +2,15 @@ import ninetoothed import ninetoothed.language as ntl -from ninetoothed.language import libdevice from ninetoothed import Tensor -from ninetoothed import Symbol +from ninetoothed.language import libdevice def _pow(x, norm, dtype): pow_dtype = dtype if dtype != ntl.float16 else ntl.float32 return ntl.cast(libdevice.pow(ntl.cast(x, pow_dtype), norm), dtype) + def arrangement_ceil_mode( *tensors, kernel_size_d, @@ -40,7 +40,6 @@ def arrangement_ceil_mode( return input_arranged, output_arranged, norm_type, kernel_size_flatted - def application_ceil_mode(input, output, norm_type, kernel_size_flatted): # input: (block_size, nearest_pow2) # output: (block_size, ) @@ -55,7 +54,7 @@ def application_ceil_mode(input, output, norm_type, kernel_size_flatted): # x_pow = _pow(input, norm_type, dtype) # acc = ntl.sum(x_pow, axis=0) # output = _pow(acc, 1.0 / norm_type, dtype) - + # 我把 ceil_mode 和普通的实现区分开来了 # 为了通过测试，下面使用的是与 torch 实现一致的版本 dtype = input.dtype @@ -64,10 +63,21 @@ def application_ceil_mode(input, output, norm_type, kernel_size_flatted): input_masked = ntl.where(~mask, 0, input) x_pow = _pow(input_masked, norm_type, dtype) acc_sim = ntl.sum(x_pow, 1) / cnt * kernel_size_flatted - output = _pow(acc_sim, 1.0 / norm_type, dtype) - + output = _pow(acc_sim, 1.0 / norm_type, dtype) # noqa: F841 + -def premake_ceil_mode(ndim, kernel_size_d, kernel_size_h, kernel_size_w, stride_d, stride_h, stride_w, block_size=None, ceil_mode=False, dtype=None): +def premake_ceil_mode( + ndim, + kernel_size_d, + kernel_size_h, + kernel_size_w, + stride_d, + stride_h, + stride_w, + block_size=None, + ceil_mode=False, + dtype=None, +): arrangement_ = functools.partial( arrangement_ceil_mode, kernel_size_d=kernel_size_d, @@ -81,24 +91,39 @@ def premake_ceil_mode(ndim, kernel_size_d, kernel_size_h, kernel_size_w, stride_ ) tensors = ( - Tensor(ndim, dtype=dtype, other=float("inf")), # input - Tensor(ndim, dtype=dtype), # output - Tensor(0, dtype=dtype), # norm_type - Tensor(0, dtype=dtype), # kernel_size_flatted + Tensor(ndim, dtype=dtype, other=float("inf")), # input + Tensor(ndim, dtype=dtype), # output + Tensor(0, dtype=dtype), # norm_type + Tensor(0, dtype=dtype), # kernel_size_flatted ) return arrangement_, application_ceil_mode, tensors - -def arrangement(input, output, norm_type, kernel_size_d, kernel_size_h, kernel_size_w, stride_d, stride_h, stride_w, block_size, ceil_mode): +def arrangement( + input, + output, + norm_type, + kernel_size_d, + kernel_size_h, + kernel_size_w, + stride_d, + stride_h, + stride_w, + block_size, + ceil_mode, +): if block_size is None: block_size = ninetoothed.block_size() # input: (N, C, D_in, H_in, W_in) output: (N, C, D_out, H_out, W_out) # ref. example 里的 max_pool2d arrangement - input_arranged = input.tile((1, 1, kernel_size_d, kernel_size_h, kernel_size_w), (1, 1, stride_d, stride_h, stride_w), floor_mode=not ceil_mode) + input_arranged = input.tile( + (1, 1, kernel_size_d, kernel_size_h, kernel_size_w), + (1, 1, stride_d, stride_h, stride_w), + floor_mode=not ceil_mode, + ) # => (N, C, D_out, H_out, W_out), dtype=(1, 1, k_d, k_h, k_w) input_arranged = input_arranged.ravel() # => (N, C, D_out, H_out, W_out, 1, 1, k_d, k_h, k_w) @@ -116,29 +141,41 @@ def arrangement(input, output, norm_type, kernel_size_d, kernel_size_h, kernel_s input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) # => (..., 1), dtype=(block_size, nearest_pow2) - output_arranged = output.tile((1, 1, 1, 1, 1)) + output_arranged = output.tile((1, 1, 1, 1, 1)) # => (N, C, D_out, H_out, W_out), dtype=(1, 1, 1, 1, 1) output_arranged = output_arranged.ravel() # => (N, C, D_out, H_out, W_out, 1, 1, 1, 1) output_arranged = output_arranged.flatten(end_dim=5).flatten(start_dim=1) # => (N*C*D_out*H_out*W_out, 1) - output_arranged = output_arranged.tile((block_size, -1)) + output_arranged = output_arranged.tile((block_size, -1)) # => (..., 1), dtype=(block_size, 1) - output_arranged.dtype = output_arranged.dtype.squeeze(1) + output_arranged.dtype = output_arranged.dtype.squeeze(1) # => (..., 1), dtype=(block_size, ) return input_arranged, output_arranged, norm_type + def application(input, output, norm_type): # input: (block_size, nearest_pow2) # output: (block_size, ) dtype = input.dtype x_pow = _pow(input, norm_type, dtype) acc = ntl.sum(x_pow, axis=1) - output = _pow(acc, 1.0 / norm_type, dtype) + output = _pow(acc, 1.0 / norm_type, dtype) # noqa: F841 -def premake(ndim, kernel_size_d, kernel_size_h, kernel_size_w, stride_d, stride_h, stride_w, block_size=None, ceil_mode=False, dtype=None): +def premake( + ndim, + kernel_size_d, + kernel_size_h, + kernel_size_w, + stride_d, + stride_h, + stride_w, + block_size=None, + ceil_mode=False, + dtype=None, +): arrangement_ = functools.partial( arrangement, kernel_size_d=kernel_size_d, @@ -152,9 +189,9 @@ def premake(ndim, kernel_size_d, kernel_size_h, kernel_size_w, stride_d, stride_ ) tensors = ( - Tensor(ndim, dtype=dtype, other=0), # input - Tensor(ndim, dtype=dtype), # output - Tensor(0, dtype=dtype), # norm_type + Tensor(ndim, dtype=dtype, other=0), # input + Tensor(ndim, dtype=dtype), # output + Tensor(0, dtype=dtype), # norm_type ) return arrangement_, application, tensors diff --git a/src/ntops/kernels/max.py b/src/ntops/kernels/max.py index 84c43ae..134db40 100644 --- a/src/ntops/kernels/max.py +++ b/src/ntops/kernels/max.py @@ -2,7 +2,6 @@ import ninetoothed import ninetoothed.language as ntl -import math from ninetoothed import Tensor from ntops.kernels.reduction import arrangement @@ -34,11 +33,12 @@ def premake(ndim, dim, dtype=None, block_size=None): ndim, dtype=dtype, other=float("-inf"), shape_options={"constexpr": True} ), Tensor(ndim, dtype=dtype), - Tensor(ndim, dtype=ninetoothed.int32) + Tensor(ndim, dtype=ninetoothed.int32), ) return arrangement_, application, tensors + def arrangement_all_elements(input, output, block_size=None): if block_size is None: block_size = ninetoothed.block_size() @@ -47,14 +47,18 @@ def arrangement_all_elements(input, output, block_size=None): output = output.tile((1,)) return input, output + def application_all_elements(input, output): output[0] = ntl.max(input, 0) + def premake_all_elements(ndim, dtype=None, block_size=None): arrangement_ = functools.partial(arrangement_all_elements, block_size=block_size) tensors = ( - Tensor(ndim, dtype=dtype, other=float("-inf"), shape_options={"constexpr": True}), + Tensor( + ndim, dtype=dtype, other=float("-inf"), shape_options={"constexpr": True} + ), Tensor(1, dtype=dtype), ) diff --git a/src/ntops/kernels/max_pool1d.py b/src/ntops/kernels/max_pool1d.py index 0d1d881..1d25422 100644 --- a/src/ntops/kernels/max_pool1d.py +++ b/src/ntops/kernels/max_pool1d.py @@ -1,13 +1,15 @@ import functools + import ninetoothed import ninetoothed.language as ntl from ninetoothed import Tensor + def arrangement(input, output, kernel_size, stride, block_size, ceil_mode): if block_size is None: block_size = ninetoothed.block_size() - # input: (N, C, L_in) + # input: (N, C, L_in) # output: (N, C, L_out) # 1. 窗口切分 @@ -15,16 +17,14 @@ def arrangement(input, output, kernel_size, stride, block_size, ceil_mode): # strides: (1, 1, stride) -> 在 L 维度上步长为 stride # floor_mode=not ceil_mode: 决定是否丢弃最后不足一个 kernel 的部分 input_arranged = input.tile( - (1, 1, kernel_size), - (1, 1, stride), - floor_mode=not ceil_mode + (1, 1, kernel_size), (1, 1, stride), floor_mode=not ceil_mode ) # => (N, C, L_out), dtype=(1, 1, k) # 2. 展平与重排 input_arranged = input_arranged.ravel() # => (N, C, L_out, 1, 1, k) - + input_arranged = input_arranged.flatten(end_dim=3).flatten(start_dim=1) # => (N*C*L_out, k) @@ -33,14 +33,14 @@ def arrangement(input, output, kernel_size, stride, block_size, ceil_mode): nearest_pow2 = 1 << (kernel_size - 1).bit_length() input_arranged = input_arranged.tile((1, nearest_pow2)) # => (..., 1), dtype=(1, nearest_pow2) - + input_arranged.dtype = input_arranged.dtype.squeeze(0) input_arranged = input_arranged.tile((block_size, -1)) input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) # => (..., 1), dtype=(block_size, nearest_pow2) # 4. Output 对齐 - output_arranged = output.tile((1, 1, 1)) # (N, C, L_out) + output_arranged = output.tile((1, 1, 1)) # (N, C, L_out) output_arranged = output_arranged.ravel() output_arranged = output_arranged.flatten(end_dim=3).flatten(start_dim=1) output_arranged = output_arranged.tile((block_size, -1)) @@ -48,12 +48,14 @@ def arrangement(input, output, kernel_size, stride, block_size, ceil_mode): return input_arranged, output_arranged + def application(input, output): # input: (block_size, nearest_pow2) # output: (block_size, ) - + # 直接取 Max，padding 值为 -inf，不影响结果 - output = ntl.max(input, axis=1) + output = ntl.max(input, axis=1) # noqa: F841 + def premake(ndim, kernel_size, stride, block_size=None, ceil_mode=False, dtype=None): arrangement_ = functools.partial( @@ -66,8 +68,8 @@ def premake(ndim, kernel_size, stride, block_size=None, ceil_mode=False, dtype=N tensors = ( # input: MaxPool 填充负无穷 - Tensor(ndim, dtype=dtype, other=float("-inf")), - Tensor(ndim, dtype=dtype), # output + Tensor(ndim, dtype=dtype, other=float("-inf")), + Tensor(ndim, dtype=dtype), # output ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/max_pool3d.py b/src/ntops/kernels/max_pool3d.py index 26aab33..d24eb28 100644 --- a/src/ntops/kernels/max_pool3d.py +++ b/src/ntops/kernels/max_pool3d.py @@ -1,27 +1,27 @@ import functools + import ninetoothed import ninetoothed.language as ntl from ninetoothed import Tensor + def arrangement(input, output, k_d, k_h, k_w, s_d, s_h, s_w, block_size, ceil_mode): if block_size is None: block_size = ninetoothed.block_size() - # input: (N, C, D_in, H_in, W_in) + # input: (N, C, D_in, H_in, W_in) # output: (N, C, D_out, H_out, W_out) # 1. 窗口切分 (增加 Depth 维度) input_arranged = input.tile( - (1, 1, k_d, k_h, k_w), - (1, 1, s_d, s_h, s_w), - floor_mode=not ceil_mode + (1, 1, k_d, k_h, k_w), (1, 1, s_d, s_h, s_w), floor_mode=not ceil_mode ) # => (N, C, D_out, H_out, W_out), dtype=(1, 1, k_d, k_h, k_w) # 2. 展平与重排 input_arranged = input_arranged.ravel() # => (N, C, D_out, H_out, W_out, 1, 1, k_d, k_h, k_w) - + # 注意：这里 end_dim=5，因为前面有 N,C,D,H,W 5个维度需要合并作为 batch 处理 input_arranged = input_arranged.flatten(end_dim=5).flatten(start_dim=1) # => (N*C*D_out*H_out*W_out, k_d*k_h*k_w) @@ -31,14 +31,14 @@ def arrangement(input, output, k_d, k_h, k_w, s_d, s_h, s_w, block_size, ceil_mo nearest_pow2 = 1 << (k_d * k_h * k_w - 1).bit_length() input_arranged = input_arranged.tile((1, nearest_pow2)) # => (..., 1), dtype=(1, nearest_pow2) - + input_arranged.dtype = input_arranged.dtype.squeeze(0) input_arranged = input_arranged.tile((block_size, -1)) input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1) # => (..., 1), dtype=(block_size, nearest_pow2) # 4. Output 对齐 - output_arranged = output.tile((1, 1, 1, 1, 1)) # (N, C, D, H, W) + output_arranged = output.tile((1, 1, 1, 1, 1)) # (N, C, D, H, W) output_arranged = output_arranged.ravel() output_arranged = output_arranged.flatten(end_dim=5).flatten(start_dim=1) output_arranged = output_arranged.tile((block_size, -1)) @@ -46,26 +46,34 @@ def arrangement(input, output, k_d, k_h, k_w, s_d, s_h, s_w, block_size, ceil_mo return input_arranged, output_arranged + def application(input, output): # input: (block_size, nearest_pow2) # output: (block_size, ) - + # Max Pooling 标准操作 - output = ntl.max(input, axis=1) + output = ntl.max(input, axis=1) # noqa: F841 + -def premake(ndim, k_d, k_h, k_w, s_d, s_h, s_w, block_size=None, ceil_mode=False, dtype=None): +def premake( + ndim, k_d, k_h, k_w, s_d, s_h, s_w, block_size=None, ceil_mode=False, dtype=None +): arrangement_ = functools.partial( arrangement, - k_d=k_d, k_h=k_h, k_w=k_w, - s_d=s_d, s_h=s_h, s_w=s_w, + k_d=k_d, + k_h=k_h, + k_w=k_w, + s_d=s_d, + s_h=s_h, + s_w=s_w, block_size=block_size, ceil_mode=ceil_mode, ) tensors = ( # input: MaxPool 填充负无穷 - Tensor(ndim, dtype=dtype, other=float("-inf")), - Tensor(ndim, dtype=dtype), # output + Tensor(ndim, dtype=dtype, other=float("-inf")), + Tensor(ndim, dtype=dtype), # output ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/maximum.py b/src/ntops/kernels/maximum.py index db85057..d92b42f 100644 --- a/src/ntops/kernels/maximum.py +++ b/src/ntops/kernels/maximum.py @@ -1,34 +1,38 @@ import functools + import ninetoothed import ninetoothed.language as ntl from ninetoothed import Tensor + def arrangement_elementwise(input, other, output, block_size=None): if block_size is None: block_size = ninetoothed.block_size() - + # 逐元素操作的核心策略：将多维 Tensor 视为展平的 1D 数组 # 使用 tile 切分数据块 input = input.flatten().tile((block_size,)) other = other.flatten().tile((block_size,)) output = output.flatten().tile((block_size,)) - + return input, other, output + def application(input, other, output): # 调用 DSL 的 maximum 原语 # 注意：在 Triton/DSL 中，maximum 的 NaN 行为取决于底层实现 val = ntl.maximum(input, other) - + # 生成索引并写回 indices = ntl.arange(0, input.shape[0]) output[indices] = val + def premake(ndim, dtype=None, block_size=None): arrangement_ = functools.partial(arrangement_elementwise, block_size=block_size) tensors = ( - Tensor(ndim, dtype=dtype), # input - Tensor(ndim, dtype=dtype), # other - Tensor(ndim, dtype=dtype), # output + Tensor(ndim, dtype=dtype), # input + Tensor(ndim, dtype=dtype), # other + Tensor(ndim, dtype=dtype), # output ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/mean.py b/src/ntops/kernels/mean.py index 9a0bcf7..b706df8 100644 --- a/src/ntops/kernels/mean.py +++ b/src/ntops/kernels/mean.py @@ -1,8 +1,11 @@ import functools + import ninetoothed.language as ntl from ninetoothed import Tensor + # 假设 arrangement 逻辑是通用的，直接复用你例子中的 reduction 模块 -from ntops.kernels.reduction import arrangement +from ntops.kernels.reduction import arrangement + def application(input, output): # 均值计算的第一步是求和。 @@ -14,32 +17,37 @@ def application(input, output): # 将累加结果转换为输出类型（通常是 float） output[0] = ntl.cast(accumulator, output.dtype.dtype) + def premake(ndim, dim, dtype=None, block_size=None): arrangement_ = functools.partial(arrangement, dim=dim, block_size=block_size) - + # Mean 算子要求输出必须是浮点数。 # 如果传入的 dtype 是整数，这里需要根据 DSL 的特性处理， # 这里假设 dtype 已经是转换好的 float 类型 (由 Torch 端传入) tensors = ( - Tensor(ndim, dtype=dtype), # Input - Tensor(ndim, dtype=dtype), # Output + Tensor(ndim, dtype=dtype), # Input + Tensor(ndim, dtype=dtype), # Output ) return arrangement_, application, tensors + # --- Global Reduction (All Elements) 的实现 --- + def arrangement_all_elements(input, output, block_size=None): input = input.flatten().tile((block_size,)) output = output.tile((1,)) return input, output + def application_all_elements(input, output): output[0] = ntl.sum(input, 0) + def premake_all_elements(ndim, dtype=None, block_size=None): arrangement_ = functools.partial(arrangement_all_elements, block_size=block_size) tensors = ( Tensor(ndim, dtype=dtype), Tensor(1, dtype=dtype), ) - return arrangement_, application_all_elements, tensors \ No newline at end of file + return arrangement_, application_all_elements, tensors diff --git a/src/ntops/kernels/median.py b/src/ntops/kernels/median.py index 871b42f..d25249a 100644 --- a/src/ntops/kernels/median.py +++ b/src/ntops/kernels/median.py @@ -1,34 +1,37 @@ import functools -import triton.language as tl # [FIX 1] 导入标准 triton language + import ninetoothed.language as ntl +import triton.language as tl # [FIX 1] 导入标准 triton language from ninetoothed import Tensor + from ntops.kernels.reduction import arrangement + def application(input, values, indices, loop_k): val_block = input[0] # 初始化用于查找的 working_val working_val = -val_block sentinel = float("-inf") - + # 原始索引 idx_block = ntl.arange(0, val_block.shape[0]) # 初始化结果 buffer (Scalar) final_val = ntl.zeros([], dtype=val_block.dtype) - + # [FIX 2] 使用 tl.int32 而不是字符串 "int32" - final_idx = ntl.zeros([], dtype=tl.int32) + final_idx = ntl.zeros([], dtype=tl.int32) # 循环 for i in range(loop_k + 1): current_max_val = ntl.max(working_val, axis=0) - current_max_idx = ntl.argmax(working_val, axis=0) # 返回 int32 + current_max_idx = ntl.argmax(working_val, axis=0) # 返回 int32 if i == loop_k: real_val = -current_max_val final_val = ntl.cast(real_val, values.dtype.dtype) - final_idx = current_max_idx # int32 -> int32 + final_idx = current_max_idx # int32 -> int32 # Mask 逻辑 mask_selected = idx_block == current_max_idx @@ -40,17 +43,16 @@ def application(input, values, indices, loop_k): # Cast int32 -> int64 (或输出需要的类型) indices[0] = ntl.cast(final_idx, indices.dtype.dtype) -def premake( - ndim, dim, loop_k, dtype=None, indices_dtype=None, block_size=None -): + +def premake(ndim, dim, loop_k, dtype=None, indices_dtype=None, block_size=None): arrangement_ = functools.partial(arrangement, dim=dim, block_size=block_size) pad_val = float("inf") tensors = ( - Tensor(ndim, dtype=dtype, other=pad_val), # Input - Tensor(ndim, dtype=dtype), # Output Values - Tensor(ndim, dtype=indices_dtype), # Output Indices - Tensor(0, constexpr=True, value=loop_k), # Loop bound + Tensor(ndim, dtype=dtype, other=pad_val), # Input + Tensor(ndim, dtype=dtype), # Output Values + Tensor(ndim, dtype=indices_dtype), # Output Indices + Tensor(0, constexpr=True, value=loop_k), # Loop bound ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/kernels/sort.py b/src/ntops/kernels/sort.py index 53d7e18..61e3829 100644 --- a/src/ntops/kernels/sort.py +++ b/src/ntops/kernels/sort.py @@ -30,7 +30,9 @@ def application(input, values, indices, sort_size, descending): encoded = ntl.where(valid, encoded, ntl.cast(0x7FFFFFFF, ntl.int32)) offsets = ntl.cast(offsets, ntl.int64) - key = ((ntl.cast(encoded, ntl.int64) & ntl.cast(0xFFFFFFFF, ntl.int64)) << 32) | offsets + key = ( + (ntl.cast(encoded, ntl.int64) & ntl.cast(0xFFFFFFFF, ntl.int64)) << 32 + ) | offsets sorted_key = ntl.sort(key) sorted_encoded = ntl.cast(sorted_key >> 32, ntl.int32) diff --git a/src/ntops/kernels/stack.py b/src/ntops/kernels/stack.py index 60676c4..f7f6636 100644 --- a/src/ntops/kernels/stack.py +++ b/src/ntops/kernels/stack.py @@ -1,21 +1,22 @@ import functools + import ninetoothed -import ninetoothed.language as ntl from ninetoothed import Tensor + def arrangement(input, output, block_size): if block_size is None: block_size = ninetoothed.block_size() # input: (...,) 任意维度的单个切片 # output: (...,) 对应的输出切片 - + # 1. 展平 (Flatten) # Stack/Copy 操作是 Element-wise 的，不依赖空间结构， # 所以我们可以将数据视为一维连续内存进行处理，以获得最大内存带宽。 input_arranged = input.flatten() output_arranged = output.flatten() - + # 2. 切分 (Tile) # 将一维数据切分为大小为 block_size 的块 # 形状变化: (Total_Elements) -> (Num_Blocks, Block_Size) @@ -24,13 +25,15 @@ def arrangement(input, output, block_size): return input_arranged, output_arranged + def application(input, output): # input: (block_size, ) # output: (block_size, ) - + # 简单的 Element-wise 赋值 # DSL 会将其翻译为 load(input) -> store(output) - output = input + output = input # noqa: F841 + def premake(ndim, block_size=None, dtype=None): arrangement_ = functools.partial( @@ -43,4 +46,4 @@ def premake(ndim, block_size=None, dtype=None): Tensor(ndim, dtype=dtype), # output (destination slice) ) - return arrangement_, application, tensors \ No newline at end of file + return arrangement_, application, tensors diff --git a/src/ntops/torch/__init__.py b/src/ntops/torch/__init__.py index ad6fd4c..0903971 100644 --- a/src/ntops/torch/__init__.py +++ b/src/ntops/torch/__init__.py @@ -1,8 +1,16 @@ from ntops.torch.abs import abs +from ntops.torch.acosh import acosh +from ntops.torch.adaptive_avg_pool2d import adaptive_avg_pool2d +from ntops.torch.adaptive_max_pool2d import adaptive_max_pool2d from ntops.torch.add import add from ntops.torch.addmm import addmm +from ntops.torch.addmv import addmv from ntops.torch.alpha_dropout import alpha_dropout +from ntops.torch.argsort import argsort +from ntops.torch.atan import atan from ntops.torch.avg_pool2d import avg_pool2d +from ntops.torch.batch_norm import batch_norm +from ntops.torch.bincount import bincount from ntops.torch.bitwise_and import bitwise_and from ntops.torch.bitwise_not import bitwise_not from ntops.torch.bitwise_or import bitwise_or @@ -17,6 +25,7 @@ from ntops.torch.dropout import dropout from ntops.torch.eq import eq from ntops.torch.exp import exp +from ntops.torch.fmax import fmax from ntops.torch.ge import ge from ntops.torch.gelu import gelu from ntops.torch.gt import gt @@ -25,9 +34,19 @@ from ntops.torch.isnan import isnan from ntops.torch.layer_norm import layer_norm from ntops.torch.le import le +from ntops.torch.logsumexp import logsumexp +from ntops.torch.lp_pool1d import lp_pool1d +from ntops.torch.lp_pool2d import lp_pool2d +from ntops.torch.lp_pool3d import lp_pool3d from ntops.torch.lt import lt from ntops.torch.matmul import matmul +from ntops.torch.max import max +from ntops.torch.max_pool1d import max_pool1d from ntops.torch.max_pool2d import max_pool2d +from ntops.torch.max_pool3d import max_pool3d +from ntops.torch.maximum import maximum +from ntops.torch.mean import mean +from ntops.torch.median import median from ntops.torch.mm import mm from ntops.torch.msort import msort from ntops.torch.mul import mul @@ -51,29 +70,10 @@ from ntops.torch.sin import sin from ntops.torch.softmax import softmax from ntops.torch.sort import sort +from ntops.torch.stack import stack from ntops.torch.sub import sub from ntops.torch.tanh import tanh from ntops.torch.threshold import threshold -from ntops.torch.max_pool1d import max_pool1d -from ntops.torch.max_pool3d import max_pool3d -from ntops.torch.stack import stack -from ntops.torch.mean import mean -from ntops.torch.median import median -from ntops.torch.maximum import maximum -from ntops.torch.atan import atan -from ntops.torch.batch_norm import batch_norm -from ntops.torch.bincount import bincount -from ntops.torch.adaptive_max_pool2d import adaptive_max_pool2d -from ntops.torch.acosh import acosh -from ntops.torch.adaptive_avg_pool2d import adaptive_avg_pool2d -from ntops.torch.addmv import addmv -from ntops.torch.argsort import argsort -from ntops.torch.fmax import fmax -from ntops.torch.logsumexp import logsumexp -from ntops.torch.lp_pool1d import lp_pool1d -from ntops.torch.lp_pool2d import lp_pool2d -from ntops.torch.lp_pool3d import lp_pool3d -from ntops.torch.max import max __all__ = [ "abs", @@ -130,7 +130,7 @@ "softmax", "sort", "sub", - "tanh", + "tanh", "max_pool1d", "max_pool3d", "stack", diff --git a/src/ntops/torch/acosh.py b/src/ntops/torch/acosh.py index 9e7cf3b..9659aee 100644 --- a/src/ntops/torch/acosh.py +++ b/src/ntops/torch/acosh.py @@ -7,14 +7,14 @@ def acosh(input, *, out=None): """ 计算反双曲余弦函数 acosh(x) = ln(x + sqrt(x² - 1)) - + 参数: input: 输入张量 out: 可选的输出张量，如果提供，结果将存储在此张量中 - + 返回: 计算acosh后的张量 - + 数值稳定性: - 对于x < 1的值，返回NaN - 对于x = 1的值，返回0 @@ -23,39 +23,43 @@ def acosh(input, *, out=None): """ # 确定输出数据类型 tensor_dtype = out.dtype if out is not None else input.dtype - + # 创建临时输出张量 temp_out = torch.empty_like(input, dtype=tensor_dtype, device=input.device) - + # 设置块大小，优化GPU/TPU性能 block_size = 256 - + # 缓存并获取acosh内核 # 注意：acosh是element-wise操作，不需要指定归约维度 # 使用input.ndim作为虚拟维度参数，保持接口一致性 kernel = _cached_make( - ntops.kernels.acosh.premake, + ntops.kernels.acosh.premake, input.ndim, # ndim - 0, # dummy dim (element-wise operation doesn't reduce any dimension) + 0, # dummy dim (element-wise operation doesn't reduce any dimension) dtype=input.dtype, - block_size=block_size + block_size=block_size, ) - + # 执行内核计算 kernel(input, temp_out) - + # 处理out参数 if out is not None: # 确保out张量的形状和数据类型与结果匹配 if out.shape != temp_out.shape: - raise RuntimeError(f"Expected out tensor to have shape {temp_out.shape}, but got {out.shape}") + raise RuntimeError( + f"Expected out tensor to have shape {temp_out.shape}, but got {out.shape}" + ) if out.dtype != temp_out.dtype: - raise RuntimeError(f"Expected out tensor to have dtype {temp_out.dtype}, but got {out.dtype}") - + raise RuntimeError( + f"Expected out tensor to have dtype {temp_out.dtype}, but got {out.dtype}" + ) + # 复制结果到out张量 out.copy_(temp_out) return out - + return temp_out @@ -63,6 +67,7 @@ def acosh(input, *, out=None): def _acosh_tensor_method(self, *, out=None): return acosh(self, out=out) + # 注册到PyTorch Tensor torch.Tensor.acosh = _acosh_tensor_method @@ -72,6 +77,7 @@ def arcosh(input, *, out=None): """acosh的别名，为了与某些数学库保持一致""" return acosh(input, out=out) + def arccosh(input, *, out=None): """acosh的别名，为了与NumPy等库保持一致""" - return acosh(input, out=out) \ No newline at end of file + return acosh(input, out=out) diff --git a/src/ntops/torch/adaptive_avg_pool2d.py b/src/ntops/torch/adaptive_avg_pool2d.py index 2ab800f..81ca939 100644 --- a/src/ntops/torch/adaptive_avg_pool2d.py +++ b/src/ntops/torch/adaptive_avg_pool2d.py @@ -1,32 +1,36 @@ import torch + import ntops from ntops.torch.utils import _cached_make + def adaptive_avg_pool2d(input, output_size): - assert input.ndim == 4 or input.ndim == 3, "Input tensor must be 4-dimensional (N, C, H, W) or 3-dimensional (C, H, W)" - + assert input.ndim == 4 or input.ndim == 3, ( + "Input tensor must be 4-dimensional (N, C, H, W) or 3-dimensional (C, H, W)" + ) + if input.ndim == 3: input = input.unsqueeze(0) - + if isinstance(output_size, int): output_size = (output_size, output_size) - + H_in, W_in = input.shape[-2], input.shape[-1] H_out, W_out = output_size - + # 计算 stride 和 kernel_size (固定窗口策略) # 对于 Input % Output == 0 的情况，这与 PyTorch 行为完全一致 stride_h = H_in // H_out stride_w = W_in // W_out - + kernel_h = H_in - (H_out - 1) * stride_h kernel_w = W_in - (W_out - 1) * stride_w - + output_shape = (input.shape[0], input.shape[1], H_out, W_out) output = torch.empty(output_shape, dtype=input.dtype, device=input.device) - + block_size = 1024 - + kernel = _cached_make( ntops.kernels.adaptive_avg_pool2d.premake, input.ndim, @@ -35,11 +39,11 @@ def adaptive_avg_pool2d(input, output_size): stride_h, stride_w, block_size=block_size, - dtype=input.dtype + dtype=input.dtype, ) - + # 传入实际的窗口面积作为除数 area = kernel_h * kernel_w kernel(input, output, area) - - return output \ No newline at end of file + + return output diff --git a/src/ntops/torch/adaptive_max_pool2d.py b/src/ntops/torch/adaptive_max_pool2d.py index 79098a7..9bc7a56 100644 --- a/src/ntops/torch/adaptive_max_pool2d.py +++ b/src/ntops/torch/adaptive_max_pool2d.py @@ -1,37 +1,41 @@ import torch + import ntops from ntops.torch.utils import _cached_make + def adaptive_max_pool2d(input, output_size): - assert input.ndim == 4 or input.ndim == 3, "Input tensor must be 4-dimensional (N, C, H, W) or 3-dimensional (C, H, W)" - + assert input.ndim == 4 or input.ndim == 3, ( + "Input tensor must be 4-dimensional (N, C, H, W) or 3-dimensional (C, H, W)" + ) + if input.ndim == 3: input = input.unsqueeze(0) - + if isinstance(output_size, int): output_size = (output_size, output_size) - + H_in, W_in = input.shape[-2], input.shape[-1] H_out, W_out = output_size - + # 计算 stride 和 kernel_size # 逻辑: stride = input // output # kernel = input - (output - 1) * stride # 注意：这种固定 stride/kernel 的方式要求 input 大小比较规整， # 对于 PyTorch 完全动态的 adaptive pool (特别是不可整除的情况)， # DSL 需要支持动态窗口才能完全对齐。这里实现的是基于固定窗口的近似/通用实现。 - + stride_h = H_in // H_out stride_w = W_in // W_out - + kernel_h = H_in - (H_out - 1) * stride_h kernel_w = W_in - (W_out - 1) * stride_w - + output_shape = (input.shape[0], input.shape[1], H_out, W_out) output = torch.empty(output_shape, dtype=input.dtype, device=input.device) - + block_size = 1024 - + kernel = _cached_make( ntops.kernels.adaptive_max_pool2d.premake, input.ndim, @@ -40,9 +44,9 @@ def adaptive_max_pool2d(input, output_size): stride_h, stride_w, block_size=block_size, - dtype=input.dtype + dtype=input.dtype, ) - + kernel(input, output) - - return output \ No newline at end of file + + return output diff --git a/src/ntops/torch/addmv.py b/src/ntops/torch/addmv.py index 300b481..757da85 100644 --- a/src/ntops/torch/addmv.py +++ b/src/ntops/torch/addmv.py @@ -6,7 +6,7 @@ def addmv(input, mat, vec, *, beta=1, alpha=1, out=None): m, k = mat.shape - + if vec.shape[0] != k: raise RuntimeError(f"size mismatch, got {vec.shape[0]}, expected {k}") @@ -24,4 +24,4 @@ def addmv(input, mat, vec, *, beta=1, alpha=1, out=None): # 传入 2D view kernel(input_2d, mat, vec_2d, beta, alpha, out_2d, _get_matmul_input_precision()) - return out \ No newline at end of file + return out diff --git a/src/ntops/torch/argsort.py b/src/ntops/torch/argsort.py index 97181e0..c26f04b 100644 --- a/src/ntops/torch/argsort.py +++ b/src/ntops/torch/argsort.py @@ -3,11 +3,13 @@ import ntops from ntops.torch.utils import _cached_make + def next_power_of_2(n): if n == 0: return 1 return 1 << (n - 1).bit_length() + def get_optimal_block_size(dim_size): target_size = next_power_of_2(dim_size) if target_size > 1024: @@ -16,6 +18,7 @@ def get_optimal_block_size(dim_size): target_size = 32 return target_size + def argsort(input, dim=-1, descending=False, *, out=None): """ Args: @@ -31,16 +34,16 @@ def argsort(input, dim=-1, descending=False, *, out=None): if dim is None: # 如果 dim 为 None，torch 通常会 flatten 后排序，这里为了简化暂不支持或视为 dim=0 raise NotImplementedError("argsort currently requires a specific dim") - + if dim < 0: dim += input.ndim - + target_dim = dim dim_size = input.shape[target_dim] - + # 逻辑 Input (无需 flatten，直接作为 ndim tensor 处理，依赖 arrangement 划分) input_logic = input - + # 计算 Block Size block_size = get_optimal_block_size(dim_size) @@ -48,14 +51,14 @@ def argsort(input, dim=-1, descending=False, *, out=None): # argsort 只需要返回 indices，但我们的 kernel 同时也计算 values (类似 sort) # 我们可以分配一个临时的 values tensor output_shape = list(input.shape) - + if out is not None: indices_logic = out else: indices_logic = torch.empty( output_shape, dtype=indices_dtype, device=input.device ) - + # 临时存储排序后的值，Kernel 需要用到它来回写 values_logic = torch.empty(output_shape, dtype=dtype, device=input.device) @@ -65,8 +68,8 @@ def argsort(input, dim=-1, descending=False, *, out=None): ntops.kernels.argsort.premake, input_logic.ndim, target_dim, - dim_size, # 这里 k = dim_size - descending, # argsort 默认 ascending (False) + dim_size, # 这里 k = dim_size + descending, # argsort 默认 ascending (False) dtype, indices_dtype, block_size, @@ -75,4 +78,4 @@ def argsort(input, dim=-1, descending=False, *, out=None): # 启动 Kernel kernel(input_logic, values_logic, indices_logic, dim_size, descending) - return indices_logic \ No newline at end of file + return indices_logic diff --git a/src/ntops/torch/atan.py b/src/ntops/torch/atan.py index 4005fdf..e32ef93 100644 --- a/src/ntops/torch/atan.py +++ b/src/ntops/torch/atan.py @@ -7,14 +7,14 @@ def atan(input, *, out=None): """ 计算反正切函数 atan(x) - + 参数: input: 输入张量 out: 可选的输出张量 - + 返回: 计算 atan 后的张量 - + 数值特性: - 定义域: (-inf, +inf) - 值域: (-pi/2, +pi/2) @@ -22,35 +22,39 @@ def atan(input, *, out=None): """ # 确定输出数据类型 tensor_dtype = out.dtype if out is not None else input.dtype - + # 创建临时输出张量 temp_out = torch.empty_like(input, dtype=tensor_dtype, device=input.device) - + # 设置块大小 block_size = 256 - + # 缓存并获取 atan 内核 kernel = _cached_make( - ntops.kernels.atan.premake, + ntops.kernels.atan.premake, input.ndim, # ndim - 0, # dummy dim + 0, # dummy dim dtype=input.dtype, - block_size=block_size + block_size=block_size, ) - + # 执行内核计算 kernel(input, temp_out) - + # 处理 out 参数 if out is not None: if out.shape != temp_out.shape: - raise RuntimeError(f"Expected out tensor to have shape {temp_out.shape}, but got {out.shape}") + raise RuntimeError( + f"Expected out tensor to have shape {temp_out.shape}, but got {out.shape}" + ) if out.dtype != temp_out.dtype: - raise RuntimeError(f"Expected out tensor to have dtype {temp_out.dtype}, but got {out.dtype}") - + raise RuntimeError( + f"Expected out tensor to have dtype {temp_out.dtype}, but got {out.dtype}" + ) + out.copy_(temp_out) return out - + return temp_out @@ -58,6 +62,7 @@ def atan(input, *, out=None): def _atan_tensor_method(self, *, out=None): return atan(self, out=out) + # 注册到 PyTorch Tensor torch.Tensor.atan = _atan_tensor_method @@ -65,4 +70,4 @@ def _atan_tensor_method(self, *, out=None): # 别名支持 (兼容 NumPy 命名习惯) def arctan(input, *, out=None): """atan 的别名""" - return atan(input, out=out) \ No newline at end of file + return atan(input, out=out) diff --git a/src/ntops/torch/batch_norm.py b/src/ntops/torch/batch_norm.py index 3fc1d7b..4f8e2a2 100644 --- a/src/ntops/torch/batch_norm.py +++ b/src/ntops/torch/batch_norm.py @@ -1,8 +1,9 @@ -import math import torch + import ntops from ntops.torch.utils import _cached_make + def batch_norm(input, weight=None, bias=None, eps=1e-5, training=True): ndim = input.ndim if ndim < 2: @@ -11,7 +12,7 @@ def batch_norm(input, weight=None, bias=None, eps=1e-5, training=True): # 假设 dim=1 是 Channel，其余为 Batch 和 Spatial channel_dim = 1 reduction_dims = tuple(d for d in range(ndim) if d != channel_dim) - + num_elements = 1 for d in reduction_dims: num_elements *= input.shape[d] @@ -33,12 +34,9 @@ def batch_norm(input, weight=None, bias=None, eps=1e-5, training=True): output = torch.empty_like(input) kernel = _cached_make( - ntops.kernels.batch_norm.premake, - ndim, - reduction_dims, - num_elements + ntops.kernels.batch_norm.premake, ndim, reduction_dims, num_elements ) kernel(input, weight_expanded, bias_expanded, eps, output, num_elements) - return output \ No newline at end of file + return output diff --git a/src/ntops/torch/bincount.py b/src/ntops/torch/bincount.py index 439dc90..b4d1f34 100644 --- a/src/ntops/torch/bincount.py +++ b/src/ntops/torch/bincount.py @@ -1,7 +1,9 @@ import torch + import ntops from ntops.torch.utils import _cached_make + def bincount(input, weights=None, minlength=0): if input.ndim != 1: raise ValueError("input must be 1-dimensional") @@ -11,19 +13,19 @@ def bincount(input, weights=None, minlength=0): raise ValueError("weights should have the same shape as input") T = input.shape[0] - + # 计算输出大小 S if T > 0: max_val = input.max().item() S = max(int(max_val) + 1, minlength) else: S = minlength - + # 计算 T 的下一个 2 的幂，用于 arange 和 tiling T_pow2 = 1 << (T - 1).bit_length() if T > 0 else 1 - + # S 不需要严格的 2 的幂用于 Kernel 逻辑，但作为参数传递保持一致性 - S_pow2 = 1 + S_pow2 = 1 # 处理 Weights 和 Output 类型 # Torch 语义: weights 为 None 时返回 Long (int64), 否则返回 weights 的类型 @@ -35,19 +37,21 @@ def bincount(input, weights=None, minlength=0): # 准备 Output Tensor output = torch.zeros(S, dtype=out_dtype, device=input.device) - + # 准备 Bin IDs (辅助 Tensor，用于告知每个 Block 它负责哪些 Bin) bin_ids = torch.arange(S, dtype=input.dtype, device=input.device) # Block Size: 决定每个 Grid 处理多少个 Output Bin - block_size = 128 - - kernel = _cached_make(ntops.kernels.bincount.premake, - T_pow2=T_pow2, - S_pow2=S_pow2, - dtype=out_dtype, - block_size=block_size) - + block_size = 128 + + kernel = _cached_make( + ntops.kernels.bincount.premake, + T_pow2=T_pow2, + S_pow2=S_pow2, + dtype=out_dtype, + block_size=block_size, + ) + kernel(input, weights, output, bin_ids, T, S) - return output \ No newline at end of file + return output diff --git a/src/ntops/torch/fmax.py b/src/ntops/torch/fmax.py index b0945f5..4381860 100644 --- a/src/ntops/torch/fmax.py +++ b/src/ntops/torch/fmax.py @@ -1,36 +1,34 @@ import torch + import ntops -import ninetoothed from ntops.torch.utils import _cached_make + def fmax(input, other, out=None): input_b, other_b = torch.broadcast_tensors(input, other) - + # 必须 contiguous，因为 kernel 使用 flatten() 视作 1D 数组处理 input_b = input_b.contiguous() other_b = other_b.contiguous() - + output_shape = input_b.shape - + if out is None: out = torch.empty(output_shape, dtype=input.dtype, device=input.device) else: assert out.shape == output_shape if not out.is_contiguous(): - # 如果用户传入的 out 不连续，会导致写入错误，这里最好报错或处理 - # 但为了性能，通常假设用户知道自己在做什么，或者在这里强制 check - raise RuntimeError("Output tensor must be contiguous for fmax kernel") + # 如果用户传入的 out 不连续，会导致写入错误，这里最好报错或处理 + # 但为了性能，通常假设用户知道自己在做什么，或者在这里强制 check + raise RuntimeError("Output tensor must be contiguous for fmax kernel") block_size = 1024 - + # 注意：这里传递 input_b.ndim 给 premake # 虽然我们在 kernel 里 flatten 了，但 Tensor(ndim) 的定义有助于 ninetoothed 校验参数 kernel = _cached_make( - ntops.kernels.fmax.premake, - input_b.ndim, - input_b.dtype, - block_size + ntops.kernels.fmax.premake, input_b.ndim, input_b.dtype, block_size ) - + kernel(input_b, other_b, out) - return out \ No newline at end of file + return out diff --git a/src/ntops/torch/lp_pool1d.py b/src/ntops/torch/lp_pool1d.py index 5635013..257b9d5 100644 --- a/src/ntops/torch/lp_pool1d.py +++ b/src/ntops/torch/lp_pool1d.py @@ -1,4 +1,5 @@ import math + import torch import ntops @@ -49,7 +50,7 @@ def lp_pool1d(input, norm_type, kernel_size, stride=None, ceil_mode=False): stride, ceil_mode=ceil_mode, dtype=input.dtype, - block_size=block_size + block_size=block_size, ) kernel(input, output, norm_type, kernel_size) diff --git a/src/ntops/torch/lp_pool2d.py b/src/ntops/torch/lp_pool2d.py index 2bd8c5d..d5038f9 100644 --- a/src/ntops/torch/lp_pool2d.py +++ b/src/ntops/torch/lp_pool2d.py @@ -1,12 +1,21 @@ import math + import torch import ntops from ntops.torch.utils import _cached_make -def lp_pool2d(input, norm_type, kernel_size: int | tuple[int, int], stride: None | int | tuple[int, int] = None, ceil_mode=False): - assert input.ndim == 4 or input.ndim == 3, "Input tensor must be 4-dimensional (N, C, H_in, W_in) or 3-dimensional (C, H_in, W_in)" +def lp_pool2d( + input, + norm_type, + kernel_size: int | tuple[int, int], + stride: None | int | tuple[int, int] = None, + ceil_mode=False, +): + assert input.ndim == 4 or input.ndim == 3, ( + "Input tensor must be 4-dimensional (N, C, H_in, W_in) or 3-dimensional (C, H_in, W_in)" + ) if input.ndim == 3: input = input.unsqueeze(0) # 添加 batch 维度 @@ -28,7 +37,7 @@ def lp_pool2d(input, norm_type, kernel_size: int | tuple[int, int], stride: None else: H_out = math.floor((H_in - kernel_size[0] + stride[0]) / stride[0]) W_out = math.floor((W_in - kernel_size[1] + stride[1]) / stride[1]) - + output_shape = (input.shape[0], input.shape[1], H_out, W_out) output = torch.empty(output_shape, dtype=input.dtype, device=input.device) @@ -36,30 +45,29 @@ def lp_pool2d(input, norm_type, kernel_size: int | tuple[int, int], stride: None block_size = 1024 if ceil_mode: kernel = _cached_make( - ntops.kernels.lp_pool2d.premake_ceil_mode, - input.ndim, - kernel_size[0], + ntops.kernels.lp_pool2d.premake_ceil_mode, + input.ndim, + kernel_size[0], kernel_size[1], stride[0], stride[1], block_size=block_size, ceil_mode=ceil_mode, - dtype=input.dtype + dtype=input.dtype, ) kernel(input, output, norm_type, kernel_size[0] * kernel_size[1]) else: kernel = _cached_make( - ntops.kernels.lp_pool2d.premake, - input.ndim, - kernel_size[0], + ntops.kernels.lp_pool2d.premake, + input.ndim, + kernel_size[0], kernel_size[1], stride[0], stride[1], block_size=block_size, ceil_mode=ceil_mode, - dtype=input.dtype + dtype=input.dtype, ) kernel(input, output, norm_type) - - + return output diff --git a/src/ntops/torch/lp_pool3d.py b/src/ntops/torch/lp_pool3d.py index 3438e24..67504b7 100644 --- a/src/ntops/torch/lp_pool3d.py +++ b/src/ntops/torch/lp_pool3d.py @@ -1,12 +1,21 @@ import math + import torch import ntops from ntops.torch.utils import _cached_make -def lp_pool3d(input, norm_type, kernel_size: int | tuple[int, int, int], stride: None | int | tuple[int, int, int] = None, ceil_mode=False): - assert input.ndim == 5 or input.ndim == 4, "Input tensor must be 4-dimensional (N, C, D_in, H_in, W_in) or 3-dimensional (C, D_in, H_in, W_in)" +def lp_pool3d( + input, + norm_type, + kernel_size: int | tuple[int, int, int], + stride: None | int | tuple[int, int, int] = None, + ceil_mode=False, +): + assert input.ndim == 5 or input.ndim == 4, ( + "Input tensor must be 4-dimensional (N, C, D_in, H_in, W_in) or 3-dimensional (C, D_in, H_in, W_in)" + ) if input.ndim == 4: input = input.unsqueeze(0) # 添加 batch 维度 @@ -30,7 +39,7 @@ def lp_pool3d(input, norm_type, kernel_size: int | tuple[int, int, int], stride: D_out = math.floor((D_in - kernel_size[0] + stride[0]) / stride[0]) H_out = math.floor((H_in - kernel_size[1] + stride[1]) / stride[1]) W_out = math.floor((W_in - kernel_size[2] + stride[2]) / stride[2]) - + output_shape = (N, C, D_out, H_out, W_out) output = torch.empty(output_shape, dtype=input.dtype, device=input.device) @@ -38,8 +47,8 @@ def lp_pool3d(input, norm_type, kernel_size: int | tuple[int, int, int], stride: block_size = 256 if ceil_mode: kernel = _cached_make( - ntops.kernels.lp_pool3d.premake_ceil_mode, - input.ndim, + ntops.kernels.lp_pool3d.premake_ceil_mode, + input.ndim, kernel_size[0], kernel_size[1], kernel_size[2], @@ -50,12 +59,14 @@ def lp_pool3d(input, norm_type, kernel_size: int | tuple[int, int, int], stride: ceil_mode=ceil_mode, dtype=input.dtype, ) - kernel(input, output, norm_type, kernel_size[0] * kernel_size[1] * kernel_size[2]) + kernel( + input, output, norm_type, kernel_size[0] * kernel_size[1] * kernel_size[2] + ) else: kernel = _cached_make( - ntops.kernels.lp_pool3d.premake, - input.ndim, - kernel_size[0], + ntops.kernels.lp_pool3d.premake, + input.ndim, + kernel_size[0], kernel_size[1], kernel_size[2], stride[0], @@ -63,9 +74,8 @@ def lp_pool3d(input, norm_type, kernel_size: int | tuple[int, int, int], stride: stride[2], block_size=block_size, ceil_mode=ceil_mode, - dtype=input.dtype + dtype=input.dtype, ) kernel(input, output, norm_type) - - + return output diff --git a/src/ntops/torch/max.py b/src/ntops/torch/max.py index 822d867..2c35066 100644 --- a/src/ntops/torch/max.py +++ b/src/ntops/torch/max.py @@ -1,10 +1,10 @@ +import math + import torch import ntops -import ninetoothed from ntops.torch.utils import _cached_make -import builtins -import math + def max(input, dim: int | None = None, keepdim=False, *, out=None): if dim is None: @@ -14,17 +14,24 @@ def max(input, dim: int | None = None, keepdim=False, *, out=None): block_size = 1024 while current.numel() > 1: output_shape = (math.ceil(current.numel() / block_size),) - output = torch.empty(output_shape, dtype=current.dtype, device=current.device) - kernel = _cached_make(ntops.kernels.max.premake_all_elements, current.ndim, current.dtype, block_size) + output = torch.empty( + output_shape, dtype=current.dtype, device=current.device + ) + kernel = _cached_make( + ntops.kernels.max.premake_all_elements, + current.ndim, + current.dtype, + block_size, + ) kernel(current, output) current = output - + result = current.view(()) - + if out is not None: out.copy_(result) return out - + return result else: output_shape = list(input.shape) diff --git a/src/ntops/torch/max_pool1d.py b/src/ntops/torch/max_pool1d.py index 9a4d438..5af6454 100644 --- a/src/ntops/torch/max_pool1d.py +++ b/src/ntops/torch/max_pool1d.py @@ -1,28 +1,39 @@ import math + import torch import torch.nn.functional as F + import ntops from ntops.torch.utils import _cached_make + def max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False): # 维度检查 is_3d = input.ndim == 3 if input.ndim == 2: - input = input.unsqueeze(0) # (C, L) -> (1, C, L) - - assert input.ndim == 3, "Input tensor must be 3-dimensional (N, C, L) or 2-dimensional (C, L)" + input = input.unsqueeze(0) # (C, L) -> (1, C, L) + + assert input.ndim == 3, ( + "Input tensor must be 3-dimensional (N, C, L) or 2-dimensional (C, L)" + ) # 参数标准化 if stride is None: stride = kernel_size - + # 处理 Tuple 参数 (虽然 1d 通常是 int，但 torch 允许 (int,)) - if isinstance(kernel_size, tuple): kernel_size = kernel_size[0] - if isinstance(stride, tuple): stride = stride[0] - if isinstance(padding, tuple): padding = padding[0] - if isinstance(dilation, tuple): dilation = dilation[0] + if isinstance(kernel_size, tuple): + kernel_size = kernel_size[0] + if isinstance(stride, tuple): + stride = stride[0] + if isinstance(padding, tuple): + padding = padding[0] + if isinstance(dilation, tuple): + dilation = dilation[0] - assert dilation == 1, "Currently only dilation=1 is supported in this DSL implementation" + assert dilation == 1, ( + "Currently only dilation=1 is supported in this DSL implementation" + ) # 处理 Explicit Padding # 如果有 padding，先用 -inf 填充 input @@ -42,7 +53,7 @@ def max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode output = torch.empty(output_shape, dtype=input.dtype, device=input.device) block_size = 1024 - + kernel = _cached_make( ntops.kernels.max_pool1d.premake, input.ndim, @@ -50,7 +61,7 @@ def max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode stride, block_size=block_size, ceil_mode=ceil_mode, - dtype=input.dtype + dtype=input.dtype, ) kernel(input, output) @@ -58,4 +69,4 @@ def max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode if not is_3d: output = output.squeeze(0) - return output \ No newline at end of file + return output diff --git a/src/ntops/torch/max_pool3d.py b/src/ntops/torch/max_pool3d.py index 1bf4931..965bd2d 100644 --- a/src/ntops/torch/max_pool3d.py +++ b/src/ntops/torch/max_pool3d.py @@ -1,26 +1,33 @@ import math + import torch import torch.nn.functional as F + import ntops from ntops.torch.utils import _cached_make + def max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False): assert input.ndim == 5, "Input tensor must be 5-dimensional (N, C, D, H, W)" - assert dilation == 1 or dilation == (1, 1, 1), "Currently only dilation=1 is supported" + assert dilation == 1 or dilation == (1, 1, 1), ( + "Currently only dilation=1 is supported" + ) # --- 参数归一化处理 --- def _triple(x): - if isinstance(x, int): return (x, x, x) - if len(x) == 1: return (x[0], x[0], x[0]) + if isinstance(x, int): + return (x, x, x) + if len(x) == 1: + return (x[0], x[0], x[0]) return x k_d, k_h, k_w = _triple(kernel_size) - + if stride is None: s_d, s_h, s_w = k_d, k_h, k_w else: s_d, s_h, s_w = _triple(stride) - + pad_d, pad_h, pad_w = _triple(padding) # --- 精确计算输出尺寸和所需的右侧 Padding --- @@ -30,22 +37,22 @@ def _calc_dim_and_pad(in_dim, k, s, p, ceil_mode): out_dim = math.ceil((in_dim + 2 * p - k) / s) + 1 else: out_dim = math.floor((in_dim + 2 * p - k) / s) + 1 - + # 2. PyTorch 特殊边界检查 (Ceil Mode 下排除纯 Padding 的窗口) # 逻辑: 如果最后一个窗口的起始位置 >= 原始输入长度 + 左padding，则该窗口无效 if ceil_mode: if (out_dim - 1) * s >= in_dim + p: out_dim -= 1 - + # 3. 反推需要的总长度，以满足 DSL 的 floor tiling # 我们需要: (out_dim - 1) * s + k needed_len = (out_dim - 1) * s + k - + # 4. 计算右侧需要补多少 # 当前已有: in_dim + p (左侧 padding) current_len = in_dim + p pad_right = max(0, needed_len - current_len) - + return out_dim, pad_right D_in, H_in, W_in = input.shape[-3], input.shape[-2], input.shape[-1] @@ -58,28 +65,30 @@ def _calc_dim_and_pad(in_dim, k, s, p, ceil_mode): # F.pad 顺序: (W_left, W_right, H_top, H_bot, D_front, D_back) if any(p > 0 for p in [pad_w, pad_r_w, pad_h, pad_r_h, pad_d, pad_r_d]): input = F.pad( - input, - (pad_w, pad_r_w, pad_h, pad_r_h, pad_d, pad_r_d), - value=float("-inf") + input, (pad_w, pad_r_w, pad_h, pad_r_h, pad_d, pad_r_d), value=float("-inf") ) output_shape = (input.shape[0], input.shape[1], D_out, H_out, W_out) output = torch.empty(output_shape, dtype=input.dtype, device=input.device) block_size = 1024 - + # 注意: 这里 ceil_mode 永远传 False # 因为我们已经通过 Explicit Padding 确保了 floor 切分能得到正确的 output size kernel = _cached_make( ntops.kernels.max_pool3d.premake, input.ndim, - k_d, k_h, k_w, - s_d, s_h, s_w, + k_d, + k_h, + k_w, + s_d, + s_h, + s_w, block_size=block_size, - ceil_mode=False, - dtype=input.dtype + ceil_mode=False, + dtype=input.dtype, ) kernel(input, output) - return output \ No newline at end of file + return output diff --git a/src/ntops/torch/maximum.py b/src/ntops/torch/maximum.py index 510efc7..fc4a2a9 100644 --- a/src/ntops/torch/maximum.py +++ b/src/ntops/torch/maximum.py @@ -1,44 +1,45 @@ import torch + import ntops -import ninetoothed -from ntops.torch.utils import _cached_make + # 假设上面的 kernel 代码保存在 ntops.kernels.maximum 中 -import ntops.kernels.maximum +import ntops.kernels.maximum +from ntops.torch.utils import _cached_make + def maximum(input, other, out=None): # 1. 处理广播机制 (Broadcasting) # 使 input 和 other 具有相同的形状 input_b, other_b = torch.broadcast_tensors(input, other) - + # 2. 确保内存连续 (Contiguous) # Triton/DSL kernel 通常假设数据在内存中是紧凑排列的 input_b = input_b.contiguous() other_b = other_b.contiguous() - + output_shape = input_b.shape - + # 3. 准备输出 Tensor if out is None: out = torch.empty(output_shape, dtype=input.dtype, device=input.device) else: # 简单的形状检查 - assert out.shape == output_shape, f"Output shape mismatch: expected {output_shape}, got {out.shape}" + assert out.shape == output_shape, ( + f"Output shape mismatch: expected {output_shape}, got {out.shape}" + ) if not out.is_contiguous(): raise RuntimeError("Output tensor must be contiguous for maximum kernel") # 4. 设置 Block Size # 这里的 1024 是经验值，通常可以根据硬件或 heuristic 动态调整 block_size = 1024 - + # 5. 获取并编译 Kernel kernel = _cached_make( - ntops.kernels.maximum.premake, - input_b.ndim, - input_b.dtype, - block_size + ntops.kernels.maximum.premake, input_b.ndim, input_b.dtype, block_size ) - + # 6. 执行 Kernel kernel(input_b, other_b, out) - - return out \ No newline at end of file + + return out diff --git a/src/ntops/torch/mean.py b/src/ntops/torch/mean.py index a666636..3dce8b9 100644 --- a/src/ntops/torch/mean.py +++ b/src/ntops/torch/mean.py @@ -1,15 +1,20 @@ import math + import torch + import ntops -from ntops.torch.utils import _cached_make + # 引入上面定义的 kernel -import ntops.kernels.mean +import ntops.kernels.mean +from ntops.torch.utils import _cached_make + def next_power_of_2(n): if n == 0: return 1 return 1 << (n - 1).bit_length() + def get_optimal_block_size(dim_size): target_size = next_power_of_2(dim_size) if target_size > 1024: @@ -18,6 +23,7 @@ def get_optimal_block_size(dim_size): target_size = 32 return target_size + def mean( input, dim: int | tuple[int] | list[int] | None = None, @@ -44,34 +50,36 @@ def mean( dims = (dim,) else: dims = tuple(dim) - + num_elements = 1 for d in dims: num_elements *= input.shape[d] # 3. Kernel 计算 (本质是 Sum，但是使用 computation_dtype) - + # --- Case A: Global Mean (所有元素) --- if dim is None: current = input block_size = get_optimal_block_size(current.numel()) - + # 递归规约 (Global Reduction) while current.numel() > 1: output_len = math.ceil(current.numel() / block_size) - output = torch.empty((output_len,), dtype=computation_dtype, device=current.device) - + output = torch.empty( + (output_len,), dtype=computation_dtype, device=current.device + ) + kernel = _cached_make( ntops.kernels.mean.premake_all_elements, current.ndim, - computation_dtype, # 确保 kernel 使用浮点 + computation_dtype, # 确保 kernel 使用浮点 block_size, ) kernel(current, output) current = output - + result_sum = current.view(()) - + # 4. 执行除法 result = result_sum.div(num_elements) @@ -87,16 +95,18 @@ def mean( if d < 0: d += input.ndim output_shape[d] = 1 - - temp_out = torch.empty(output_shape, dtype=computation_dtype, device=input.device) + + temp_out = torch.empty( + output_shape, dtype=computation_dtype, device=input.device + ) block_size = get_optimal_block_size(output_shape[dims[0]]) kernel = _cached_make( - ntops.kernels.mean.premake, - input.ndim, - dims, - computation_dtype, # 确保 kernel 使用浮点 - block_size + ntops.kernels.mean.premake, + input.ndim, + dims, + computation_dtype, # 确保 kernel 使用浮点 + block_size, ) kernel(input, temp_out) @@ -110,7 +120,7 @@ def mean( final_shape = list(output_shape) for d in dims_to_remove: del final_shape[d] - + if not final_shape: temp_out = temp_out.view(()) else: @@ -120,4 +130,4 @@ def mean( out.copy_(temp_out) return out - return temp_out \ No newline at end of file + return temp_out diff --git a/src/ntops/torch/median.py b/src/ntops/torch/median.py index 2f16eb6..801939e 100644 --- a/src/ntops/torch/median.py +++ b/src/ntops/torch/median.py @@ -1,12 +1,15 @@ import torch + import ntops from ntops.torch.utils import _cached_make + def next_power_of_2(n): if n == 0: return 1 return 1 << (n - 1).bit_length() + def get_optimal_block_size(dim_size): target_size = next_power_of_2(dim_size) if target_size > 1024: @@ -15,6 +18,7 @@ def get_optimal_block_size(dim_size): target_size = 32 return target_size + def median(input, dim=-1, keepdim=False, *, out=None): """ Args: @@ -31,13 +35,13 @@ def median(input, dim=-1, keepdim=False, *, out=None): # 处理 dim if dim is None: raise NotImplementedError("median currently requires a specific dim") - + if dim < 0: dim += input.ndim - + target_dim = dim dim_size = input.shape[target_dim] - + # 计算中位数在排序后的索引位置 (PyTorch 默认行为：(N-1) // 2) # 比如 N=3, idx=1; N=4, idx=1 median_rank = (dim_size - 1) // 2 @@ -50,23 +54,25 @@ def median(input, dim=-1, keepdim=False, *, out=None): # 但 Kernel 这里的 arrangement 是一对一的 (Input Block -> Output Block) # 通常 reduction kernel 输出的 shape 在 target_dim 上是 1 或被 squeeze # 这里我们按照 keepdim=True 的形状申请，最后根据参数 squeeze - + output_shape = list(input.shape) - output_shape[target_dim] = 1 # 结果只有 1 个元素 - + output_shape[target_dim] = 1 # 结果只有 1 个元素 + if out is not None: values_logic, indices_logic = out else: values_logic = torch.empty(output_shape, dtype=dtype, device=input.device) - indices_logic = torch.empty(output_shape, dtype=indices_dtype, device=input.device) + indices_logic = torch.empty( + output_shape, dtype=indices_dtype, device=input.device + ) # 构建 Kernel # 注意：我们将 median_rank 作为 loop_k 传入 premake kernel = _cached_make( - ntops.kernels.median.premake, # 假设上面的 kernel 代码在这个路径 + ntops.kernels.median.premake, # 假设上面的 kernel 代码在这个路径 input_logic.ndim, target_dim, - median_rank, # loop_k + median_rank, # loop_k dtype, indices_dtype, block_size, @@ -80,4 +86,4 @@ def median(input, dim=-1, keepdim=False, *, out=None): values_logic = values_logic.squeeze(dim) indices_logic = indices_logic.squeeze(dim) - return torch.return_types.median((values_logic, indices_logic)) \ No newline at end of file + return torch.return_types.median((values_logic, indices_logic)) diff --git a/src/ntops/torch/stack.py b/src/ntops/torch/stack.py index 76ae6a7..b73d7ca 100644 --- a/src/ntops/torch/stack.py +++ b/src/ntops/torch/stack.py @@ -1,20 +1,24 @@ import torch + import ntops from ntops.torch.utils import _cached_make + def stack(tensors, dim=0): if not tensors: raise ValueError("stack expects a non-empty list of tensors") - + first_tensor = tensors[0] target_dtype = first_tensor.dtype target_device = first_tensor.device input_shape = first_tensor.shape ndim = first_tensor.ndim - + # 1. 校验所有输入张量的形状和类型必须一致 for t in tensors: - assert t.shape == input_shape, f"Shape mismatch: expected {input_shape}, got {t.shape}" + assert t.shape == input_shape, ( + f"Shape mismatch: expected {input_shape}, got {t.shape}" + ) assert t.dtype == target_dtype, "Dtype mismatch in input tensors" assert t.device == target_device, "Device mismatch in input tensors" @@ -22,20 +26,20 @@ def stack(tensors, dim=0): # stack 会在 dim 维度插入大小为 len(tensors) 的新维度 output_shape = list(input_shape) output_shape.insert(dim, len(tensors)) - + # 3. 分配输出内存 output = torch.empty(output_shape, dtype=target_dtype, device=target_device) - + # 4. 准备 Kernel # 注意：这里我们使用 input 的 ndim，因为 Kernel 处理的是切片 block_size = 1024 kernel = _cached_make( - ntops.kernels.stack.premake, # 假设上面的代码保存在这里 + ntops.kernels.stack.premake, # 假设上面的代码保存在这里 ndim, block_size=block_size, - dtype=target_dtype + dtype=target_dtype, ) - + # 5. 循环执行 (Iterative Execution) # 对于每个输入张量，将其拷贝到 output 对应的切片中 for i, t in enumerate(tensors): @@ -44,5 +48,5 @@ def stack(tensors, dim=0): # Kernel 就像处理两个普通张量一样处理它们。 output_slice = output.select(dim, i) kernel(t, output_slice) - - return output \ No newline at end of file + + return output diff --git a/tests/test_acosh.py b/tests/test_acosh.py index 3dbd32e..f3e7b9e 100644 --- a/tests/test_acosh.py +++ b/tests/test_acosh.py @@ -1,5 +1,3 @@ -import random - import pytest import torch @@ -18,19 +16,19 @@ def test_acosh_basic(shape, dtype, device, rtol, atol): # 创建输入张量，确保值在定义域内（x >= 1） # 生成[1, 10]范围内的随机数，避免数值不稳定区域 input_tensor = 1.0 + 9.0 * torch.rand(shape, dtype=dtype, device=device) - + # 计算参考输出（PyTorch内置实现） reference_output = torch.acosh(input_tensor) - + # 计算ntops实现 ntops_output = ntops.torch.acosh(input_tensor) - + # 验证数值精度 assert torch.allclose(ntops_output, reference_output, rtol=rtol, atol=atol) - + # 验证输出形状与输入形状一致 assert ntops_output.shape == input_tensor.shape - + # 验证数据类型一致性 assert ntops_output.dtype == input_tensor.dtype @@ -45,26 +43,30 @@ def test_acosh_boundary_values(dtype, device): """ # 测试点：1.0（精确边界）, 1.000001（接近1）, 100.0（大值）, 0.5（小于1，应返回NaN） test_values = torch.tensor([1.0, 1.000001, 100.0, 0.5], dtype=dtype, device=device) - + # 计算参考输出 reference_output = torch.acosh(test_values) - + # 计算ntops实现 ntops_output = ntops.torch.acosh(test_values) - + # 验证数值精度（跳过NaN比较） valid_mask = ~torch.isnan(reference_output) - assert torch.allclose(ntops_output[valid_mask], reference_output[valid_mask], - rtol=1e-3 if dtype == torch.float16 else 1e-5, - atol=1e-3 if dtype == torch.float16 else 1e-6) - + assert torch.allclose( + ntops_output[valid_mask], + reference_output[valid_mask], + rtol=1e-3 if dtype == torch.float16 else 1e-5, + atol=1e-3 if dtype == torch.float16 else 1e-6, + ) + # 验证NaN处理 nan_mask = torch.isnan(reference_output) assert torch.all(torch.isnan(ntops_output[nan_mask])) - + # 验证acosh(1) = 0 - assert torch.allclose(ntops_output[0], torch.tensor(0.0, dtype=dtype, device=device), - atol=1e-6) + assert torch.allclose( + ntops_output[0], torch.tensor(0.0, dtype=dtype, device=device), atol=1e-6 + ) @skip_if_cuda_not_available @@ -76,19 +78,19 @@ def test_acosh_with_contiguous_out(dtype, dim): 验证out参数的正确性和数据一致性 """ device = "cuda" - + # 创建输入张量（3D） input_tensor = 1.0 + 9.0 * torch.rand(4, 5, 6, dtype=dtype, device=device) - + # 计算参考输出 reference_output = torch.acosh(input_tensor) - + # 创建连续的输出张量 out = torch.empty_like(reference_output) - + # 使用ntops的acosh，传入out参数 result = ntops.torch.acosh(input_tensor, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -105,20 +107,22 @@ def test_acosh_with_strided_output(dtype): 验证内核对非标准内存布局的处理能力 """ device = "cuda" - + # 创建输入张量 input_tensor = 1.0 + 9.0 * torch.rand(4, 5, 6, dtype=dtype, device=device) - + # 计算参考输出 reference_output = torch.acosh(input_tensor) - + # 创建一个具有不同strides的输出张量 large_tensor = torch.empty(4, 5, 8, dtype=dtype, device=device) - out = large_tensor[:, :, :6] # shape (4, 5, 6) 但strides为(40, 8, 1)而不是标准的(30, 6, 1) - + out = large_tensor[ + :, :, :6 + ] # shape (4, 5, 6) 但strides为(40, 8, 1)而不是标准的(30, 6, 1) + # 使用ntops的acosh，传入strided输出 result = ntops.torch.acosh(input_tensor, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -135,22 +139,22 @@ def test_acosh_with_transposed_output(dtype): 验证内核对转置张量内存布局的处理能力 """ device = "cuda" - + # 创建输入张量 input_tensor = 1.0 + 9.0 * torch.rand(3, 4, 5, dtype=dtype, device=device) - + # 计算参考输出 reference_output = torch.acosh(input_tensor) - + # 创建一个转置的输出张量（非连续） out_base = torch.empty(5, 4, 3, dtype=dtype, device=device) out = out_base.permute(2, 1, 0) # shape (3, 4, 5)，但内存布局非连续 - + assert not out.is_contiguous(), "out should be non-contiguous before calling acosh" - + # 使用ntops的acosh result = ntops.torch.acosh(input_tensor, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -166,24 +170,21 @@ def test_acosh_with_strided_slice_output(dtype): 验证内核对复杂内存布局的处理能力 """ device = "cuda" - + # 创建输入张量 input_tensor = 1.0 + 9.0 * torch.rand(2, 8, 6, dtype=dtype, device=device) - - # 计算参考输出 - reference_output = torch.acosh(input_tensor) - + # 创建一个更大的张量，使用步进切片得到非连续的子张量 large_tensor = torch.empty(2, 8, 12, dtype=dtype, device=device) out = large_tensor[:, ::2, ::2] # shape (2, 4, 6)，strides非标准 - + # 调整输入以匹配输出的第二维（仅用于测试，实际使用不需要） input_tensor_adjusted = input_tensor[:, ::2, :] # shape (2, 4, 6) reference_output_adjusted = torch.acosh(input_tensor_adjusted) - + # 使用ntops的acosh result = ntops.torch.acosh(input_tensor_adjusted, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -201,21 +202,23 @@ def test_acosh_out_vs_no_out(dtype, keep_out): 验证两种调用方式的数值一致性 """ device = "cuda" - + # 创建输入张量 input_tensor = 1.0 + 9.0 * torch.rand(4, 5, 6, dtype=dtype, device=device) - + # 不使用out参数 result_no_out = ntops.torch.acosh(input_tensor) - + # 使用out参数 if keep_out: out = torch.empty_like(result_no_out) result_with_out = ntops.torch.acosh(input_tensor, out=out) assert result_with_out is out else: - result_with_out = ntops.torch.acosh(input_tensor, out=torch.empty_like(result_no_out)) - + result_with_out = ntops.torch.acosh( + input_tensor, out=torch.empty_like(result_no_out) + ) + # 验证结果一致 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -230,31 +233,37 @@ def test_acosh_special_values(dtype, device="cuda"): 包括：NaN, inf, -inf等 """ # 创建包含特殊值的输入张量 - special_values = torch.tensor([ - float('nan'), # NaN - float('inf'), # inf - -float('inf'), # -inf - 1.0, # 边界值 - 2.0, # 正常值 - 0.5 # 小于1（应返回NaN） - ], dtype=dtype, device=device) - + special_values = torch.tensor( + [ + float("nan"), # NaN + float("inf"), # inf + -float("inf"), # -inf + 1.0, # 边界值 + 2.0, # 正常值 + 0.5, # 小于1（应返回NaN） + ], + dtype=dtype, + device=device, + ) + # 计算ntops实现 ntops_output = ntops.torch.acosh(special_values) - + # 验证特殊值处理 # NaN输入应该产生NaN输出 assert torch.isnan(ntops_output[0]) - + # inf输入应该产生inf输出 assert torch.isinf(ntops_output[1]) and ntops_output[1] > 0 - + # -inf输入应该产生NaN（因为acosh定义域是x >= 1） assert torch.isnan(ntops_output[2]) - + # 1.0应该产生0.0 - assert torch.allclose(ntops_output[3], torch.tensor(0.0, dtype=dtype, device=device), atol=1e-6) - + assert torch.allclose( + ntops_output[3], torch.tensor(0.0, dtype=dtype, device=device), atol=1e-6 + ) + # 0.5应该产生NaN assert torch.isnan(ntops_output[5]) @@ -268,21 +277,21 @@ def test_acosh_large_tensors(shape, dtype): 验证内存使用和计算效率 """ device = "cuda" - + # 创建大输入张量 input_tensor = 1.0 + 9.0 * torch.rand(shape, dtype=dtype, device=device) - + # 计算参考输出 reference_output = torch.acosh(input_tensor) - + # 计算ntops实现 ntops_output = ntops.torch.acosh(input_tensor) - + # 验证数值精度 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 assert torch.allclose(ntops_output, reference_output, rtol=rtol, atol=atol) - + # 验证内存效率（简单检查，不实际测量内存） assert ntops_output.shape == input_tensor.shape @@ -295,27 +304,27 @@ def test_acosh_inplace_consistency(dtype): 虽然acosh本身不是原地操作，但需要确保与原地操作链的兼容性 """ device = "cuda" - + # 创建输入张量 input_tensor = 1.0 + 9.0 * torch.rand(4, 5, 6, dtype=dtype, device=device) - + # 复制输入用于原地操作测试 input_copy = input_tensor.clone() - + # 链式操作：先进行一些原地操作，然后计算acosh input_copy.add_(1.0) # 原地加法 input_copy.mul_(0.5) # 原地乘法 - + # 确保值仍在定义域内 input_copy = torch.clamp(input_copy, min=1.0) - + # 计算参考输出 reference_output = torch.acosh(input_copy) - + # 计算ntops实现 ntops_output = ntops.torch.acosh(input_copy) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 - assert torch.allclose(ntops_output, reference_output, rtol=rtol, atol=atol) \ No newline at end of file + assert torch.allclose(ntops_output, reference_output, rtol=rtol, atol=atol) diff --git a/tests/test_adaptive_avg_pool2d.py b/tests/test_adaptive_avg_pool2d.py index 446eb0c..844a983 100644 --- a/tests/test_adaptive_avg_pool2d.py +++ b/tests/test_adaptive_avg_pool2d.py @@ -1,9 +1,12 @@ import random + import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available + @skip_if_cuda_not_available @pytest.mark.parametrize("output_size", [(1, 1), (2, 2), (4, 4)]) @pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) @@ -12,29 +15,27 @@ def test_adaptive_avg_pool2d(output_size, dtype): batch = random.randint(1, 3) channels = random.randint(1, 4) - + # 构造能整除的 Shape 以保证与 PyTorch 行为一致 # 因为 DSL 使用固定 stride/kernel，而 PyTorch 在不能整除时使用变长 stride base_h, base_w = output_size height = base_h * random.randint(2, 6) width = base_w * random.randint(2, 6) - input_tensor = torch.randn((batch, channels, height, width), device=device, dtype=dtype) + input_tensor = torch.randn( + (batch, channels, height, width), device=device, dtype=dtype + ) # Ntops implementation - ntops_output = ntops.torch.adaptive_avg_pool2d( - input_tensor, - output_size - ) + ntops_output = ntops.torch.adaptive_avg_pool2d(input_tensor, output_size) # Reference implementation reference_output = torch.nn.functional.adaptive_avg_pool2d( - input_tensor, - output_size + input_tensor, output_size ) # float16 容差稍大 atol = 1e-3 if dtype == torch.float32 else 5e-3 rtol = 1e-3 if dtype == torch.float32 else 5e-3 - - assert torch.allclose(ntops_output, reference_output, atol=atol, rtol=rtol) \ No newline at end of file + + assert torch.allclose(ntops_output, reference_output, atol=atol, rtol=rtol) diff --git a/tests/test_adaptive_max_pool2d.py b/tests/test_adaptive_max_pool2d.py index 434db30..f2c91fc 100644 --- a/tests/test_adaptive_max_pool2d.py +++ b/tests/test_adaptive_max_pool2d.py @@ -1,9 +1,12 @@ import random + import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available + @skip_if_cuda_not_available @pytest.mark.parametrize("output_size", [(1, 1), (2, 2), (3, 3), (5, 7)]) def test_adaptive_max_pool2d(output_size): @@ -18,18 +21,16 @@ def test_adaptive_max_pool2d(output_size): height = base_h * random.randint(1, 5) width = base_w * random.randint(1, 5) - input_tensor = torch.randn((batch, channels, height, width), device=device, dtype=dtype) + input_tensor = torch.randn( + (batch, channels, height, width), device=device, dtype=dtype + ) # Ntops implementation - ntops_output = ntops.torch.adaptive_max_pool2d( - input_tensor, - output_size - ) + ntops_output = ntops.torch.adaptive_max_pool2d(input_tensor, output_size) # Reference implementation reference_output = torch.nn.functional.adaptive_max_pool2d( - input_tensor, - output_size + input_tensor, output_size ) - assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3) \ No newline at end of file + assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3) diff --git a/tests/test_addmv.py b/tests/test_addmv.py index 6fcf93d..feacc4e 100644 --- a/tests/test_addmv.py +++ b/tests/test_addmv.py @@ -3,14 +3,16 @@ import ntops from tests.skippers import skip_if_cuda_not_available + # 假设你有一个类似的参数生成器，或者我们需要稍微修改以适应 addmv (n=1) -from tests.test_mm import generate_arguments +from tests.test_mm import generate_arguments from tests.utils import gauss + @skip_if_cuda_not_available # 我们可以复用 generate_arguments，但在 addmv 中，n 维度通常不再是矩阵列数，而是 1 (或者不需要) # 这里假设 generate_arguments 返回 (m, n, k, ...)，我们在 test 中忽略 n 或者将其视为 batch -@pytest.mark.parametrize(*generate_arguments()) +@pytest.mark.parametrize(*generate_arguments()) def test_addmv(m, n, k, dtype, device, rtol, atol): # 构造数据 # input: (m,) - 偏置向量 @@ -19,15 +21,15 @@ def test_addmv(m, n, k, dtype, device, rtol, atol): mat = torch.randn((m, k), dtype=dtype, device=device) # vec: (k,) - 输入向量 vec = torch.randn((k,), dtype=dtype, device=device) - + beta = gauss() alpha = gauss() # 计算 ninetoothed 结果 ninetoothed_output = ntops.torch.addmv(input, mat, vec, beta=beta, alpha=alpha) - + # 计算 PyTorch 参考结果 reference_output = torch.addmv(input, mat, vec, beta=beta, alpha=alpha) # 验证结果 - assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol) \ No newline at end of file + assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol) diff --git a/tests/test_argsort.py b/tests/test_argsort.py index e6a01a2..742e6bd 100644 --- a/tests/test_argsort.py +++ b/tests/test_argsort.py @@ -1,10 +1,13 @@ import random + import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available from tests.utils import _random_shape + # 生成测试参数 def generate_argsort_args(): args = [] @@ -16,7 +19,7 @@ def generate_argsort_args(): for _ in range(5): shape = _random_shape(ndim) dim = random.randint(0, ndim - 1) - + # 限制 dim_size 以免测试跑太久 (因为选择排序是 O(N^2)) # 对于 DSL Demo 这是一个合理的限制 if shape[dim] > 128: @@ -27,6 +30,7 @@ def generate_argsort_args(): args.append((shape, dim, dtype, device, rtol, atol)) return "shape, dim, dtype, device, rtol, atol", args + @skip_if_cuda_not_available @pytest.mark.parametrize(*generate_argsort_args()) @pytest.mark.parametrize("descending", [True, False]) @@ -36,23 +40,23 @@ def test_argsort(shape, dim, dtype, device, rtol, atol, descending): # 调用 ntops 实现 # 注意：PyTorch 的 argsort 默认是 ascending (descending=False) ntops_indices = ntops.torch.argsort(input_tensor, dim=dim, descending=descending) - + # 调用 PyTorch 参考实现 ref_indices = torch.argsort(input_tensor, dim=dim, descending=descending) # 验证索引对应的值是否是有序的 # 直接比较 indices 可能会失败，因为对于相同的值，排序是不稳定的 (unstable) # 所以我们 gather 出值来比较值是否一致且有序 - + ntops_gathered = torch.gather(input_tensor, dim, ntops_indices) ref_gathered = torch.gather(input_tensor, dim, ref_indices) assert torch.allclose(ntops_gathered, ref_gathered, rtol=rtol, atol=atol) - + # 额外验证：确保输出确实是排序过的 if input_tensor.numel() > 0: diffs = ntops_gathered.diff(dim=dim) if descending: assert (diffs <= atol).all(), "Result is not sorted descending" else: - assert (diffs >= -atol).all(), "Result is not sorted ascending" \ No newline at end of file + assert (diffs >= -atol).all(), "Result is not sorted ascending" diff --git a/tests/test_atan.py b/tests/test_atan.py index b20f49b..f2813b3 100644 --- a/tests/test_atan.py +++ b/tests/test_atan.py @@ -1,6 +1,7 @@ +import math + import pytest import torch -import math import ntops from tests.skippers import skip_if_cuda_not_available @@ -15,10 +16,10 @@ def test_atan_basic(shape, dtype, device, rtol, atol): """ # atan 定义域是全实数，生成范围在 [-10, 10] 的随机数 input_tensor = (torch.rand(shape, dtype=dtype, device=device) - 0.5) * 20.0 - + reference_output = torch.atan(input_tensor) ntops_output = ntops.torch.atan(input_tensor) - + assert torch.allclose(ntops_output, reference_output, rtol=rtol, atol=atol) assert ntops_output.shape == input_tensor.shape assert ntops_output.dtype == input_tensor.dtype @@ -31,31 +32,31 @@ def test_atan_boundary_values(dtype, device): """ 边界值测试：0, 1, -1, inf, -inf """ - test_values = torch.tensor([ - 0.0, - 1.0, - -1.0, - float('inf'), - -float('inf') - ], dtype=dtype, device=device) - + test_values = torch.tensor( + [0.0, 1.0, -1.0, float("inf"), -float("inf")], dtype=dtype, device=device + ) + reference_output = torch.atan(test_values) ntops_output = ntops.torch.atan(test_values) - + # 验证数值精度 # float16 精度较低，适当放宽误差 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 - + assert torch.allclose(ntops_output, reference_output, rtol=rtol, atol=atol) - + # 显式验证特殊数学性质 # atan(0) = 0 assert torch.abs(ntops_output[0]) < 1e-6 # atan(inf) = pi/2 - assert torch.abs(ntops_output[3] - (math.pi / 2)) < (1e-3 if dtype == torch.float16 else 1e-6) + assert torch.abs(ntops_output[3] - (math.pi / 2)) < ( + 1e-3 if dtype == torch.float16 else 1e-6 + ) # atan(-inf) = -pi/2 - assert torch.abs(ntops_output[4] - (-math.pi / 2)) < (1e-3 if dtype == torch.float16 else 1e-6) + assert torch.abs(ntops_output[4] - (-math.pi / 2)) < ( + 1e-3 if dtype == torch.float16 else 1e-6 + ) @skip_if_cuda_not_available @@ -67,16 +68,16 @@ def test_atan_strided_output(dtype): device = "cuda" input_tensor = torch.randn(4, 5, 6, dtype=dtype, device=device) reference_output = torch.atan(input_tensor) - + # 创建非连续输出张量 large_tensor = torch.empty(4, 5, 8, dtype=dtype, device=device) out = large_tensor[:, :, :6] - + result = ntops.torch.atan(input_tensor, out=out) - + rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 - + assert torch.allclose(result, reference_output, rtol=rtol, atol=atol) assert result is out assert not out.is_contiguous() @@ -91,17 +92,17 @@ def test_atan_inplace_compatibility(dtype): device = "cuda" input_tensor = torch.randn(10, 10, dtype=dtype, device=device) input_copy = input_tensor.clone() - + # 链式操作 input_copy.mul_(2.0) input_copy.add_(0.5) - + reference_output = torch.atan(input_copy) ntops_output = ntops.torch.atan(input_copy) - + rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 - + assert torch.allclose(ntops_output, reference_output, rtol=rtol, atol=atol) @@ -112,9 +113,9 @@ def test_atan_nan_handling(dtype): 测试 NaN 输入处理 """ device = "cuda" - input_tensor = torch.tensor([float('nan'), 1.0], dtype=dtype, device=device) - + input_tensor = torch.tensor([float("nan"), 1.0], dtype=dtype, device=device) + output = ntops.torch.atan(input_tensor) - + assert torch.isnan(output[0]) - assert not torch.isnan(output[1]) \ No newline at end of file + assert not torch.isnan(output[1]) diff --git a/tests/test_batch_norm.py b/tests/test_batch_norm.py index 7d65564..ff521e0 100644 --- a/tests/test_batch_norm.py +++ b/tests/test_batch_norm.py @@ -1,22 +1,22 @@ import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available from tests.utils import generate_arguments + @skip_if_cuda_not_available @pytest.mark.parametrize("eps", (1e-5, 1e-3)) @pytest.mark.parametrize("affine", (True, False)) @pytest.mark.parametrize(*generate_arguments()) -def test_batch_norm( - shape, dtype, device, rtol, atol, affine, eps -): +def test_batch_norm(shape, dtype, device, rtol, atol, affine, eps): if len(shape) < 2: return - + input = torch.randn(shape, dtype=dtype, device=device) C = shape[1] - + if affine: weight = torch.randn(C, dtype=dtype, device=device) bias = torch.randn(C, dtype=dtype, device=device) @@ -28,17 +28,17 @@ def test_batch_norm( ninetoothed_output = ntops.torch.batch_norm( input, weight=weight, bias=bias, eps=eps, training=True ) - + # 调用 PyTorch 参考 # 必须指定 training=True 以强制从当前 batch 计算统计量 reference_output = torch.nn.functional.batch_norm( - input, - running_mean=None, - running_var=None, - weight=weight, - bias=bias, - training=True, - eps=eps + input, + running_mean=None, + running_var=None, + weight=weight, + bias=bias, + training=True, + eps=eps, ) - assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol) \ No newline at end of file + assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol) diff --git a/tests/test_bincount.py b/tests/test_bincount.py index 64a8c8e..df60d52 100644 --- a/tests/test_bincount.py +++ b/tests/test_bincount.py @@ -1,47 +1,52 @@ import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available + def generate_bincount_args(): # 参数组合: (size, has_weights, minlength) cases = [ (100, False, 0), (100, True, 0), (100, True, 50), - (1000, False, 2000), # minlength > max_val - (50, True, 10), # minlength < max_val - (0, False, 5), # Empty input + (1000, False, 2000), # minlength > max_val + (50, True, 10), # minlength < max_val + (0, False, 5), # Empty input ] return "size, has_weights, minlength", cases + @skip_if_cuda_not_available @pytest.mark.parametrize(*generate_bincount_args()) def test_bincount(size, has_weights, minlength): device = "cuda" if torch.cuda.is_available() else "cpu" - + # 生成随机非负整数输入 if size > 0: max_val = size # 保证值有一定的分布 - input_tensor = torch.randint(0, max_val, (size,), device=device, dtype=torch.int32) + input_tensor = torch.randint( + 0, max_val, (size,), device=device, dtype=torch.int32 + ) else: input_tensor = torch.tensor([], device=device, dtype=torch.int32) - + weights = None if has_weights and size > 0: weights = torch.randn(size, device=device, dtype=torch.float32) elif has_weights and size == 0: weights = torch.tensor([], device=device, dtype=torch.float32) - + # 运行参考实现 (PyTorch) ref_out = torch.bincount(input_tensor, weights=weights, minlength=minlength) - + # 运行 ntops 实现 ntops_out = ntops.torch.bincount(input_tensor, weights=weights, minlength=minlength) - + # 比较结果 # 注意: 浮点数比较需要一定容差 if ntops_out.is_floating_point(): assert torch.allclose(ntops_out, ref_out, atol=1e-4) else: - assert torch.equal(ntops_out, ref_out) \ No newline at end of file + assert torch.equal(ntops_out, ref_out) diff --git a/tests/test_fmax.py b/tests/test_fmax.py index ffde063..517151f 100644 --- a/tests/test_fmax.py +++ b/tests/test_fmax.py @@ -1,10 +1,11 @@ -import random import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available from tests.utils import generate_arguments + @skip_if_cuda_not_available @pytest.mark.parametrize(*generate_arguments()) def test_fmax_elementwise(shape, dtype, device, rtol, atol): @@ -14,7 +15,7 @@ def test_fmax_elementwise(shape, dtype, device, rtol, atol): # ntops 实现 ntops_result = ntops.torch.fmax(input_tensor, other_tensor) - + # PyTorch 参考实现 (torch.fmax 或 torch.maximum) # torch.fmax: 忽略 NaN # torch.maximum: 传播 NaN @@ -23,6 +24,7 @@ def test_fmax_elementwise(shape, dtype, device, rtol, atol): assert torch.allclose(ntops_result, reference_result, rtol=rtol, atol=atol) + @skip_if_cuda_not_available @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) def test_fmax_broadcasting(dtype): @@ -30,7 +32,7 @@ def test_fmax_broadcasting(dtype): device = "cuda" shape_a = (4, 1, 32) shape_b = (1, 64, 32) - + input_tensor = torch.randn(shape_a, dtype=dtype, device=device) other_tensor = torch.randn(shape_b, dtype=dtype, device=device) @@ -38,4 +40,4 @@ def test_fmax_broadcasting(dtype): reference_result = torch.maximum(input_tensor, other_tensor) assert ntops_result.shape == (4, 64, 32) - assert torch.allclose(ntops_result, reference_result) \ No newline at end of file + assert torch.allclose(ntops_result, reference_result) diff --git a/tests/test_logsumexp.py b/tests/test_logsumexp.py index 9f688e8..5011d63 100644 --- a/tests/test_logsumexp.py +++ b/tests/test_logsumexp.py @@ -30,27 +30,31 @@ def test_logsumexp(shape, dtype, device, rtol, atol, keepdim): def test_logsumexp_with_strided_output(dtype, keepdim): """测试 logsumexp 使用非连续（strided）输出张量的情况""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(4, 5, 6, dtype=dtype, device=device) dim = 2 - + # 计算参考输出 reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个具有不同 strides 的输出张量 if keepdim: # 创建一个更大的张量，然后切片得到具有非标准 strides 的子张量 large_tensor = torch.empty(4, 5, 3, dtype=dtype, device=device) - out = large_tensor[:, :, :1] # shape (4, 5, 1) 但 strides 为 (15, 3, 1) 而不是标准的 (5, 1, 1) + out = large_tensor[ + :, :, :1 + ] # shape (4, 5, 1) 但 strides 为 (15, 3, 1) 而不是标准的 (5, 1, 1) else: # 创建一个更大的张量，然后切片得到具有非标准 strides 的子张量 large_tensor = torch.empty(4, 5, 2, dtype=dtype, device=device) - out = large_tensor[:, :, 0] # shape (4, 5) 但 strides 为 (10, 2) 而不是标准的 (5, 1) - + out = large_tensor[ + :, :, 0 + ] # shape (4, 5) 但 strides 为 (10, 2) 而不是标准的 (5, 1) + # 使用 ntops 的 logsumexp，传入 strided 输出 result = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -63,22 +67,22 @@ def test_logsumexp_with_strided_output(dtype, keepdim): def test_logsumexp_with_transposed_output(dtype): """测试 logsumexp 使用转置（非连续）输出张量的情况""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(3, 4, 5, dtype=dtype, device=device) dim = 1 keepdim = True - + # 计算参考输出 reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个转置的输出张量（非连续） out_base = torch.empty(1, 3, 5, dtype=dtype, device=device) out = out_base.transpose(0, 1) # shape (3, 1, 5)，但内存布局非连续 - + # 使用 ntops 的 logsumexp result = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -91,26 +95,27 @@ def test_logsumexp_with_transposed_output(dtype): def test_logsumexp_with_strided_slice_output(dtype): """测试 logsumexp 使用步进切片（strided slice）输出张量的情况""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(2, 8, 6, dtype=dtype, device=device) dim = 2 keepdim = True - - # 计算参考输出 - reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个更大的张量，使用步进切片得到非连续的子张量 large_tensor = torch.empty(2, 8, 4, dtype=dtype, device=device) out = large_tensor[:, ::2, :1] # shape (2, 4, 1)，strides 非标准 - + # 调整输入以匹配输出的第二维 input_tensor_adjusted = input_tensor[:, ::2, :] # shape (2, 4, 6) - reference_output_adjusted = torch.logsumexp(input_tensor_adjusted, dim=dim, keepdim=keepdim) - + reference_output_adjusted = torch.logsumexp( + input_tensor_adjusted, dim=dim, keepdim=keepdim + ) + # 使用 ntops 的 logsumexp - result = ntops.torch.logsumexp(input_tensor_adjusted, dim=dim, keepdim=keepdim, out=out) - + result = ntops.torch.logsumexp( + input_tensor_adjusted, dim=dim, keepdim=keepdim, out=out + ) + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -125,19 +130,19 @@ def test_logsumexp_with_strided_slice_output(dtype): def test_logsumexp_with_contiguous_out(dtype, keepdim, dim): """测试 logsumexp 使用正常的连续（contiguous）输出张量""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(4, 5, 6, dtype=dtype, device=device) - + # 计算参考输出 reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个连续的输出张量 out = torch.empty_like(reference_output) - + # 使用 ntops 的 logsumexp，传入连续输出 result = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -152,25 +157,29 @@ def test_logsumexp_with_contiguous_out(dtype, keepdim, dim): def test_logsumexp_with_out_different_strides_dim0(dtype, keepdim): """测试 logsumexp 在 dim=0 时使用 strided 输出张量""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(6, 4, 5, dtype=dtype, device=device) dim = 0 - + # 计算参考输出 reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个具有不同 strides 的输出张量 if keepdim: large_tensor = torch.empty(3, 4, 5, dtype=dtype, device=device) - out = large_tensor[:1, :, :] # shape (1, 4, 5) 但 strides 为 (20, 5, 1) 而不是标准的 (20, 5, 1) + out = large_tensor[ + :1, :, : + ] # shape (1, 4, 5) 但 strides 为 (20, 5, 1) 而不是标准的 (20, 5, 1) else: large_tensor = torch.empty(4, 5, 2, dtype=dtype, device=device) - out = large_tensor[:, :, 0] # shape (4, 5) 但 strides 为 (10, 2) 而不是标准的 (5, 1) - + out = large_tensor[ + :, :, 0 + ] # shape (4, 5) 但 strides 为 (10, 2) 而不是标准的 (5, 1) + # 使用 ntops 的 logsumexp result = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -184,25 +193,29 @@ def test_logsumexp_with_out_different_strides_dim0(dtype, keepdim): def test_logsumexp_with_out_different_strides_dim1(dtype, keepdim): """测试 logsumexp 在 dim=1 时使用 strided 输出张量""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(3, 8, 5, dtype=dtype, device=device) dim = 1 - + # 计算参考输出 reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个具有不同 strides 的输出张量 if keepdim: large_tensor = torch.empty(3, 2, 5, dtype=dtype, device=device) - out = large_tensor[:, :1, :] # shape (3, 1, 5) 但 strides 为 (10, 5, 1) 而不是标准的 (5, 5, 1) + out = large_tensor[ + :, :1, : + ] # shape (3, 1, 5) 但 strides 为 (10, 5, 1) 而不是标准的 (5, 5, 1) else: large_tensor = torch.empty(3, 5, 3, dtype=dtype, device=device) - out = large_tensor[:, :, 0] # shape (3, 5) 但 strides 为 (15, 3) 而不是标准的 (5, 1) - + out = large_tensor[ + :, :, 0 + ] # shape (3, 5) 但 strides 为 (15, 3) 而不是标准的 (5, 1) + # 使用 ntops 的 logsumexp result = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -215,25 +228,27 @@ def test_logsumexp_with_out_different_strides_dim1(dtype, keepdim): def test_logsumexp_with_out_permuted(dtype): """测试 logsumexp 使用 permute 后的输出张量（非连续）""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(3, 4, 5, 6, dtype=dtype, device=device) dim = 2 keepdim = True - + # 计算参考输出 reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个 permute 后的输出张量（非连续） # 使用 permute(3, 2, 1, 0) 来确保输出是非连续的 out_base = torch.empty(6, 1, 4, 3, dtype=dtype, device=device) out = out_base.permute(3, 2, 1, 0) # shape (3, 4, 1, 6)，内存布局非连续 - - assert not out.is_contiguous(), "out should be non-contiguous before calling logsumexp" - + + assert not out.is_contiguous(), ( + "out should be non-contiguous before calling logsumexp" + ) + # 使用 ntops 的 logsumexp result = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -246,23 +261,23 @@ def test_logsumexp_with_out_permuted(dtype): def test_logsumexp_with_out_multiple_strides(dtype): """测试 logsumexp 使用多个维度都有非标准 strides 的输出张量""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(2, 6, 8, dtype=dtype, device=device) dim = 2 keepdim = True - + # 计算参考输出 reference_output = torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 创建一个多个维度都有非标准 strides 的输出张量 # 使用步进切片在多个维度上创建非连续张量 large_tensor = torch.empty(4, 12, 3, dtype=dtype, device=device) out = large_tensor[::2, ::2, :1] # shape (2, 6, 1)，所有维度的 strides 都非标准 - + # 使用 ntops 的 logsumexp result = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out) - + # 验证结果 rtol = 1e-3 if dtype == torch.float16 else 1e-5 atol = 1e-3 if dtype == torch.float16 else 1e-6 @@ -277,20 +292,20 @@ def test_logsumexp_with_out_multiple_strides(dtype): def test_logsumexp_out_vs_no_out(dtype, keepdim): """测试使用 out 参数和不使用 out 参数的结果一致性""" device = "cuda" - + # 创建输入张量 input_tensor = torch.randn(4, 5, 6, dtype=dtype, device=device) dim = 1 - + # 不使用 out 参数 result_no_out = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim) - + # 使用连续的 out 参数 out_contiguous = torch.empty_like(result_no_out) - result_with_out = ntops.torch.logsumexp(input_tensor, dim=dim, keepdim=keepdim, out=out_contiguous) - + result_with_out = ntops.torch.logsumexp( + input_tensor, dim=dim, keepdim=keepdim, out=out_contiguous + ) + # 验证结果一致 assert torch.allclose(result_no_out, result_with_out, rtol=1e-6, atol=1e-6) assert result_with_out is out_contiguous - - diff --git a/tests/test_lp_pool1d.py b/tests/test_lp_pool1d.py index f76207f..87f9338 100644 --- a/tests/test_lp_pool1d.py +++ b/tests/test_lp_pool1d.py @@ -43,4 +43,6 @@ def test_lp_pool1d(norm_type, ceil_mode, use_stride): ceil_mode=ceil_mode, ) - assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3, equal_nan=True) + assert torch.allclose( + ntops_output, reference_output, atol=1e-3, rtol=1e-3, equal_nan=True + ) diff --git a/tests/test_lp_pool2d.py b/tests/test_lp_pool2d.py index d341d53..e802ffb 100644 --- a/tests/test_lp_pool2d.py +++ b/tests/test_lp_pool2d.py @@ -35,11 +35,18 @@ def test_lp_pool2d(norm_type, ceil_mode, use_tuple_kernel, use_stride): s_w = random.randint(1, max(1, kernel_size[1])) stride = (s_h, s_w) else: - stride = random.randint(1, max(1, kernel_size if isinstance(kernel_size, int) else min(kernel_size))) + stride = random.randint( + 1, + max( + 1, kernel_size if isinstance(kernel_size, int) else min(kernel_size) + ), + ) else: stride = None - input_tensor = torch.randn((batch, channels, height, width), device=device, dtype=dtype) + input_tensor = torch.randn( + (batch, channels, height, width), device=device, dtype=dtype + ) ntops_output = ntops.torch.lp_pool2d( input_tensor, @@ -56,4 +63,6 @@ def test_lp_pool2d(norm_type, ceil_mode, use_tuple_kernel, use_stride): ceil_mode=ceil_mode, ) - assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3, equal_nan=True) + assert torch.allclose( + ntops_output, reference_output, atol=1e-3, rtol=1e-3, equal_nan=True + ) diff --git a/tests/test_lp_pool3d.py b/tests/test_lp_pool3d.py index 3862177..d226edc 100644 --- a/tests/test_lp_pool3d.py +++ b/tests/test_lp_pool3d.py @@ -43,7 +43,9 @@ def test_lp_pool3d(norm_type, ceil_mode, use_tuple_kernel, use_stride): else: stride = None - input_tensor = torch.randn((batch, channels, depth, height, width), device=device, dtype=dtype) + input_tensor = torch.randn( + (batch, channels, depth, height, width), device=device, dtype=dtype + ) ntops_output = ntops.torch.lp_pool3d( input_tensor, @@ -60,4 +62,6 @@ def test_lp_pool3d(norm_type, ceil_mode, use_tuple_kernel, use_stride): ceil_mode=ceil_mode, ) - assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3, equal_nan=True) + assert torch.allclose( + ntops_output, reference_output, atol=1e-3, rtol=1e-3, equal_nan=True + ) diff --git a/tests/test_max.py b/tests/test_max.py index 35b6310..f763055 100644 --- a/tests/test_max.py +++ b/tests/test_max.py @@ -18,8 +18,12 @@ def test_max_dim(shape, dtype, device, rtol, atol, keepdim): if random.choice((True, False)): dim = dim - input_tensor.ndim - ntops_values, ntops_indices = ntops.torch.max(input_tensor, dim=dim, keepdim=keepdim) - reference_values, reference_indices = torch.max(input_tensor, dim=dim, keepdim=keepdim) + ntops_values, ntops_indices = ntops.torch.max( + input_tensor, dim=dim, keepdim=keepdim + ) + reference_values, reference_indices = torch.max( + input_tensor, dim=dim, keepdim=keepdim + ) assert torch.allclose(ntops_values, reference_values, rtol=rtol, atol=atol) assert torch.equal(ntops_indices, reference_indices) diff --git a/tests/test_max_pool1d.py b/tests/test_max_pool1d.py index bb16487..52c9eff 100644 --- a/tests/test_max_pool1d.py +++ b/tests/test_max_pool1d.py @@ -1,17 +1,20 @@ import random + import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available + @skip_if_cuda_not_available @pytest.mark.parametrize("ceil_mode", [False, True]) @pytest.mark.parametrize("padding", [0, 1, 2]) -@pytest.mark.parametrize("kernel_stride", [(2, 2), (3, 2), (3, 1)]) # (kernel, stride) +@pytest.mark.parametrize("kernel_stride", [(2, 2), (3, 2), (3, 1)]) # (kernel, stride) def test_max_pool1d(ceil_mode, padding, kernel_stride): device = "cuda" dtype = torch.float32 - + kernel_size, stride = kernel_stride batch = random.randint(1, 4) @@ -27,7 +30,7 @@ def test_max_pool1d(ceil_mode, padding, kernel_stride): kernel_size=kernel_size, stride=stride, padding=padding, - ceil_mode=ceil_mode + ceil_mode=ceil_mode, ) except Exception as e: # 某些极端情况下的尺寸计算可能会导致 torch 原生报错， @@ -41,7 +44,7 @@ def test_max_pool1d(ceil_mode, padding, kernel_stride): kernel_size=kernel_size, stride=stride, padding=padding, - ceil_mode=ceil_mode + ceil_mode=ceil_mode, ) except RuntimeError: # 如果 Torch 因为输出尺寸计算报错 (例如 padding 此时过大导致 L_out < 0 等)， @@ -49,4 +52,4 @@ def test_max_pool1d(ceil_mode, padding, kernel_stride): # 这里简单起见，如果 torch 挂了，我们跳过这次 assert return - assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3) \ No newline at end of file + assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3) diff --git a/tests/test_max_pool3d.py b/tests/test_max_pool3d.py index 204776f..6466c4d 100644 --- a/tests/test_max_pool3d.py +++ b/tests/test_max_pool3d.py @@ -1,9 +1,12 @@ import random + import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available + @skip_if_cuda_not_available @pytest.mark.parametrize("ceil_mode", [False, True]) @pytest.mark.parametrize("padding", [0, 1]) @@ -11,7 +14,7 @@ def test_max_pool3d(ceil_mode, padding, use_tuple): device = "cuda" dtype = torch.float32 - + batch = random.randint(1, 2) channels = random.randint(1, 3) depth = random.randint(8, 16) @@ -25,7 +28,9 @@ def test_max_pool3d(ceil_mode, padding, use_tuple): kernel_size = random.randint(2, 3) stride = random.randint(1, 2) - input_tensor = torch.randn((batch, channels, depth, height, width), device=device, dtype=dtype) + input_tensor = torch.randn( + (batch, channels, depth, height, width), device=device, dtype=dtype + ) # Ntops implementation ntops_output = ntops.torch.max_pool3d( @@ -33,7 +38,7 @@ def test_max_pool3d(ceil_mode, padding, use_tuple): kernel_size=kernel_size, stride=stride, padding=padding, - ceil_mode=ceil_mode + ceil_mode=ceil_mode, ) # Reference implementation @@ -42,7 +47,7 @@ def test_max_pool3d(ceil_mode, padding, use_tuple): kernel_size=kernel_size, stride=stride, padding=padding, - ceil_mode=ceil_mode + ceil_mode=ceil_mode, ) - assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3) \ No newline at end of file + assert torch.allclose(ntops_output, reference_output, atol=1e-3, rtol=1e-3) diff --git a/tests/test_maximum.py b/tests/test_maximum.py index 0be30f2..53bbfd3 100644 --- a/tests/test_maximum.py +++ b/tests/test_maximum.py @@ -1,9 +1,11 @@ import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available from tests.utils import generate_arguments + @skip_if_cuda_not_available @pytest.mark.parametrize(*generate_arguments()) def test_maximum_elementwise(shape, dtype, device, rtol, atol): @@ -15,13 +17,14 @@ def test_maximum_elementwise(shape, dtype, device, rtol, atol): # 1. 运行你的 DSL 实现 # 注意：确保 ntops.torch.maximum 已经正确暴露 ntops_result = ntops.torch.maximum(input_tensor, other_tensor) - + # 2. 运行 PyTorch 参考实现 reference_result = torch.maximum(input_tensor, other_tensor) # 3. 验证结果 assert torch.allclose(ntops_result, reference_result, rtol=rtol, atol=atol) + @skip_if_cuda_not_available @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) def test_maximum_broadcasting(dtype): @@ -33,7 +36,7 @@ def test_maximum_broadcasting(dtype): shape_a = (4, 1, 32) shape_b = (1, 64, 32) expected_shape = (4, 64, 32) - + input_tensor = torch.randn(shape_a, dtype=dtype, device=device) other_tensor = torch.randn(shape_b, dtype=dtype, device=device) @@ -45,18 +48,19 @@ def test_maximum_broadcasting(dtype): # 检查数值正确性 assert torch.allclose(ntops_result, reference_result) + @skip_if_cuda_not_available def test_maximum_out_variant(): """测试 out= 参数""" device = "cuda" if not torch.cuda.is_available(): pytest.skip("CUDA not available") - + x = torch.randn(100, device=device) y = torch.randn(100, device=device) out = torch.empty_like(x) - + ntops.torch.maximum(x, y, out=out) expected = torch.maximum(x, y) - - assert torch.allclose(out, expected) \ No newline at end of file + + assert torch.allclose(out, expected) diff --git a/tests/test_mean.py b/tests/test_mean.py index 26ef10f..2ebba2f 100644 --- a/tests/test_mean.py +++ b/tests/test_mean.py @@ -1,13 +1,14 @@ import random + import pytest import torch -import ntops -# 假设你在 ntops 包里导出了 mean -from ntops.torch.mean import mean as ntops_mean +# 假设你在 ntops 包里导出了 mean +from ntops.torch.mean import mean as ntops_mean from tests.skippers import skip_if_cuda_not_available from tests.utils import generate_arguments + @skip_if_cuda_not_available @pytest.mark.parametrize(*generate_arguments()) @pytest.mark.parametrize("keepdim", (False, True)) @@ -28,6 +29,7 @@ def test_mean_dim(shape, dtype, device, rtol, atol, keepdim): assert torch.allclose(ntops_value, reference_value, rtol=rtol, atol=atol) + @skip_if_cuda_not_available @pytest.mark.parametrize(*generate_arguments()) def test_mean_global(shape, dtype, device, rtol, atol): @@ -42,15 +44,18 @@ def test_mean_global(shape, dtype, device, rtol, atol): assert torch.allclose(ntops_value, reference_value, rtol=rtol, atol=atol) + @skip_if_cuda_not_available def test_mean_int_input(): # 特别测试：整数输入应该产生浮点输出 device = "cuda" shape = (1024, 1024) input_tensor = torch.randint(0, 10, shape, device=device, dtype=torch.int32) - + ntops_value = ntops_mean(input_tensor) - reference_value = torch.mean(input_tensor.float()) # torch.mean(int)在旧版本可能报错或行为不同，通常需转float - + reference_value = torch.mean( + input_tensor.float() + ) # torch.mean(int)在旧版本可能报错或行为不同，通常需转float + assert ntops_value.is_floating_point() - assert torch.allclose(ntops_value, reference_value) \ No newline at end of file + assert torch.allclose(ntops_value, reference_value) diff --git a/tests/test_median.py b/tests/test_median.py index 93bc0d7..a94b2d3 100644 --- a/tests/test_median.py +++ b/tests/test_median.py @@ -1,10 +1,13 @@ import random + import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available from tests.utils import _random_shape + def generate_median_args(): args = [] # Median 也是一种选择排序，复杂度较高，测试时控制 dim_size @@ -16,7 +19,7 @@ def generate_median_args(): for _ in range(5): shape = _random_shape(ndim) dim = random.randint(0, ndim - 1) - + # 限制 dim_size，原因同 argsort if shape[dim] > 128: shape_list = list(shape) @@ -26,6 +29,7 @@ def generate_median_args(): args.append((shape, dim, dtype, device, rtol, atol)) return "shape, dim, dtype, device, rtol, atol", args + @skip_if_cuda_not_available @pytest.mark.parametrize(*generate_median_args()) def test_median(shape, dim, dtype, device, rtol, atol): @@ -33,7 +37,7 @@ def test_median(shape, dim, dtype, device, rtol, atol): # 1. 调用 ntops 实现 ntops_vals, ntops_idxs = ntops.torch.median(input_tensor, dim=dim, keepdim=False) - + # 2. 调用 PyTorch 参考实现 ref_vals, ref_idxs = torch.median(input_tensor, dim=dim, keepdim=False) @@ -44,8 +48,12 @@ def test_median(shape, dim, dtype, device, rtol, atol): # 4. 验证索引 (Indices) # 注意：如果有重复值，median 的 index 可能不唯一。 # 更稳健的测试方法是：用 ntops 返回的 index 去 input 取值，看取出来的值是否等于 ref_val - gathered_vals = torch.gather(input_tensor, dim, ntops_idxs.unsqueeze(dim)).squeeze(dim) - assert torch.allclose(gathered_vals, ref_vals, rtol=rtol, atol=atol), "Gathered values from indices mismatch" - + gathered_vals = torch.gather(input_tensor, dim, ntops_idxs.unsqueeze(dim)).squeeze( + dim + ) + assert torch.allclose(gathered_vals, ref_vals, rtol=rtol, atol=atol), ( + "Gathered values from indices mismatch" + ) + # 如果数据没有重复值，可以直接比较 index (可选) - # assert torch.equal(ntops_idxs, ref_idxs) \ No newline at end of file + # assert torch.equal(ntops_idxs, ref_idxs) diff --git a/tests/test_stack.py b/tests/test_stack.py index 592cc71..bf47170 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -1,27 +1,29 @@ import random + import pytest import torch + import ntops from tests.skippers import skip_if_cuda_not_available + @skip_if_cuda_not_available @pytest.mark.parametrize("dim", [0, 1, 2]) def test_stack(dim): device = "cuda" dtype = torch.float32 - + # 构造测试参数 num_tensors = random.randint(2, 5) shape = (random.randint(16, 64), random.randint(16, 64), random.randint(16, 64)) - + # 确保 dim 不越界 (ndim + 1 因为 stack 会增加维度) if dim > len(shape): return # 创建输入列表 tensors = [ - torch.randn(shape, device=device, dtype=dtype) - for _ in range(num_tensors) + torch.randn(shape, device=device, dtype=dtype) for _ in range(num_tensors) ] # Reference implementation @@ -32,6 +34,6 @@ def test_stack(dim): # 验证一致性 assert torch.allclose(ntops_output, reference_output, atol=1e-6, rtol=1e-6) - + # 验证内存是否连续（Stack 的输出通常是连续的） - assert ntops_output.is_contiguous() \ No newline at end of file + assert ntops_output.is_contiguous()