Skip to content

Vectorise the tensor conversion layer - #13

Open
rodvals wants to merge 1 commit into
fix/multiple-fixesfrom
perf/vectorise-tensor-conversions
Open

Vectorise the tensor conversion layer#13
rodvals wants to merge 1 commit into
fix/multiple-fixesfrom
perf/vectorise-tensor-conversions

Conversation

@rodvals

@rodvals rodvals commented Jul 28, 2026

Copy link
Copy Markdown
Member

Dependency

Stacked on #9 and targets fix/multiple-fixes, because it must preserve the parsing semantics that pull request establishes: stored zeros are not edges, weights are kept, edge direction is not reversed, and isolated nodes survive. Merge after its parent, then rebase onto master and retarget.
It shares no commits with #12 and can merge before or after it.

Problem

Six conversion functions in utils/parsing.py stepped over one sparse entry at a time in Python, calling .item() per entry, so the cost was dominated by crossing the Python and torch boundary rather than by the work itself. Aggregation also carried three near identical branches for sum, max and min.

Changes

  • aggregate: one coo_matrix(...).tocsr() for sum, and a sort plus reduceat for max and min; three branches become one
  • per-layer graphs: one add_edge_list call per layer with a weight property, reciprocal edges collapsed by sorting
  • Laplacian: replica strengths via Tensor.index_add_, accumulated in float64 and cast back
  • tensor from graph list: per-layer COO blocks concatenated once
  • layer aggregation: accumulate from the first layer rather than from a fresh zero matrix

No conversion changes what it returns.

Verification

  • tests/test_parsing_vectorisation.py: Every conversion is compared against a reference implementation carrying the original per-entry algorithm, so equivalence is demonstrated rather than assumed.
  • measured at 400 nodes, 3 layers and 29791 non-zeros, base branch against this one: aggregate sum 323.1 ms to 2.2 ms, max 371.3ms to 3.4ms, min 361.9ms to 8.4ms, Laplacian 853.2ms to 5.9ms

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant