并行编程实战—CUDA中的Tile编程之五元素级运算和广播

原创于 2026-09-06 12:17:53 发布 · 209 阅读

4 ·

本内容遵循CC 4.0 BY-SA版权协议

收录于

并行编程

一、逐元素运算

既然Tile支持自己处理数据分配，那么就必须要有一些数据处理的标准。否则什么样的数据都往里面塞，到最后也不知道算出个什么。对Tile来说，它支持标准的逐元素算术运算，但如果两个操作数的形状不同时，较小的操作数需要进行广播从而匹配较大的操作数的形状，然后才能进行计算。
也就是说，Tile中的运算，必须是满足特定的“相似”的要求，才能够进行展开。

二、Broadcasting

这种“相似”其实就是以广播的形式进行展开的。这种广播和NumPy的语义保持了一致。即标量会被复制到Tile的每个元素；单例维度（长度为1）会被拉伸以匹配另一操作数的对应维度；低秩操作数会通过将缺失的前导维度视为单例维度，对齐到高秩操作数的尾部维度。如果两个对应的维度都是非单例且不相等，则运算是非法的。可以总结为：

维度对齐
对齐的方向从尾部（最右）维度开始向左逐维进行比较，当维度不足时，在左侧补1直到两数组维度相同
单维兼容
即对齐的维度必须满足相等或其中一方为1（可扩展）或一方维度不存在
错误判定
如果任一维度不相等或一方不为1，则抛出异常
输出操作数
操作数双方各维度对应的最大值

下面看一个例子：

A(5,)和B() 可广播，输出结果： C（5,） #标量自动扩展
A(2, 3)和B(3,)，可广播，输出结果：C(2, 3)#B补为 (1,3)
A(2, 1)和B(1, 3)，可广播，输出结果：C(2, 3)
A(2, 3)和B(3, 2)，不可广播，最右侧维度3!=2，且无1
A(4, 1, 6)和B(7, 5)，不可广播。#B补后为 (1,7,5) ，5!=6

三、算术运算符

有了广播以后，就可以对Tile进行逐元素的算术操作了。这和线代里的点积和叉积的规则是类似的。Tile和Tile按照广播进行处理形成一个新的Tile。标量与Tile操作时，标题广播到每个元素。即在可操作运算的前提下，不同的操作数会选择保留更多的操作数的那种类型。

Tile与Tile运算
输出的Tile类型为精度更高或范围更大的类型：如int和float，输出为float。short和int，输出为int
标量和Tile
当标量与Tile中的类型一致时，以Tile中元素的类型为准；如果不一致，需要进行缩窄操作标量才能匹配时，Python会将输出提升为容纳两者的类型（即整体的扩大），而C++则认为其非良构而拒绝处理。
举一个例子：Tile中元素类型为整型int，而标量为3.7，Python会提升为float（这和C++中的默认运算操作有些类似。）,而C++则抛出一个错误。

所以强烈建议在实际的应用中尽量保持操作数双主的类型匹配，如果不匹配，则进行显式的类型转换。

四、例程

下面看一个简单的例程：

#!/usr/bin/env python3

from __future__ import annotations

import argparse
from typing import Any, Callable

import cupy as cp
import cuda.tile as ct


DEFAULT_ROWS = 7
DEFAULT_COLS = 5
DEFAULT_TILE = 4


@ct.kernel
def same_shape_mul_kernel(x, y, out, tile_rows: ct.Constant[int], tile_cols: ct.Constant[int]):
    block_row = ct.bid(0)
    block_col = ct.bid(1)
    x_tile = ct.load(x, index=(block_row, block_col), shape=(tile_rows, tile_cols), padding_mode=ct.PaddingMode.ZERO)
    y_tile = ct.load(y, index=(block_row, block_col), shape=(tile_rows, tile_cols), padding_mode=ct.PaddingMode.ZERO)
    ct.store(out, index=(block_row, block_col), tile=x_tile * y_tile)


