2D vector sum using usm_ndarrays does not work in kernels

[diptorupd]

Two dimensional indexing over dpctl.tensor.usm_ndarray data types inside a kernel does not work. Incorrect result is generated.

Reproducer:

import dpctl
import dpctl.tensor as dpt
import numba_dpex as dpex

@dpex.kernel
def data_parallel_sum(a, b, c):
    """
    A two-dimensional vector addition example using the ``kernel`` decorator.
    """
    i = dpex.get_global_id(0)
    j = dpex.get_global_id(1)
    c[i, j] = a[i, j] + b[i, j]

# Array dimensions
X = 8
Y = 8
global_size = X, Y

device = dpctl.select_default_device()
a_dpt = dpt.arange(X * Y, dtype=dpt.float32, device=device)
a_dpt = dpt.reshape(a_dpt, (X, Y))
b_dpt = dpt.arange(X * Y, dtype=dpt.float32, device=device)
b_dpt = dpt.reshape(b_dpt, (X, Y))
c_dpt = dpt.empty_like(a_dpt)
c_dpt = dpt.reshape(c_dpt, (X, Y))

data_parallel_sum[global_size](a, b, c)

[diptorupd]

Fixed. The issue was in the example and the reproducer. The output of the reshape was not assigned to the array variable. Refer the updated vector_sum2D.py.

[diptorupd]

A cleaned up LLVM IR for the kernel is as follows. As can be seen the indexing is being calculated incorrectly.

For shape (X,Y),

global id (0,0) : 0 (OK serendipitous)
global id (0,1) : Y (should be 1)
global id (0,2) : 2+2Y (should be 2)

; ModuleID = '/home/diptorupd/Desktop/vector_sum2d.ll'
source_filename = "/home/diptorupd/Desktop/vector_sum2d.ll"
target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
target triple = "spir64-unknown-unknown"

@_data_parallel_sum_ = common local_unnamed_addr global i8* null

; Function Attrs: alwaysinline
define spir_func i32 @_data_parallel_sum_(
    i8** nocapture writeonly %.ret, 
    float addrspace(1)* nocapture readnone %arg.a.0, 
    float addrspace(1)* nocapture readnone %arg.a.1, 
    i64 %arg.a.2, 
    i64 %arg.a.3, 
    float addrspace(1)* nocapture readonly %arg.a.4, 
    i64 %arg.a.5.0, 
    i64 %arg.a.5.1, 
    i64 %arg.a.6.0, 
    i64 %arg.a.6.1,
    float addrspace(1)* nocapture readnone %arg.b.0,
    float addrspace(1)* nocapture readnone %arg.b.1,
    i64 %arg.b.2,
    i64 %arg.b.3,
    float addrspace(1)* nocapture readonly %arg.b.4,
    i64 %arg.b.5.0,
    i64 %arg.b.5.1,
    i64 %arg.b.6.0, 
    i64 %arg.b.6.1,
    float addrspace(1)* nocapture readnone %arg.c.0,
    float addrspace(1)* nocapture readnone %arg.c.1, 
    i64 %arg.c.2, 
    i64 %arg.c.3, 
    float addrspace(1)* nocapture writeonly %arg.c.4,
    i64 %arg.c.5.0,
    i64 %arg.c.5.1,
    i64 %arg.c.6.0,
    i64 %arg.c.6.1
) local_unnamed_addr #0 {
entry:
  %.62 = tail call spir_func i64 @_Z13get_global_idj(i32 0)
  %.84 = tail call spir_func i64 @_Z13get_global_idj(i32 1)
  %.120 = icmp slt i64 %.62, 0
  %.121 = select i1 %.120, i64 %arg.a.5.0, i64 0
  %.122 = add i64 %.121, %.62
  %.123 = icmp slt i64 %.84, 0
  %.124 = select i1 %.123, i64 %arg.a.5.1, i64 0
  %.125 = add i64 %.124, %.84
  %.142 = mul i64 %.122, %arg.a.5.1
  %.145 = add i64 %.125, %.142
  %.146 = getelementptr float, float addrspace(1)* %arg.a.4, i64 %.145
  %.147 = load float, float addrspace(1)* %.146, align 4
  %.184 = select i1 %.120, i64 %arg.b.5.0, i64 0
  %.185 = add i64 %.184, %.62
  %.187 = select i1 %.123, i64 %arg.b.5.1, i64 0
  %.188 = add i64 %.187, %.84
  %.205 = mul i64 %.185, %arg.b.5.1
  %.208 = add i64 %.188, %.205
  %.209 = getelementptr float, float addrspace(1)* %arg.b.4, i64 %.208
  %.210 = load float, float addrspace(1)* %.209, align 4
  %.220 = fadd float %.147, %.210
  %.262 = select i1 %.120, i64 %arg.c.5.0, i64 0
  %.263 = add i64 %.262, %.62
  %.265 = select i1 %.123, i64 %arg.c.5.1, i64 0
  %.266 = add i64 %.265, %.84
  %.283 = mul i64 %.263, %arg.c.5.1
  %.286 = add i64 %.266, %.283
  %.287 = getelementptr float, float addrspace(1)* %arg.c.4, i64 %.286
  store float %.220, float addrspace(1)* %.287, align 4
  store i8* null, i8** %.ret, align 8
  ret i32 0
}

declare spir_func i64 @_Z13get_global_idj(i32) local_unnamed_addr

attributes #0 = { alwaysinline }

[diptorupd]

The above points to some issue with the overload for __getitem__ that we are invoking.

[fcharras]

