Value Clipping

原始题目：LeetGPU - Value Clipping

题目描述

编写一个 GPU 程序，对一维输入向量执行裁剪操作。给定输入张量（形状 $[N]$，类型 float32），将每个元素裁剪到指定范围 $[lo, hi]$ 内：

$$\text{clip}(x, lo, hi) = \begin{cases} lo, & \text{if } x < lo \\ x, & \text{if } lo \le x \le hi \\ hi, & \text{if } x > hi \end{cases}$$

该操作确保所有值落在指定范围内，广泛用于 ML 中的激活值稳定化和量化前预处理。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 张量中。

示例

示例 1

Input:  input = [1.5, -2.0, 3.0, 4.5], lo = 0.0, hi = 3.5
Output: [1.5, 0.0, 3.0, 3.5]

示例 2

Input:  input = [-1.0, 2.0, 5.0], lo = -0.5, hi = 2.5
Output: [-0.5, 2.0, 2.5]

约束条件

• $1 \le N \le 100{,}000$。
• $-10^6 \le input[i] \le 10^6$。
• $lo \le hi$。
• 性能测试在 $N = 100{,}000$ 的规模下进行。

解题思路

裁剪操作是典型的逐元素 Map 操作。每个元素经过两次比较（下界和上界），可以使用 min + max 的组合或 clamp。在 CUDA 中 fminf 和 fmaxf 是硬件支持的指令。与 ReLU 类似，这是一个内存带宽受限的内核。

代码实现

CUDA

#include <cuda_runtime.h>

// 解法一：基础逐元素裁剪
__global__ void clip_kernel(const float* input, float* output, float lo, float hi, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        output[i] = fmaxf(fminf(input[i], hi), lo);
    }
}

// 解法二：Grid-Stride Loop
__global__ void clip_grid_stride(const float* input, float* output, float lo, float hi, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < N; i += stride) {
        output[i] = fmaxf(fminf(input[i], hi), lo);
    }
}

extern "C" void solve(const float* input, float* output, float lo, float hi, int N) {
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
    clip_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, lo, hi, N);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl
import torch


@triton.jit
def clip_kernel(
    input_ptr, output_ptr,
    lo, hi,
    N: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < N
    x = tl.load(input_ptr + idx, mask=mask)
    tl.store(output_ptr + idx, tl.minimum(tl.maximum(x, lo), hi), mask=mask)


def solve(input: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
    N = input.numel()
    output = torch.empty_like(input)
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    clip_kernel[grid](input, output, lo, hi, N, BLOCK_SIZE=1024)
    return output