@ct.kernel
def broadcast_mul_kernel(x, y_column, out, tile_rows: ct.Constant[int], tile_cols: ct.Constant[int]):
    block_row = ct.bid(0)
    block_col = ct.bid(1)
    x_tile = ct.load(x, index=(block_row, block_col), shape=(tile_rows, tile_cols), padding_mode=ct.PaddingMode.ZERO)
    y_tile = ct.load(y_column, index=(block_row, 0), shape=(tile_rows, 1), padding_mode=ct.PaddingMode.ZERO)
    ct.store(out, index=(block_row, block_col), tile=x_tile * y_tile)


@ct.kernel
def rank_promotion_add_kernel(
    x,
    y,
    out,
    batch_tile: ct.Constant[int],
    rows_tile: ct.Constant[int],
    cols_tile: ct.Constant[int],
):

    x_tile = ct.load(x, index=(0, 0), shape=(rows_tile, cols_tile), padding_mode=ct.PaddingMode.ZERO)
    y_tile = ct.load(
        y,
        index=(0, 0, 0),
        shape=(batch_tile, 1, cols_tile),
        padding_mode=ct.PaddingMode.ZERO,
    )
    x_tile = ct.broadcast_to(x_tile, (batch_tile, rows_tile, cols_tile))
    ct.store(out, index=(0, 0, 0), tile=x_tile + y_tile)


@ct.kernel
def scalar_mul_kernel(x, scalar: ct.Constant[int], out, tile: ct.Constant[int]):
    block = ct.bid(0)
    x_tile = ct.load(x, index=(block,), shape=(tile,), padding_mode=ct.PaddingMode.ZERO)
    ct.store(out, index=(block,), tile=x_tile * scalar)


@ct.kernel
def dtype_promotion_mul_kernel(x, y, out, tile: ct.Constant[int]):
    block = ct.bid(0)
    x_tile = ct.load(x, index=(block,), shape=(tile,), padding_mode=ct.PaddingMode.ZERO)
    y_tile = ct.load(y, index=(block,), shape=(tile,), padding_mode=ct.PaddingMode.ZERO)
    ct.store(out, index=(block,), tile=x_tile * y_tile)


@ct.kernel
def scalar_dtype_promotion_kernel(x, int_scalar: ct.Constant[int], float_scalar: ct.Constant[float], out_int, out_float, tile: ct.Constant[int]):
    block = ct.bid(0)
    x_tile = ct.load(x, index=(block,), shape=(tile,), padding_mode=ct.PaddingMode.ZERO)
    ct.store(out_int, index=(block,), tile=x_tile + int_scalar)
    ct.store(out_float, index=(block,), tile=x_tile + float_scalar)


def _ceil_div(value: int, divisor: int) -> int:
    return (value + divisor - 1) // divisor


def _next_power_of_two(value: int) -> int:
    return 1 << (value - 1).bit_length()


def _launch(kernel: Callable[..., Any], args: tuple[Any, ...], grid: tuple[int, ...], stream: Any) -> None:
    ct.launch(stream, grid, kernel, args)
    stream.synchronize()


def _assert_equal(label: str, actual: cp.ndarray, expected: cp.ndarray) -> None:
    try:
        cp.testing.assert_array_equal(actual, expected)
    except AssertionError as exc:
        max_error = float(cp.max(cp.abs(actual.astype(cp.float64) - expected.astype(cp.float64))).get())
        raise RuntimeError(f"{label} failue (max {max_error:g})") from exc
    print(f"[PASS] {label}")


def _device_name(device: int) -> str:
    name = cp.cuda.runtime.getDeviceProperties(device)["name"]
    return name.decode() if isinstance(name, bytes) else str(name)


