[Spielopoly]

This PR adds a Python code-generation backend and cuTile support to DaCe. Supported computations in DaCe programs can now be lowered into tile-based GPU kernels, with generated Python code handling host execution, data movement, and kernel launches.

The implementation introduces three main components:

• Python backend: Generates executable Python for control flow, maps, tasklets, and nested SDFGs, with NumPy and CuPy array management. This backend is also usable independently of cuTile.
• cuTile code generation: Adds cuTile implementations of the existing tile operations, including arithmetic, conditional selection, gathers and scatters, and reductions.
• VectorizeCuTile: Provides a single entry point that combines canonicalization, multidimensional vectorization, GPU data management, and cuTile lowering. It supports symbolic problem sizes and masked handling of incomplete tiles.

Generated programs can combine cuTile kernels with CuPy implementations of BLAS, linear algebra, and reductions.

For example, the following program computes a scaled addition of two matrices using symbolic dimensions and two-dimensional tiles:

import dace
import numpy as np
from dace.transformation.passes.vectorization import VectorizeCuTile

M = dace.symbol("M")
N = dace.symbol("N")

@dace.program
def scaled_add(
    alpha: dace.float32,
    A: dace.float32[M, N],
    B: dace.float32[M, N],
    C: dace.float32[M, N],
):
    C[:] = alpha * A + B

sdfg = scaled_add.to_sdfg()
VectorizeCuTile(widths=(16, 32)).apply_pass(sdfg, {})
compiled = sdfg.compile()

rng = np.random.default_rng(42)
shape = (257, 513)
a = rng.random(shape, dtype=np.float32)
b = rng.random(shape, dtype=np.float32)
c = np.empty_like(a)

compiled(
    alpha=np.float32(2.5),
    A=a,
    B=b,
    C=c,
    M=shape[0],
    N=shape[1],
)

VectorizeCuTile lowers the computation into cuTile Python kernels operating on 16 × 32 tiles. The generated program handles transfers for NumPy inputs and masks incomplete tiles at the array boundaries. Setting use_gpu_storage=True allows callers to pass CuPy arrays directly and keep data on the GPU between calls. Tile widths must be powers of two; problem sizes need not be divisible by those widths.

The example produces a cuTile Python kernel like the following (simplified for clarity):

import cuda.tile as ct
import cupy as cp

@ct.kernel
def scaled_add_kernel(A, B, alpha, C, M, N):
    M = ct.load(M, (0,), shape=()).item()
    N = ct.load(N, (0,), shape=()).item()
    alpha = ct.load(alpha, (0,), shape=()).item()

    tile_m = ct.bid(0)
    tile_n = ct.bid(1)

    rows = ct.arange(16, dtype=ct.int32) + tile_m * 16
    cols = ct.arange(32, dtype=ct.int32) + tile_n * 32
    rows = ct.broadcast_to(rows[:, None], (16, 32))
    cols = ct.broadcast_to(cols[None, :], (16, 32))
    mask = (rows < M) & (cols < N)

    a = ct.load(
        A, index=(tile_m, tile_n), shape=(16, 32),
        padding_mode=ct.PaddingMode.ZERO,
    )
    b = ct.load(
        B, index=(tile_m, tile_n), shape=(16, 32),
        padding_mode=ct.PaddingMode.ZERO,
    )
    a = ct.where(mask, a, 0)
    b = ct.where(mask, b, 0)

    scaled = ct.where(mask, alpha * a, False)
    result = ct.where(mask, scaled + b, False)
    ct.scatter(C, (rows, cols), result, mask=mask)


def scaled_add(A, B, C, M, N, alpha):
    A_gpu = cp.asarray(A)
    B_gpu = cp.asarray(B)
    C_gpu = cp.empty((M, N), dtype=cp.float32)
    grid = ((M + 15) // 16, (N + 31) // 32, 1)

    if min(grid) > 0:
        ct.launch(
            cp.cuda.get_current_stream(),
            grid,
            scaled_add_kernel,
            (
                A_gpu,
                B_gpu,
                cp.asarray(alpha, dtype=cp.float32).reshape(1),
                C_gpu,
                cp.asarray(M, dtype=cp.int32).reshape(1),
                cp.asarray(N, dtype=cp.int32).reshape(1),
            ),
        )

    cp.cuda.get_current_stream().synchronize()
    C_gpu.get(out=C)

[Copilot]

Copilot wasn't able to review this pull request because it exceeds the maximum number of lines (20,000). Try reducing the number of changed lines and requesting a review from Copilot again.