Unify row_mul_mat! and row_mul_vec! into single methods#2545
Open
imreddyTeja wants to merge 1 commit into
Open
Unify row_mul_mat! and row_mul_vec! into single methods#2545imreddyTeja wants to merge 1 commit into
imreddyTeja wants to merge 1 commit into
Conversation
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oirTEPPBEj5BWwhgYKMzn
imreddyTeja
force-pushed
the
tr/unified-column-helpers-main
branch
from
97ce300 to
d2a3537
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
The commit in this PR was authored by Claude (Claude Code), working on behalf of @imreddyTeja.
Summary
ext/cuda/column_matrix_helpers.jlcontained 8 nearly identicalrow_mul_mat!methods (one per pair of matrix staggering shapes) and 4 nearly identicalrow_mul_vec!methods. This PR replaces them with one method each, shrinking the file from 431 to 119 lines (−376/+65).The per-shape index arithmetic is factored into three small compile-time shape-trait helpers:
column_slot(v, d, shape)— converts entrydof the matrix row computed by threadvinto the shared-memory slot of the value it multiplies (v + d, with±halfcorrections for the mixed-staggering shapes)n_column_slots(shape)— number of valid slots in the matrix's column space (blockDim().xfor face columns, one fewer for center columns)product_shape(shape1, shape2)— the shape ofmatrix1 * matrix2, used for the product-row bounds check inrow_mul_mat!All shapes are singleton types, so everything still resolves at compile time.
One deliberate behavior nuance: in the two CenterToFace-product cases, the old outer bounds check (
v + pd + half < 1) admitted product column 0, which lies outside the matrix. The unified slot-based check zeroes that entry instead, matching the CPU referencemultiply_matrix_at_indexinsrc/MatrixFields/matrix_multiplication.jl.Performance
This change is performance neutral. The unified helpers compile down to essentially the same machine code as the 12 hand-written methods: generated PTX is bit-identical for the same-staggered shape combinations (after normalizing session-dependent pointer constants and label numbering), and the mixed-staggering combinations differ only by register renumbering and equivalent bounds computations (e.g.
< blockDimvs<= blockDim - 1), with equal or slightly lower register counts. Measured A100 kernel timings across all 12 shape combinations plus a chained product are unchanged (within run-to-run noise).Verification (on an A100)
test/MatrixFields/matrix_fields_broadcastingscripts (scalar 1–17, non-scalar 1–5) produce identical results before and after the refactor.< blockDimvs<= blockDim - 1), with equal or slightly lower register counts.🤖 Generated with Claude Code
https://claude.ai/code/session_011oirTEPPBEj5BWwhgYKMzn