def run_all_demos(rows: int = DEFAULT_ROWS, cols: int = DEFAULT_COLS, tile: int = DEFAULT_TILE, device: int = 0) -> None:
    if rows <= 0 or cols <= 0 or tile <= 0:
        raise ValueError("rows, cols, and tile must be positive")
    if tile & (tile - 1):
        raise ValueError("tile must be a power of two for cuTile")
    if device < 0:
        raise ValueError("device must be non-negative")

    cp.cuda.Device(device).use()
    stream = cp.cuda.get_current_stream()
    print(f"GPU: {_device_name(device)} (device {device})")
    print(f"shape=({rows}, {cols}), tile={tile}\n")

    grid_2d = (_ceil_div(rows, tile), _ceil_div(cols, tile))
    x = cp.full((rows, cols), 7, dtype=cp.int32)
    y = cp.full((rows, cols), 3, dtype=cp.int32)
    out = cp.empty_like(x)
    _launch(same_shape_mul_kernel, (x, y, out, tile, tile), grid_2d, stream)
    _assert_equal("1. same-shape multiplication", out, x * y)

    y_column = cp.full((rows, 1), 3, dtype=cp.int32)
    _launch(broadcast_mul_kernel, (x, y_column, out, tile, tile), grid_2d, stream)
    _assert_equal("2. column broadcasting", out, x * y_column)

    batch = 3
    rank_rows_tile = _next_power_of_two(rows)
    rank_cols_tile = _next_power_of_two(cols)
    rank_batch_tile = _next_power_of_two(batch)
    x_rank = cp.full((rows, cols), 3, dtype=cp.int32)
    y_rank = cp.full((batch, 1, cols), 5, dtype=cp.int32)
    out_rank = cp.empty((batch, rows, cols), dtype=cp.int32)
    x_rank_padded = cp.zeros((rank_rows_tile, rank_cols_tile), dtype=cp.int32)
    x_rank_padded[:rows, :cols] = x_rank
    y_rank_padded = cp.zeros((rank_batch_tile, 1, rank_cols_tile), dtype=cp.int32)
    y_rank_padded[:batch, :, :cols] = y_rank
    out_rank_padded = cp.empty((rank_batch_tile, rank_rows_tile, rank_cols_tile), dtype=cp.int32)
    _launch(
        rank_promotion_add_kernel,
        (x_rank_padded, y_rank_padded, out_rank_padded, rank_batch_tile, rank_rows_tile, rank_cols_tile),
        (1,),
        stream,
    )
    out_rank[...] = out_rank_padded[:batch, :rows, :cols]
    _assert_equal("3. rank-promotion addition", out_rank, x_rank + y_rank)

    x_vec = cp.full((rows * cols,), 7, dtype=cp.int32)
    out_vec = cp.empty_like(x_vec)
    _launch(scalar_mul_kernel, (x_vec, 2, out_vec, tile), (_ceil_div(x_vec.size, tile),), stream)
    _assert_equal("4. scalar multiplication", out_vec, x_vec * 2)

    x_i32 = cp.full((rows * cols,), 7, dtype=cp.int32)
    x_i64 = cp.full((rows * cols,), 3, dtype=cp.int64)
    out_i64 = cp.empty_like(x_i64)
    _launch(dtype_promotion_mul_kernel, (x_i32, x_i64, out_i64, tile), (_ceil_div(x_i32.size, tile),), stream)
    _assert_equal("5. int32/int64 promotion", out_i64, x_i32 * x_i64)

    out_int = cp.empty_like(x_vec)
    out_float = cp.empty(x_vec.shape, dtype=cp.float32)
    _launch(scalar_dtype_promotion_kernel, (x_vec, 2, 2.5, out_int, out_float, tile), (_ceil_div(x_vec.size, tile),), stream)
    _assert_equal("6a. integer scalar promotion", out_int, x_vec + 2)
    _assert_equal("6b. float scalar promotion", out_float, x_vec.astype(cp.float32) + cp.float32(2.5))
    print("\ncuTile test end。")


sameShapeMulKernel = same_shape_mul_kernel
broadcastMulKernel = broadcast_mul_kernel
rankPromotionAddKernel = rank_promotion_add_kernel
scalarMulKernel = scalar_mul_kernel
dtypePromotionMulKernel = dtype_promotion_mul_kernel
scalarDtypePromoKernel = scalar_dtype_promotion_kernel
runAllDemos = run_all_demos

def _positive_int(value: str) -> int:
    number = int(value)
    if number <= 0:
        raise argparse.ArgumentTypeError("must be greater than zero")
    return number

def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--rows", type=_positive_int, default=DEFAULT_ROWS)
    parser.add_argument("--cols", type=_positive_int, default=DEFAULT_COLS)
    parser.add_argument("--tile", type=_positive_int, default=DEFAULT_TILE)
    parser.add_argument("--device", type=int, default=0)
    args = parser.parse_args()
    run_all_demos(args.rows, args.cols, args.tile, args.device)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())