ILGPU分块矩阵乘法Kernel输出错误，请求排查原因

我从ILGPU示例中获取了**分块形式（tiled form）**的矩阵乘法Kernel代码，用来计算以下矩阵的乘积：

a = |1  2  3  4|
    |5  6  7  8|

b = | 9  10|
    |11  12|
    |13  14|
    |15  16|

根据计算，正确结果应为：

130   140
   322   348

但我的代码输出结果为：

Result of matrix multiplication:
  31   34
 111  122

我使用了与示例完全相同的Kernel和tile-size，请问错误出在哪里？

源代码 #

using ILGPU;
using ILGPU.Runtime;
using System;

public static class Program
{
    // Define tile size for matrix multiplication (2x2 tiles for simplicity)
    // This determines how many elements will be loaded into shared memory at a time.
    const int TILE_SIZE = 2;

    // Kernel function to perform tiled matrix multiplication on the GPU
    static void MatrixMultiplyTiledKernel(
            ArrayView2D<int, Stride2D.DenseX> aView, // Input matrix A in GPU memory
            ArrayView2D<int, Stride2D.DenseX> bView, // Input matrix B in GPU memory
            ArrayView2D<int, Stride2D.DenseX> cView) // Output matrix C in GPU memory
    {
        // Get the global thread index in the 2D grid
        Index2D global = Grid.GlobalIndex.XY;

        // Get the thread indices within the thread group (local indices)
        int x = Group.IdxX; // Local row index within the tile
        int y = Group.IdxY; // Local column index within the tile

        // Allocate shared memory for storing tiles of matrix A and matrix B
        var aTile = SharedMemory.Allocate2D<int, Stride2D.DenseX>(
            new Index2D(TILE_SIZE, TILE_SIZE),
            new Stride2D.DenseX(TILE_SIZE)
        );
        var bTile = SharedMemory.Allocate2D<int, Stride2D.DenseX>(
            new Index2D(TILE_SIZE, TILE_SIZE),
            new Stride2D.DenseX(TILE_SIZE)
        );

        // Initialize the accumulation variable for the result of C[global.X, global.Y]
        var sum = 0;

        // Loop over the tiles of A and B matrices
        // The loop increments by TILE_SIZE to process one tile at a time
        for (var i = 0; i < aView.IntExtent.X; i += TILE_SIZE)
        {
            // Load the corresponding tile of A into shared memory
            if (global.X < aView.IntExtent.X && y + i < aView.IntExtent.Y)
                aTile[x, y] = aView[global.X, y + i];
            else
                aTile[x, y] = 0; // Pad with zeros if out of bounds

            // Load the corresponding tile of B into shared memory
            if (x + i < bView.IntExtent.X && global.Y < bView.IntExtent.Y)
                bTile[x, y] = bView[x + i, global.Y];
            else
                bTile[x, y] = 0; // Pad with zeros if out of bounds

            // Ensure all threads in the thread group have finished loading tiles
            Group.Barrier();

            // Perform computation on the tiles
            for (var k = 0; k < TILE_SIZE; k++)
            {
                // Multiply elements of A and B tiles and accumulate the result
                sum += aTile[new Index2D(x, k)] * bTile[new Index2D(k, y)];
            }

            // Synchronize threads before loading the next tile
            Group.Barrier();
        }

        // Write the computed result to the output matrix C
        if (global.X < cView.IntExtent.X && global.Y < cView.IntExtent.Y)
        {
            cView[global] = sum; // Store the result in the appropriate position
        }
    }

