From d2a35374714fbe9adb7aa34d9d3fd46bd8786393 Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Fri, 10 Jul 2026 18:01:57 -0700 Subject: [PATCH] Unify row_mul_mat! and row_mul_vec! into single methods Replace the 8 row_mul_mat! methods and 4 row_mul_vec! methods in ext/cuda/column_matrix_helpers.jl with one method each, driven by small shape-trait helpers (column_slot, n_column_slots, product_shape) derived from the AbstractMatrixShape types. Product-row entries outside the product matrix are zeroed, matching multiply_matrix_at_index on the CPU. Verified on GPU (A100): all matrix_fields_broadcasting tests pass, and kernel timings/PTX are unchanged from the multi-method version. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_011oirTEPPBEj5BWwhgYKMzn --- ext/cuda/column_matrix_helpers.jl | 441 +++++------------------------- 1 file changed, 65 insertions(+), 376 deletions(-) diff --git a/ext/cuda/column_matrix_helpers.jl b/ext/cuda/column_matrix_helpers.jl index f832769a08..cf6f815ead 100644 --- a/ext/cuda/column_matrix_helpers.jl +++ b/ext/cuda/column_matrix_helpers.jl @@ -1,250 +1,73 @@ -# row_mul_mat! handles banded matrix * banded matrix. There are 8 methods, but they all have the -# same structure, so we they could be written as a single method. -# The others can be obtained by copy-pasting and changing the indices appropriately. -# Note that these are all specialized for CUDA.blockDim().x faces , so the indices are hardcoded. -Base.@propagate_inbounds function row_mul_mat!( - ::Type{P}, - mat1_row, - matrix2, - ::FaceToCenter, - ::CenterToFace, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) - li = 1i32 - ri = CUDA.blockDim().x - 1i32 - zero_entry = zero(eltype(prod_eltype)) - prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd < li || v + pd > ri - zero_entry - else - UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d - if ld2 <= pd - mat1_row_d <= ud2 && - (0i32 < v + mat1_row_d + half <= CUDA.blockDim().x) - @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] - else - zero_entry - end - end - end - end - return BandMatrixRow{pd1}(prod_entries...) -end - -Base.@propagate_inbounds function row_mul_mat!( - ::Type{P}, - mat1_row, - matrix2, - ::CenterToFace, - ::FaceToCenter, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) - li = 1i32 - ri = CUDA.blockDim().x - zero_entry = zero(eltype(prod_eltype)) - prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd < li || v + pd > ri - zero_entry - else - UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d - if ld2 <= pd - mat1_row_d <= ud2 && - (0i32 < v + mat1_row_d - half < CUDA.blockDim().x) - @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d - half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] - else - zero_entry - end - end - end - end - return BandMatrixRow{pd1}(prod_entries...) -end - -Base.@propagate_inbounds function row_mul_mat!( - ::Type{P}, - mat1_row, - matrix2, - ::CenterToCenter, - ::CenterToCenter, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) - li = 1i32 - ri = CUDA.blockDim().x - 1i32 - zero_entry = zero(eltype(prod_eltype)) - prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd < li || v + pd > ri - zero_entry - else - UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d - if ld2 <= pd - mat1_row_d <= ud2 && - (0i32 < v + mat1_row_d <= CUDA.blockDim().x - 1i32) - @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] - else - zero_entry - end - end - end - end - return BandMatrixRow{pd1}(prod_entries...) -end +# Helpers for multiplying a single row of a banded matrix with a banded matrix or vector +# stored in shared memory. These are specialized for columns with CUDA.blockDim().x face +# levels, computed by threads `v = 1:CUDA.blockDim().x` (one column per `threadIdx().y`). +# +# Shared memory holds one value per thread: values on cell centers are stored at slot +# `v = center index` (so the last slot is unused), and values on cell faces at slot +# `v = face index + half`. -Base.@propagate_inbounds function row_mul_mat!( - ::Type{P}, - mat1_row, - matrix2, - ::FaceToFace, - ::FaceToFace, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) +# `column_slot(v, d, shape)` is the shared-memory slot of the value that entry `d` of the +# matrix row computed by thread `v` multiplies, for a matrix with the given `shape`: the row +# index of thread `v` is `v` for center rows and `v - half` for face rows, the column index +# is `row index + d`, and the slot of a column index is `index + half` for face columns and +# `index` for center columns. +Base.@propagate_inbounds column_slot(v, d, ::Union{CenterToCenter, FaceToFace}) = v + d +Base.@propagate_inbounds column_slot(v, d, ::FaceToCenter) = v + d + half +Base.@propagate_inbounds column_slot(v, d, ::CenterToFace) = v + d - half - li = 1i32 - ri = CUDA.blockDim().x +# Number of valid slots in the column space of a matrix with the given shape. +@inline n_column_slots(::Union{FaceToCenter, FaceToFace}) = CUDA.blockDim().x +@inline n_column_slots(::Union{CenterToFace, CenterToCenter}) = + CUDA.blockDim().x - 1i32 - zero_entry = zero(eltype(prod_eltype)) - prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd < li || v + pd > ri - zero_entry - else - UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d - if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d <= CUDA.blockDim().x) - @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] - else - zero_entry - end - end - end - end - return BandMatrixRow{pd1}(prod_entries...) -end - -Base.@propagate_inbounds function row_mul_mat!( - ::Type{P}, - mat1_row, - matrix2, - ::FaceToCenter, - ::FaceToFace, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) - li = 1i32 - ri = CUDA.blockDim().x - zero_entry = zero(eltype(prod_eltype)) - prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd + half < li || v + pd + half > ri - zero_entry - else - UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d - if ld2 <= pd - mat1_row_d <= ud2 && - (0i32 < v + mat1_row_d + half <= CUDA.blockDim().x) - @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] - else - zero_entry - end - end - end - end - return BandMatrixRow{pd1}(prod_entries...) -end +# Shape of `matrix1 * matrix2`: its rows match `matrix1`'s rows and its columns match +# `matrix2`'s columns. +@inline product_shape( + ::Union{FaceToCenter, CenterToCenter}, + ::Union{CenterToFace, CenterToCenter}, +) = CenterToCenter() +@inline product_shape( + ::Union{FaceToCenter, CenterToCenter}, + ::Union{FaceToCenter, FaceToFace}, +) = FaceToCenter() +@inline product_shape( + ::Union{CenterToFace, FaceToFace}, + ::Union{CenterToFace, CenterToCenter}, +) = CenterToFace() +@inline product_shape( + ::Union{CenterToFace, FaceToFace}, + ::Union{FaceToCenter, FaceToFace}, +) = FaceToFace() +# row_mul_mat! handles banded matrix * banded matrix. Entries of the product row whose +# column index lies outside the product matrix are set to zero, matching +# `multiply_matrix_at_index` in `src/MatrixFields/matrix_multiplication.jl`. Base.@propagate_inbounds function row_mul_mat!( ::Type{P}, mat1_row, matrix2, - ::CenterToFace, - ::CenterToCenter, + shape1::MatrixFields.AbstractMatrixShape, + shape2::MatrixFields.AbstractMatrixShape, ) where {P} - prod_eltype = P v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) - li = 1i32 - ri = CUDA.blockDim().x - zero_entry = zero(eltype(prod_eltype)) + block_col_idx = threadIdx().y + ld1, ud1 = MatrixFields.outer_diagonals(typeof(mat1_row)) + ld2, ud2 = MatrixFields.outer_diagonals(eltype(matrix2)) + pd1, pd2 = MatrixFields.outer_diagonals(P) + prod_shape = product_shape(shape1, shape2) + mat2_offset = (block_col_idx - 1i32) * CUDA.blockDim().x + zero_entry = zero(eltype(P)) prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd + half < li || v + pd + half > ri - zero_entry + prod_slot = column_slot(v, pd, prod_shape) + if prod_slot < 1i32 || prod_slot > n_column_slots(prod_shape) + zero_entry # This entry is outside the product matrix. else UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d + mat2_slot = column_slot(v, mat1_row_d, shape1) if ld2 <= pd - mat1_row_d <= ud2 && - (0i32 < v + mat1_row_d - half < CUDA.blockDim().x) - @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d - half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] - else - zero_entry - end - end - end - end - return BandMatrixRow{pd1}(prod_entries...) -end - -Base.@propagate_inbounds function row_mul_mat!( - ::Type{P}, - mat1_row, - matrix2, - ::FaceToFace, - ::CenterToFace, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) - li = 1i32 - ri = CUDA.blockDim().x - zero_entry = zero(eltype(prod_eltype)) - prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd + half < li || v + pd + half > ri - zero_entry - else - UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d - if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d <= CUDA.blockDim().x) + 0i32 < mat2_slot <= n_column_slots(shape1) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[mat2_slot + mat2_offset][pd - mat1_row_d] else zero_entry end @@ -254,162 +77,28 @@ Base.@propagate_inbounds function row_mul_mat!( return BandMatrixRow{pd1}(prod_entries...) end -Base.@propagate_inbounds function row_mul_mat!( - ::Type{P}, - mat1_row, - matrix2, - ::CenterToCenter, - ::FaceToCenter, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - ld2, ud2 = MatrixFields.outer_diagonals(mat2_eltype) - pd1, pd2 = MatrixFields.outer_diagonals(prod_eltype) - li = 1i32 - ri = CUDA.blockDim().x - zero_entry = zero(eltype(prod_eltype)) - prod_entries = UnrolledUtilities.unrolled_map((pd1:pd2...,)) do pd - if v + pd + half < li || v + pd + half > ri - zero_entry - else - UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d - if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d < CUDA.blockDim().x) - @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] - else - zero_entry - end - end - end - end - return BandMatrixRow{pd1}(prod_entries...) -end - -# row_mul_vec! handles banded matrix * vector. There are four methods, but they all have the -# same structure, so we they could be written as a single method. -# The others can be obtained by copy-pasting and changing the indices appropriately. -# Note that these are all specialized for CUDA.blockDim().x faces , so the indices are hardcoded. -Base.@propagate_inbounds function row_mul_vec!( - ::Type{P}, - mat1_row, - matrix2, - ::FaceToCenter, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - li = 1i32 - ri = CUDA.blockDim().x - 1i32 - zero_entry = zero(prod_eltype) - return UnrolledUtilities.unrolled_mapreduce( - +, - ld1:ud1; - init = zero_entry, - ) do mat1_row_d - if (0i32 < v + mat1_row_d + half <= CUDA.blockDim().x) - @inbounds outer_or_mul( - mat1_row[mat1_row_d], - matrix2[v + mat1_row_d + half + (i - 1i32) * CUDA.blockDim().x], - ) - else - zero_entry - end - end -end - +# row_mul_vec! handles banded matrix * vector. Base.@propagate_inbounds function row_mul_vec!( ::Type{P}, mat1_row, - matrix2, - ::CenterToFace, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - li = 1i32 - ri = CUDA.blockDim().x - zero_entry = zero(prod_eltype) - return UnrolledUtilities.unrolled_mapreduce( - +, - ld1:ud1; - init = zero_entry, - ) do mat1_row_d - if (0i32 < v + mat1_row_d - half < CUDA.blockDim().x) - @inbounds outer_or_mul( - mat1_row[mat1_row_d], - matrix2[v + mat1_row_d - half + (i - 1i32) * CUDA.blockDim().x], - ) - else - zero_entry - end - end -end - -Base.@propagate_inbounds function row_mul_vec!( - ::Type{P}, - mat1_row, - matrix2, - ::CenterToCenter, -) where {P} - prod_eltype = P - v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - li = 1i32 - ri = CUDA.blockDim().x - 1i32 - zero_entry = zero(prod_eltype) - return UnrolledUtilities.unrolled_mapreduce( - +, - ld1:ud1; - init = zero_entry, - ) do mat1_row_d - if (0i32 < v + mat1_row_d <= CUDA.blockDim().x - 1i32) - @inbounds outer_or_mul( - mat1_row[mat1_row_d], - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x], - ) - else - zero_entry - end - end -end - -Base.@propagate_inbounds function row_mul_vec!( - ::Type{P}, - mat1_row, - matrix2, - ::FaceToFace, + vector2, + shape1::MatrixFields.AbstractMatrixShape, ) where {P} - prod_eltype = P v = threadIdx().x - i = threadIdx().y - mat1_eltype = typeof(mat1_row) - mat2_eltype = eltype(matrix2) - ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) - li = 1i32 - ri = CUDA.blockDim().x - zero_entry = zero(prod_eltype) + block_col_idx = threadIdx().y + ld1, ud1 = MatrixFields.outer_diagonals(typeof(mat1_row)) + vec2_offset = (block_col_idx - 1i32) * CUDA.blockDim().x + zero_entry = zero(P) return UnrolledUtilities.unrolled_mapreduce( +, ld1:ud1; init = zero_entry, ) do mat1_row_d - if (0i32 < v + mat1_row_d <= CUDA.blockDim().x) + vec2_slot = column_slot(v, mat1_row_d, shape1) + if 0i32 < vec2_slot <= n_column_slots(shape1) @inbounds outer_or_mul( mat1_row[mat1_row_d], - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x], + vector2[vec2_slot + vec2_offset], ) else zero_entry