Numba-dpex generating incorrect results for sliding window matmul kernel

[diptorupd]

Here's an example that I can't get to work: it's basically a port of this numba.cuda example "fast matmul" where threads load cooperatively two arrays of shared memories.

import sklearn_numba_dpex
import numba_dpex as dpex
import numpy as np
import dpctl.tensor as dpt

square_block_side = 2
work_group_size = (square_block_side, square_block_side)
dtype = np.float32


@dpex.kernel
def matmul(
    X,         # IN READ-ONLY    (X_n_rows, n_cols)
    y,         # IN READ-ONLY    (n_cols, y_n_rows),
    result,    # OUT             (X_n_rows, y_n_rows)
):

    X_n_rows = X.shape[0]
    Y_n_cols = y.shape[1]
    n_cols = X.shape[1]

    result_row_idx = dpex.get_global_id(0)
    result_col_idx = dpex.get_global_id(1)

    local_row_idx = dpex.get_local_id(0)
    local_col_idx = dpex.get_local_id(1)

    n_blocks_for_cols = n_cols // square_block_side
    if (n_cols % square_block_side) > 0:
        n_blocks_for_cols += 1

    X_sliding_window = dpex.local.array(shape=work_group_size, dtype=dtype)
    Y_sliding_window = dpex.local.array(shape=work_group_size, dtype=dtype)

    output = dtype(0)

    for block_idx in range(n_blocks_for_cols):
        X_sliding_window[local_row_idx, local_col_idx] = dtype(0)
        Y_sliding_window[local_row_idx, local_col_idx] = dtype(0)
        if (result_row_idx < X_n_rows) and (
            (local_col_idx + (square_block_side * block_idx)) < n_cols
        ):
            X_sliding_window[local_row_idx, local_col_idx] = X[
                result_row_idx, local_col_idx + (square_block_side * block_idx)
            ]

        if (result_col_idx < Y_n_cols) and (
            (local_row_idx + (square_block_side * block_idx)) < n_cols
        ):
            Y_sliding_window[local_row_idx, local_col_idx] = y[
                local_row_idx + (square_block_side * block_idx), result_col_idx
            ]

        dpex.barrier(dpex.CLK_LOCAL_MEM_FENCE)

        for idx in range(square_block_side):
            output += (
                X_sliding_window[local_row_idx, idx]
                * Y_sliding_window[idx, local_col_idx]
            )

        dpex.barrier(dpex.CLK_LOCAL_MEM_FENCE)

    if (result_row_idx < X_n_rows) and (result_col_idx < Y_n_cols):
        result[result_row_idx, result_col_idx] = output


def _arange_reshaped(shape, dtype):
    n_items = shape[0] * shape[1]
    return np.arange(n_items, dtype=dtype).reshape(shape)


X = _arange_reshaped((5, 5), dtype)
Y = _arange_reshaped((5, 5), dtype)

print(np.matmul(X, Y))

X = dpt.asarray(X)
Y = dpt.asarray(Y)

device = X.device.sycl_device
result = dpt.zeros((5, 5), dtype, device=device)

matmul[(6,6), (2,2)](X, Y, result)

print(result)

Output"

# expected output
[[ 150.  160.  170.  180.  190.]
 [ 400.  435.  470.  505.  540.]
 [ 650.  710.  770.  830.  890.]
 [ 900.  985. 1070. 1155. 1240.]
 [1150. 1260. 1370. 1480. 1590.]]

# kernel output
[[ 150.  160.  170.  180.  190.]
 [ 400.  435.  470.  505.  540.]
 [ 650.  710.  770.  830.  890.]
 [ 900.  985. 1070. 1155. 1240.]
 [ 700.  766.  832.  898.  964.]]

I've tried many variations of it with no success, the last row of the output always has wrong values. Note that this seems to be deterministic: the values in the last row are always the same. Maybe there's an error in my snippet but I've questioned each row of it already and tried to inspect each of it, enough to start thinking the compiled code might be wrong instead even if the kernel is written right. WDYT ?

Originally posted by @fcharras in #871 (comment)

[diptorupd]

@fcharras Here is some initial analysis: If I change the matrix to 6x6 so that none of the work items perform noops then the result is correct. If I add a print between the two barriers the result is correct even for the 5x5 matrix multiplication.

Based on my preliminary analysis, the problem seems to be a synchronization problem caused when some of the work items are doing noops. I will next look at the generated LLVM and SPIR-V code to further analyze.

[fcharras]

TY for looking at this ! I also noticed that the issue is hard to grasp, changing minor things in the example can led to a different behavior, or even no bug at all, and I couldn't notice a pattern. I hope you can find better information with the LLVM amd SPIR-V code 🤞

