diff --git a/src/level1/dot_device.mojo b/src/level1/dot_device.mojo index 87e8d5f..5e8bf5b 100644 --- a/src/level1/dot_device.mojo +++ b/src/level1/dot_device.mojo @@ -6,9 +6,8 @@ from math import ceildiv comptime TBsize = 512 -# level1.dot -# computes the dot product of two vectors -fn dot_device[ +# Pass 1: each block finds its partial sum +fn dot_device_partial[ BLOCK: Int, dtype: DType ]( @@ -17,16 +16,15 @@ fn dot_device[ incx: Int, y: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], incy: Int, - output: UnsafePointer[Scalar[dtype], MutAnyOrigin], + partial_results: UnsafePointer[Scalar[dtype], MutAnyOrigin], ): - var global_i = block_dim.x * block_idx.x + thread_idx.x var n_threads = grid_dim.x * block_dim.x var local_i = thread_idx.x shared_res = stack_allocation[ BLOCK, - Float64, + Scalar[dtype], address_space = AddressSpace.SHARED ]() @@ -34,7 +32,7 @@ fn dot_device[ for i in range(global_i, n, n_threads): thread_sum += x[i * incx] * y[i * incy] - shared_res[local_i] = Float64(thread_sum) + shared_res[local_i] = thread_sum barrier() var stride = BLOCK // 2 @@ -45,9 +43,44 @@ fn dot_device[ stride //= 2 if local_i == 0: - _ = Atomic[dtype].fetch_add(output, Scalar[dtype](shared_res[0])) + partial_results[block_idx.x] = shared_res[0] + +# Pass 2: Final reduction +fn dot_device_reduce[ + BLOCK: Int, + dtype: DType +]( + n_blocks: Int, + partial_results: UnsafePointer[Scalar[dtype], MutAnyOrigin], + d_out: UnsafePointer[Scalar[dtype], MutAnyOrigin], +): + shared_res = stack_allocation[ + BLOCK, + Scalar[dtype], + address_space = AddressSpace.SHARED + ]() + + var local_i = thread_idx.x + + if local_i < n_blocks: + shared_res[local_i] = partial_results[local_i] + else: + shared_res[local_i] = 0 + barrier() + + var stride = BLOCK // 2 + while stride > 0: + if local_i < stride: + shared_res[local_i] += shared_res[local_i + stride] + barrier() + stride //= 2 + if local_i == 0: + d_out[0] = shared_res[0] + +# level1.dot +# computes the dot product of two vectors fn blas_dot[dtype: DType]( n: Int, d_x: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], @@ -57,16 +90,29 @@ fn blas_dot[dtype: DType]( d_out: UnsafePointer[Scalar[dtype], MutAnyOrigin], ctx: DeviceContext ) raises: - + blas_error_if["blas_dot", "n < 0"](n < 0) blas_error_if["blas_copy", "incx == 0"](incx == 0) blas_error_if["blas_copy", "incy == 0"](incy == 0) - comptime kernel = dot_device[TBsize, dtype] - ctx.enqueue_function[kernel, kernel]( + # Limit number of blocks to TBsize max for second pass + var n_blocks = min(ceildiv(n, TBsize), TBsize) + var partial_results = ctx.enqueue_create_buffer[dtype](n_blocks) + + comptime kernel1 = dot_device_partial[TBsize, dtype] + ctx.enqueue_function[kernel1, kernel1]( n, d_x, incx, - d_y, incy, d_out, - grid_dim=ceildiv(n, TBsize), + d_y, incy, partial_results.unsafe_ptr(), + grid_dim=n_blocks, block_dim=TBsize, ) + + comptime kernel2 = dot_device_reduce[TBsize, dtype] + ctx.enqueue_function[kernel2, kernel2]( + n_blocks, + partial_results.unsafe_ptr(), + d_out, + grid_dim=1, + block_dim=TBsize + ) ctx.synchronize() diff --git a/src/level1/iamax_device.mojo b/src/level1/iamax_device.mojo index aba2f53..5dba4d2 100644 --- a/src/level1/iamax_device.mojo +++ b/src/level1/iamax_device.mojo @@ -1,32 +1,27 @@ from memory import stack_allocation from gpu.memory import AddressSpace -from gpu import thread_idx, block_dim, block_idx, barrier +from gpu import thread_idx, block_dim, block_idx, grid_dim, barrier from os.atomic import Atomic from gpu.host import DeviceContext from math import ceildiv comptime TBsize = 512 -# level1.iamax -# finds the index of the first element having maximum absolute value -fn iamax_device[ +# Pass 1: each block finds its local max_val and max_idx +fn iamax_device_partial[ BLOCK: Int, dtype: DType ]( n: Int, sx: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], incx: Int, - result: UnsafePointer[Scalar[DType.int64], MutAnyOrigin] + partial_vals: UnsafePointer[Scalar[dtype], MutAnyOrigin], + partial_idxs: UnsafePointer[Int64, MutAnyOrigin] ): - result[0] = -1 - if n < 1 or incx <= 0: - return - - result[0] = 0 + # Quick return if possible if n == 1: return - # Shared memory for indices and values shared_indices = stack_allocation[ BLOCK, Int, @@ -39,8 +34,8 @@ fn iamax_device[ ]() var local_tid = thread_idx.x - var global_tid = block_idx.x * block_dim.x + thread_idx.x - var n_threads = grid_dim.x * block_dim.x + var global_tid = block_idx.x * BLOCK + local_tid + var n_threads = grid_dim.x * BLOCK # Each thread finds its local max var local_max_id = -1 @@ -60,28 +55,84 @@ fn iamax_device[ barrier() # Parallel reduction to find max within the block - var stride = block_dim.x // 2 + var stride = BLOCK // 2 while stride > 0: if local_tid < stride: - var other_idx = local_tid + stride - if other_idx < BLOCK and shared_indices[other_idx] >= 0: - if shared_values[other_idx] > shared_values[local_tid]: - shared_values[local_tid] = shared_values[other_idx] - shared_indices[local_tid] = shared_indices[other_idx] - # Resolve ties by selecting smaller index - elif shared_values[other_idx] == shared_values[local_tid] and - shared_indices[other_idx] < shared_indices[local_tid]: - shared_indices[local_tid] = shared_indices[other_idx] - + var other_idx = shared_indices[local_tid + stride] + var other_val = shared_values[local_tid + stride] + var my_idx = shared_indices[local_tid] + var my_val = shared_values[local_tid] + + if other_val > my_val or (other_val == my_val and other_idx != -1 + and (my_idx == -1 or other_idx < my_idx)): + shared_indices[local_tid] = other_idx + shared_values[local_tid] = other_val barrier() stride //= 2 - # Thread 0 atomically updates the global result - # TODO: complete this to support more than one thread block if local_tid == 0: - result[0] = shared_indices[0] + partial_vals[block_idx.x] = shared_values[0] + partial_idxs[block_idx.x] = shared_indices[0] +# Pass 2: Final reduction +fn iamax_device_reduce[ + BLOCK: Int, + dtype: DType +]( + n: Int, + n_blocks: Int, + partial_vals: UnsafePointer[Scalar[dtype], MutAnyOrigin], + partial_idxs: UnsafePointer[Scalar[DType.int64], MutAnyOrigin], + d_res: UnsafePointer[Scalar[DType.int64], MutAnyOrigin] +): + # Quick return if possible + d_res[0] = 0 + if n == 1: + return + + var shared_vals = stack_allocation[ + BLOCK, + Scalar[dtype], + address_space = AddressSpace.SHARED + ]() + var shared_idxs = stack_allocation[ + BLOCK, + Int, + address_space = AddressSpace.SHARED + ]() + + var local_tid = thread_idx.x + + if local_tid < n_blocks: + shared_vals[local_tid] = partial_vals[local_tid] + shared_idxs[local_tid] = Int(partial_idxs[local_tid]) + else: + shared_vals[local_tid] = Scalar[dtype](-1) + shared_idxs[local_tid] = -1 + + barrier() + + var stride = BLOCK // 2 + while stride > 0: + if local_tid < stride: + var other_val = shared_vals[local_tid + stride] + var other_idx = shared_idxs[local_tid + stride] + var my_val = shared_vals[local_tid] + var my_idx = shared_idxs[local_tid] + + if other_val > my_val or (other_val == my_val and other_idx != -1 + and (my_idx == -1 or other_idx < my_idx)): + shared_vals[local_tid] = other_val + shared_idxs[local_tid] = other_idx + barrier() + stride //= 2 + + if local_tid == 0: + d_res[0] = shared_idxs[0] + +# level1.iamax +# finds the index of the first element having maximum absolute value fn blas_iamax[dtype: DType]( n: Int, d_v: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], @@ -89,16 +140,30 @@ fn blas_iamax[dtype: DType]( d_res: UnsafePointer[Scalar[DType.int64], MutAnyOrigin], ctx: DeviceContext ) raises: - blas_error_if["blas_iamax", "n < 0"](n<=0) blas_error_if["blas_iamax", "incx <= 0"](incx <= 0) + # Limit number of blocks to TBsize max for second pass + var n_blocks = min(ceildiv(n, TBsize), TBsize) + + var partial_vals = ctx.enqueue_create_buffer[dtype](n_blocks) + var partial_idxs = ctx.enqueue_create_buffer[DType.int64](n_blocks) - comptime kernel = iamax_device[TBsize, dtype] - ctx.enqueue_function[kernel, kernel]( + comptime kernel1 = iamax_device_partial[TBsize, dtype] + ctx.enqueue_function[kernel1, kernel1]( n, d_v, incx, + partial_vals.unsafe_ptr(), partial_idxs.unsafe_ptr(), + grid_dim=n_blocks, + block_dim=TBsize + ) + + comptime kernel2 = iamax_device_reduce[TBsize, dtype] + ctx.enqueue_function[kernel2, kernel2]( + n, n_blocks, + partial_vals.unsafe_ptr(), partial_idxs.unsafe_ptr(), d_res, - grid_dim=1, # total thread blocks - block_dim=TBsize # threads per block + grid_dim=1, + block_dim=TBsize ) + ctx.synchronize() diff --git a/src/level1/nrm2_device.mojo b/src/level1/nrm2_device.mojo index df9ae23..51b2bb8 100644 --- a/src/level1/nrm2_device.mojo +++ b/src/level1/nrm2_device.mojo @@ -2,54 +2,82 @@ from gpu import thread_idx, block_idx, block_dim, lane_id from gpu.host import DeviceContext, HostBuffer, DeviceBuffer from layout import Layout, LayoutTensor from gpu.primitives.warp import sum as warp_sum, WARP_SIZE -from math import ceildiv,sqrt -from buffer import NDBuffer, DimList -from algorithm import sum -from layout import Layout, LayoutTensor +from math import ceildiv, sqrt from os.atomic import Atomic -fn nrm2_device[ +comptime TBsize = 512 + +# Pass 1: each block finds its partial sum +fn nrm2_device_partial[ BLOCK: Int, dtype: DType ]( n: Int, x: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], incx: Int, - output: UnsafePointer[Scalar[dtype], MutAnyOrigin], + partial_results: UnsafePointer[Scalar[dtype], MutAnyOrigin], ): - global_i = block_dim.x * block_idx.x + thread_idx.x - local_i = thread_idx.x + var global_i = block_dim.x * block_idx.x + thread_idx.x + var n_threads = grid_dim.x * block_dim.x + var local_i = thread_idx.x - shared_scratch = stack_allocation[ + shared_res = stack_allocation[ BLOCK, - Float64, + Scalar[dtype], address_space = AddressSpace.SHARED ]() - # Each thread computes one partial square - var local_sum = Scalar[DType.float64](0) - if global_i < UInt(n): - v = x[global_i * incx] - vf64 = v.cast[DType.float64]() - - local_sum += vf64 * vf64 - - shared_scratch[local_i] = local_sum + var thread_sum = Scalar[dtype](0) + for i in range(global_i, n, n_threads): + var v = x[i * incx] + thread_sum += v * v + shared_res[local_i] = thread_sum barrier() var stride = BLOCK // 2 while stride > 0: if local_i < stride: - shared_scratch[local_i] += shared_scratch[local_i + stride] + shared_res[local_i] += shared_res[local_i + stride] barrier() stride //= 2 - # Lane 0 accumulates into global output if local_i == 0: - real_total = Scalar[dtype](shared_scratch[0]) - _ = Atomic[dtype].fetch_add(output, real_total) + partial_results[block_idx.x] = shared_res[0] +# Pass 2: Final reduction + sqrt +fn nrm2_device_reduce[ + BLOCK: Int, + dtype: DType +]( + n_blocks: Int, + partial_results: UnsafePointer[Scalar[dtype], MutAnyOrigin], + d_out: UnsafePointer[Scalar[dtype], MutAnyOrigin], +): + shared_res = stack_allocation[ + BLOCK, + Scalar[dtype], + address_space = AddressSpace.SHARED + ]() + + var local_i = thread_idx.x + + if local_i < n_blocks: + shared_res[local_i] = partial_results[local_i] + else: + shared_res[local_i] = 0 + + barrier() + + var stride = BLOCK // 2 + while stride > 0: + if local_i < stride: + shared_res[local_i] += shared_res[local_i + stride] + barrier() + stride //= 2 + + if local_i == 0: + d_out[0] = Scalar[dtype](sqrt(Float32(shared_res[0]))) fn blas_nrm2[dtype: DType]( n: Int, @@ -60,12 +88,26 @@ fn blas_nrm2[dtype: DType]( ) raises: blas_error_if["blas_nrm2", "n < 0"](n < 0) blas_error_if["blas_nrm2", "incx <= 0"](incx <= 0) - - comptime kernel = nrm2_device[TBsize, dtype] - ctx.enqueue_function[kernel, kernel]( - n, d_x, incx, d_out, - grid_dim=ceildiv(n, TBsize), + + # Limit number of blocks to TBsize max for second pass + var n_blocks = min(ceildiv(n, TBsize), TBsize) + var partial_results = ctx.enqueue_create_buffer[dtype](n_blocks) + + comptime kernel1 = nrm2_device_partial[TBsize, dtype] + ctx.enqueue_function[kernel1, kernel1]( + n, d_x, incx, + partial_results.unsafe_ptr(), + grid_dim=n_blocks, block_dim=TBsize, ) + + comptime kernel2 = nrm2_device_reduce[TBsize, dtype] + ctx.enqueue_function[kernel2, kernel2]( + n_blocks, + partial_results.unsafe_ptr(), + d_out, + grid_dim=1, + block_dim=TBsize, + ) + ctx.synchronize() - diff --git a/test-level1.mojo b/test-level1.mojo index cc1a683..73d40c4 100644 --- a/test-level1.mojo +++ b/test-level1.mojo @@ -1,5 +1,5 @@ from sys import argv -from testing import assert_equal, assert_almost_equal, TestSuite +from testing import assert_equal, assert_almost_equal, assert_true, TestSuite from gpu.host import DeviceContext from math import sqrt from complex import * @@ -149,7 +149,6 @@ def dot_test[ size: Int ](): with DeviceContext() as ctx: - # print("[ dot test", dtype, "]") out = ctx.enqueue_create_buffer[dtype](1) out.enqueue_fill(0) a_device = ctx.enqueue_create_buffer[dtype](size) @@ -197,11 +196,17 @@ def dot_test[ return sp_res_mojo = Scalar[dtype](py=sp_res) + var norm_a = frobenius_norm[dtype](a.unsafe_ptr(), size) + var norm_b = frobenius_norm[dtype](b.unsafe_ptr(), size) with out.map_to_host() as res_mojo: - # print("out:", res_mojo[0]) - # print("expected:", sp_res) - # may want to use assert_almost_equal with tolerance specified - assert_almost_equal(res_mojo[0], sp_res_mojo, atol=atol) + var error = abs(res_mojo[0] - sp_res_mojo) + var ok = check_gemm_error[dtype]( + 1, 1, size, + Scalar[dtype](1), Scalar[dtype](0), + norm_a, norm_b, Scalar[dtype](0), + error, + ) + assert_true(ok) def dot_test_complex[ @@ -257,11 +262,17 @@ def dot_test_complex[ return sp_res_mojo = Scalar[dtype](py=sp_res) + var norm_a = frobenius_norm[dtype](a.unsafe_ptr(), size) + var norm_b = frobenius_norm[dtype](b.unsafe_ptr(), size) with out.map_to_host() as res_mojo: - # print("out:", res_mojo[0]) - # print("expected:", sp_res) - # may want to use assert_almost_equal with tolerance specified - assert_almost_equal(res_mojo[0], sp_res_mojo, atol=atol) + var error = abs(res_mojo[0] - sp_res_mojo) + var ok = check_gemm_error[dtype]( + 1, 1, size, + Scalar[dtype](1), Scalar[dtype](0), + norm_a, norm_b, Scalar[dtype](0), + error, + ) + assert_true(ok) def dotc_test[ @@ -483,9 +494,6 @@ def nrm2_test[ # Move Mojo result from CPU to GPU and compare to SciPy sp_res_mojo = Scalar[dtype](py=sp_res) with d_res.map_to_host() as res_mojo: - res_mojo[0] = sqrt(res_mojo[0]) - # print("out:", res_mojo[0]) - # print("expected:", sp_res) assert_almost_equal(res_mojo[0], sp_res_mojo) diff --git a/test-level2.mojo b/test-level2.mojo index a7e742b..22f7f5b 100644 --- a/test-level2.mojo +++ b/test-level2.mojo @@ -263,11 +263,11 @@ def spr_test[ if dtype == DType.float32: np_AP = np.array(py_AP, dtype=np.float32) np_x = np.array(py_x, dtype=np.float32) - sp_res = sp_blas.sspr(alpha, np_x, lower=uplo, ap=np_AP, overwrite_ap=False) + sp_res = sp_blas.sspr(n, alpha, np_x, lower=uplo, ap=np_AP, overwrite_ap=False) elif dtype == DType.float64: np_AP = np.array(py_AP, dtype=np.float64) np_x = np.array(py_x, dtype=np.float64) - sp_res = sp_blas.dspr(alpha, np_x, lower=uplo, ap=np_AP, overwrite_ap=False) + sp_res = sp_blas.dspr(n, alpha, np_x, lower=uplo, ap=np_AP, overwrite_ap=False) else: print("Unsupported type: ", dtype) return @@ -971,12 +971,13 @@ def tbsv_test[ for i in range(n): py_b.append(b[i]) - var sp_res: PythonObject - var is_lower = 0 if uplo else 1 var trans_mode = 1 if trans else 0 var unit_diag = True if diag else False + var np_A: PythonObject + var np_b: PythonObject + var sp_res: PythonObject if dtype == DType.float32: np_A = np.array(py_A, dtype=np.float32).reshape(n, n) np_b = np.array(py_b, dtype=np.float32) @@ -999,6 +1000,7 @@ def tbsv_test[ print("Unsupported type: ", dtype) return + # Primary check: will pass for well-conditioned A with x_d.map_to_host() as res_mojo: var norm_diff = Scalar[dtype](0) for i in range(n): @@ -1012,6 +1014,45 @@ def tbsv_test[ norm_A, norm_b, Scalar[dtype](0), norm_diff ) + + # A is ill-conditioned, perform backward error-check. + # Compute x * A and check if it's close enough to b. + if not ok: + # Copy x into temp so we don't overwrite our solution + temp_d = ctx.enqueue_create_buffer[dtype](n) + ctx.enqueue_copy(temp_d, x_d) + ctx.synchronize() + + # Compute temp = A * x with tbmv + work = ctx.enqueue_create_buffer[dtype](n) + blas_tbmv[dtype]( + uplo, trans, diag, + n, k, + A_d.unsafe_ptr(), lda, + temp_d.unsafe_ptr(), 1, + work.unsafe_ptr(), + ctx + ) + + with temp_d.map_to_host() as ax: + # Compute residual r = A*x - b, and norm of x for scaling + var norm_r = Scalar[dtype](0) + var norm_x = Scalar[dtype](0) + for i in range(n): + var r_i = ax[i] - b[i] + norm_r += r_i * r_i + norm_x += res_mojo[i] * res_mojo[i] + norm_r = sqrt(norm_r) + norm_x = sqrt(norm_x) + + # Check residual + ok = check_gemm_error[dtype]( + 1, n, n, + Scalar[dtype](1), Scalar[dtype](0), + norm_A, norm_x, Scalar[dtype](0), + norm_r, + ) + assert_true(ok) def test_gemv():