Do you think those issues with indexing seen in the LLVM IR snippet are linked to the more recent PRs or could it affect versions <=0.19.0 ? I'm running into some unexpected behavior of simpler kernels with ==0.19.0 that makes me think of this issue (indexing in the binaries seems done wrong, despite code being written right).

[fcharras]

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 ?

[oleksandr-pavlyk]

@fcharras The equivalent code written in DPC++ produces the expected output:

#include <CL/sycl.hpp>
#include <vector>
#include <iostream>

template <typename T>
sycl::event
gemm(
  sycl::queue q,
  const T* X_ptr, const T* Y_ptr, T* R_ptr,
  size_t X_n_rows,
  size_t X_n_cols,
  size_t Y_n_cols,
  size_t gws0_blocks, size_t gws1_blocks, size_t lws0, size_t lws1,
  size_t square_block_size,
  const std::vector<sycl::event> &depends={}
)
{
  sycl::event comp_ev =
    q.submit([&](sycl::handler &cgh) {
        cgh.depends_on(depends);

        sycl::local_accessor<T, 2> X_sliding_window(
             sycl::range<2>(lws0, lws1), cgh);
        sycl::local_accessor<T, 2> Y_sliding_window(
             sycl::range<2>(lws0, lws1), cgh);

        auto gwsRange = sycl::range<2>(
             gws0_blocks*lws0, gws1_blocks*lws1);
        auto lwsRange = sycl::range<2>(
             lws0, lws1);
        size_t n_cols = X_n_cols;
        cgh.parallel_for(
           sycl::nd_range<2>(gwsRange, lwsRange),
           [=](sycl::nd_item<2> ndit) {
             size_t result_row_idx = ndit.get_global_id(0);
             size_t result_col_idx = ndit.get_global_id(1);
             size_t local_row_idx = ndit.get_local_id(0);
             size_t local_col_idx = ndit.get_local_id(1);

             size_t n_blocks_for_cols = (n_cols + square_block_size - 1) / square_block_size;

             T output(0);
             for(size_t block_idx = 0; block_idx < n_blocks_for_cols; ++block_idx) {
               size_t col_idx = local_col_idx + square_block_size * block_idx;
               T X_v = (result_row_idx < X_n_rows && col_idx < n_cols) ? X_ptr[result_row_idx * X_n_cols + col_idx] : T(0);
               X_sliding_window[ndit.get_local_id()] = X_v;

               size_t row_idx = local_row_idx + square_block_size * block_idx;
               T Y_v = (result_col_idx < Y_n_cols && row_idx < n_cols) ? Y_ptr[row_idx * Y_n_cols + result_col_idx] : T(0);

               Y_sliding_window[ndit.get_local_id()] = Y_v;

               ndit.barrier(sycl::access::fence_space::local_space);

               for(size_t idx = 0; idx < square_block_size; ++idx) {
                 output += X_sliding_window[sycl::id<2>(local_row_idx, idx)] * Y_sliding_window[sycl::id<2>(idx, local_col_idx)];
               }

               ndit.barrier(sycl::access::fence_space::local_space);

             }
             if (result_row_idx < X_n_rows && result_col_idx < Y_n_cols) {
               R_ptr[result_row_idx * Y_n_cols + result_col_idx] = output;
             }
           }
        );
    });

  return comp_ev;
}

int main(void) {
  using T = float;
  size_t n = 5;

  size_t square_block_side = 2;

  sycl::queue q{ sycl::default_selector_v };

  size_t n_elems = n * n;
  T *X_ptr = sycl::malloc_device<T>(n_elems, q);
  T *Y_ptr = sycl::malloc_device<T>(n_elems, q);
  T *R_ptr = sycl::malloc_device<T>(n_elems, q);

  sycl::event pop_ev =
    q.submit([&](sycl::handler &cgh) {
       cgh.parallel_for({n_elems}, [=](sycl::id<1> id) { T v(id[0]); X_ptr[id] = v; Y_ptr[id] = v; R_ptr[id] = T(-7); });
    });

  size_t lws0 = 2;
  size_t lws1 = 2;
  size_t gws0_blocks = (n + lws0 - 1)/ lws0;
  size_t gws1_blocks = (n + lws1 - 1)/ lws1;

  sycl::event comp_ev =
    gemm(q, X_ptr, Y_ptr, R_ptr, n, n, n, gws0_blocks, gws1_blocks, lws0, lws1, lws0, {pop_ev});

  T* host_res = new T[n_elems];

  sycl::event copy_ev = q.copy<T>(R_ptr, host_res, n_elems, {comp_ev});

  copy_ev.wait();



  for(size_t i0 = 0; i0 < n; ++i0) {
    std::cout << "[ ";
    for(size_t i1 = 0; i1 < n; ++i1) {
      std::cout << host_res[ i0 * n + i1] << " ";
    }
    std::cout << "]" << std::endl;
  }

  delete[] host_res;

  sycl::free(X_ptr, q);
  sycl::free(Y_ptr, q);
  sycl::free(R_ptr, q);

  return 0;
}

Compile it using icpx -fsycl gemm.cpp -o gemm.exe.

Output:

opavlyk@opavlyk-mobl:~/repos/gemm$ ./gemm.exe
[ 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 ]

[fcharras]

Thank you for the snippet ! the DPC++ code produces the expected output in my environment too.

Would you say that it confirms that it's a bug in numba_dpex ?

What confuses me is that our KMeans have more complicated kernels, yet this problem has never occured before. I can't tell what makes this kernel stand out from the other. After further testing I have found that:

• the issue persists if using a 1D grid of work items rather than 2D grid (here that would be gs=36, ls=4 rather than gs=(6,6), ls=(2,2) + adapting the kernel to map 1D ids to 2D indexing)

• it also persists when using a single local memory array that is twice the size