    static void Main()
    {
        // Create an ILGPU context (manages devices and resources)
        using (var context = Context.CreateDefault())
        {
            // Select the preferred accelerator (GPU or CPU)
            using (var accelerator = context.GetPreferredDevice(preferCPU: true).CreateAccelerator(context))
            {
                try
                {
                    // Initialize sample input matrices (4x4)
                    int[,] a = {
                        { 1, 2, 3, 4},
                        { 5, 6, 7, 8}
                    };
                    int[,] b = {
                        { 9, 10},
                        { 11, 12},
                        { 13, 14},
                        { 15, 16}
                    };

                    int m = a.GetLength(0);
                    int ka = a.GetLength(1);
                    int kb = b.GetLength(0);
                    int n = b.GetLength(1);

                    int[,] hostResult = new int[m, n];

                    // Define the group size (number of threads per tile) and number of groups
                    Index2D groupSize = new Index2D(TILE_SIZE, TILE_SIZE); // Threads per group (block)
                    Index2D numGroups = new Index2D((m + TILE_SIZE - 1) / TILE_SIZE, (n + TILE_SIZE - 1) / TILE_SIZE); // Total number of thread groups

                    // Allocate device memory for input matrices A, B, and output matrix C
                    MemoryBuffer2D<int, Stride2D.DenseX> aBuffer = accelerator.Allocate2DDenseX<int>(new Index2D(m, ka));
                    MemoryBuffer2D<int, Stride2D.DenseX> bBuffer = accelerator.Allocate2DDenseX<int>(new Index2D(kb, n));
                    MemoryBuffer2D<int, Stride2D.DenseX> cBuffer = accelerator.Allocate2DDenseX<int>(new Index2D(m, n));

                    try
                    {
                        // Copy input matrices from host (CPU) to device (GPU) memory
                        aBuffer.CopyFromCPU(a);
                        bBuffer.CopyFromCPU(b);

                        // Load and precompile the kernel function
                        var loadedKernel = accelerator.LoadStreamKernel<
                            ArrayView2D<int, Stride2D.DenseX>,
                            ArrayView2D<int, Stride2D.DenseX>,
                            ArrayView2D<int, Stride2D.DenseX>>(MatrixMultiplyTiledKernel);

                        // Launch the kernel function on the GPU
                        // Specify the number of thread groups and threads per group
                        loadedKernel((numGroups, groupSize), aBuffer, bBuffer, cBuffer);

                        // Wait for the GPU to complete execution
                        accelerator.Synchronize();

                        // Retrieve the result matrix from GPU memory back to host memory
                        cBuffer.CopyToCPU(hostResult);

                        // Print the result matrix
                        Console.WriteLine("Result of matrix multiplication:");
                        for (int i = 0; i < m; i++)
                        {
                            for (int j = 0; j < n; j++)
                            {
                                Console.Write($"{hostResult[i, j],4} ");
                            }
                            Console.WriteLine();
                        }
                    }
                    finally
                    {
                        // Dispose of GPU resources to free up memory
                        aBuffer.Dispose();
                        bBuffer.Dispose();
                        cBuffer.Dispose();
                    }
                }
                finally
                {
                    // Dispose of the accelerator and context
                    accelerator.Dispose();
                }
            }
        }

        // Wait for user input before closing the console
        Console.ReadLine();
    }
}

核心错误：分块循环的终止条件错误 #

当前循环条件为 for (var i = 0; i < aView.IntExtent.X; i += TILE_SIZE)，其中aView.IntExtent.X是矩阵A的行数（2），但矩阵乘法需要遍历的是A的列数/ B的行数（4）——这是矩阵乘法中累加维度的长度。

错误导致循环仅执行1次（仅处理前2列的元素），只累加了一半的乘积项，最终结果为正确值的前半部分。

修正方法 #

将循环终止条件改为遍历A的列数，即把i < aView.IntExtent.X替换为i < aView.IntExtent.Y：

// 修正：循环终止条件改为遍历A的列数（即B的行数）
for (var i = 0; i < aView.IntExtent.Y; i += TILE_SIZE)

修正后的完整Kernel代码 #

static void MatrixMultiplyTiledKernel(
        ArrayView2D<int, Stride2D.DenseX> aView, // Input matrix A in GPU memory
        ArrayView2D<int, Stride2D.DenseX> bView, // Input matrix B in GPU memory
        ArrayView2D<int, Stride2D.DenseX> cView) // Output matrix C in GPU memory
{
    // Get the global thread index in the 2D grid
    Index2D global = Grid.GlobalIndex.XY;

    // Get the thread indices within the thread group (local indices)
    int x = Group.IdxX; // Local row index within the tile
    int y = Group.IdxY; // Local column index within the tile

    // Allocate shared memory for storing tiles of matrix A and matrix B
    var aTile = SharedMemory.Allocate2D<int, Stride2D.DenseX>(
        new Index2D(TILE_SIZE, TILE_SIZE),
        new Stride2D.DenseX(TILE_SIZE)
    );
    var bTile = SharedMemory.Allocate2D<int, Stride2D.DenseX>(
        new Index2D(TILE_SIZE, TILE_SIZE),
        new Stride2D.DenseX(TILE_SIZE)
    );

    // Initialize the accumulation variable for the result of C[global.X, global.Y]
    var sum = 0;

    // 修正：循环终止条件改为遍历A的列数（即B的行数）
    for (var i = 0; i < aView.IntExtent.Y; i += TILE_SIZE)
    {
        // Load the corresponding tile of A into shared memory
        if (global.X < aView.IntExtent.X && y + i < aView.IntExtent.Y)
            aTile[x, y] = aView[global.X, y + i];
        else
            aTile[x, y] = 0; // Pad with zeros if out of bounds

        // Load the corresponding tile of B into shared memory
        if (x + i < bView.IntExtent.X && global.Y < bView.IntExtent.Y)
            bTile[x, y] = bView[x + i, global.Y];
        else
            bTile[x, y] = 0; // Pad with zeros if out of bounds

        // Ensure all threads in the thread group have finished loading tiles
        Group.Barrier();

        // Perform computation on the tiles
        for (var k = 0; k < TILE_SIZE; k++)
        {
            // Multiply elements of A and B tiles and accumulate the result
            sum += aTile[x, k] * bTile[k, y];
        }

        // Synchronize threads before loading the next tile
        Group.Barrier();
    }

    // Write the computed result to the output matrix C
    if (global.X < cView.IntExtent.X && global.Y < cView.IntExtent.Y)
    {
        cView[global] = sum; // Store the result in the appropriate position
    }
}

修正后，循环会执行2次（i=0和i=2），完整累加所有4个维度的乘积项，输出结果将与预期一致。

内容的提问来源于stack exchange，提问作者user366312