[chudur-budur]

@fcharras I was trying to write a serial equivalent of the matmul code, but my results are different.

import numpy as np

square_block_side = 2
work_group_size = (square_block_side, square_block_side)
dtype = np.float32

result_idx = np.array([[0, 0],[1, 0],[0, 1],[1, 1],[0, 2],[1, 2],
                       [0, 3],[1, 3],[0, 4],[1, 4],[0, 5],[1, 5],
                       [4, 2],[5, 2],[4, 3],[5, 3],[4, 0],[5, 0],
                       [4, 1],[5, 1],[4, 4],[5, 4],[4, 5],[5, 5],
                       [2, 0],[3, 0],[2, 1],[3, 1],[2, 2],[3, 2],
                       [2, 3],[3, 3],[2, 4],[3, 4],[2, 5],[3, 5]])

local_idx = np.array([[0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1],
                      [0, 0],[1, 0],[0, 1],[1, 1]])

def matmul_no_kernel(X, Y):
    X_n_rows = X.shape[0]
    Y_n_cols = Y.shape[1]
    n_cols = X.shape[1]

    n_blocks_for_cols = n_cols // square_block_side
    if (n_cols % square_block_side) > 0:
        n_blocks_for_cols += 1

    X_sliding_window = np.zeros(shape=work_group_size, dtype=dtype)     
    Y_sliding_window = np.zeros(shape=work_group_size, dtype=dtype) 
    
    result = np.zeros((X_n_rows,Y_n_cols), dtype=dtype)
    
    for i in range(result_idx.shape[0]):
        local_row_idx, local_col_idx = local_idx[i]
        result_row_idx, result_col_idx = result_idx[i]
        
        output = 0.0 
        for block_idx in range(n_blocks_for_cols):
            if (result_row_idx < X_n_rows) and (
                (local_col_idx + (square_block_side * block_idx)) < n_cols
            ):
                X_sliding_window[local_row_idx, local_col_idx] = X[
                    result_row_idx, local_col_idx + (square_block_side * block_idx)
                ]

            if (result_col_idx < Y_n_cols) and (
                (local_row_idx + (square_block_side * block_idx)) < n_cols
            ):
                Y_sliding_window[local_row_idx, local_col_idx] = Y[
                    local_row_idx + (square_block_side * block_idx), result_col_idx
                ]

            for idx in range(square_block_side):
                output += (
                    X_sliding_window[local_row_idx, idx]
                    * Y_sliding_window[idx, local_col_idx]
                )

        if (result_row_idx < X_n_rows) and (result_col_idx < Y_n_cols):
            result[result_row_idx, result_col_idx] = output
    return result    

def _arange_reshaped(shape, dtype):
    n_items = shape[0] * shape[1]
    return np.arange(n_items, dtype=dtype).reshape(shape)

if __name__ == "__main__":
    X = _arange_reshaped((5, 5), dtype)
    Y = _arange_reshaped((5, 5), dtype)

    print(np.matmul(X, Y))

    print(result_idx.shape)
    print(local_idx.shape)
    print(matmul_no_kernel(X, Y))

But I am getting a different result.

[[ 253.  258.  112.  156.  259.]
 [ 700.  859.  462.  957.  880.]
 [1711. 1128. 1057. 1138. 1207.]
 [1300. 1869. 1860. 2087. 2070.]
 [ 871. 1864. 1867. 2008. 2137.]]

Is my code correct?

[fcharras]

Fully serialized, this can't work @chudur-budur , it really requires the barrier steps for the synchronization in work items of the same work group to work. So the fix is not trivial here.

[fcharras]

See #906 which I think is a simpler instance of this bug.

[diptorupd]

@fcharras I am finally able to devote time to the issue. My initial intuition after going over the generated LLVM IR is that with -O3 optimization flag LLVM is applying some optimizations that are leading to incorrect indexing when looking up values from the local array. It seems to be happening because LLVM does not really know the semantics of the get_global_id and family of SPIR-V calls that we generate.

To test my hypothesis, I ran the reproducer with NUMBA_OPT=0 environment variable. With all optimizations turned off, at least this reproducer works for me on all supported devices (opencl:cpu, level_zero:gpu, opencl:gpu).

My next steps are:

a) I will extract the optimization flag that dpcpp uses and use just those when we compiler a kernel.
b) Replace the SPIR-V instructions we use with those dpcpp currently generates. (I have already done it for barrier and tested unsuccessfully).

Will update with latest findings soon.

[fcharras]

Thank you! I will experiment with NUMBA_OPT and report the findings too.