Internals: RBF Weight Building System
This document explains the architecture of the solve system, which builds stencil-wise (ELL-format) weight matrices — StencilWeights — for RBF operators. It is intended for developers who want to understand or extend the package.
Call Graph
Nodes are colored by architecture layer — indigo api.jl (routing and entry points), amber execution.jl (allocation and parallel kernel execution), and green assembly.jl (the per-stencil mathematics: build A, build b, and the local solve). The execution layer runs one _build_stencil! per evaluation point and writes each stencil's k weights directly into its column of the ELL value matrix (fill_entries!); once every batch completes, _assemble_weights wraps the filled value and index buffers into StencilWeights — no sparse assembly pass. See the layer breakdown below.
Architecture Overview
The solve system is organized into four layers, each with a clear responsibility:
src/solve/
├── types.jl # Layer 0: Shared data structures
├── assembly.jl # Layer 1: Pure mathematical operations
├── execution.jl # Layer 2: Parallel execution & allocation
└── api.jl # Layer 3: Entry points & routing| Layer | File | Purpose |
|---|---|---|
| 3 | api.jl | Entry points from operators |
| 2 | execution.jl | Memory allocation, parallel kernel execution, ELL weight assembly |
| 1 | assembly.jl | Pure mathematics: collocation matrix, RHS, stencil assembly |
| 0 | types.jl | Shared types and arity helpers |
The StencilWeights container itself lives outside this tree in two layers: src/ellsparse/ holds the self-contained EllSparse module (generic SELL-C/ELL storage, the CPU/GPU mul! apply kernels, transpose-map adjoint, and conversions), and src/stencil_weights.jl wraps it as the thin RBF-policy type StencilWeights.
System Flow
The system builds ELL-format weight matrices by computing stencil weights in parallel for each evaluation point:
┌───────────┐ ┌───────────────┐ ┌─────────────────┐ ┌─────────────────┐ ┌────────────────────┐ ┌──────────────┐
│ User Code │───>│ api.jl │───>│ execution.jl │───>│ assembly.jl │───>│ execution.jl │───>│Return to user│
│ │ │ Route request │ │Allocate & launch│ │Build A,b, solve │ │ Wrap ELL buffers as│ │ │
│ │ │ │ │ kernel │ │ A\b │ │ StencilWeights │ │ │
└───────────┘ └───────────────┘ └─────────────────┘ └─────────────────┘ └────────────────────┘ └──────────────┘Key steps: 2. Route — Operator calls _build_weights, which extracts data and applies the operator to the basis
Allocate —
allocate_ellcreates one densek × N_evalvalue matrix per operator component plus a sharedk × N_evalInt32neighbor-index matrix (build_weights_kernelfirst rejects a raggedadjlwith anArgumentError— ELL storage requires a uniform stencil size)Build A — Construct the collocation matrix for each stencil's k-nearest neighbors
Build b — Construct the RHS by applying the differential operator to the basis at the evaluation point
Solve — Compute weights via
A \ bAssemble —
fill_entries!writes each stencil's weights into its ELL column, then_assemble_weightswraps the buffers inStencilWeights(a tuple of them for multi-component operators, all sharing the index matrix)
Layer Details
Layer 0: types.jl — Data Structures
Defines shared types used throughout the solve system.
Key types:
_num_ops,_prepare_buffer— Operator arity helpers for pre-allocating RHS buffers (dispatch onTuplevs non-Tuple)
When to modify: Adding new shared data structures or arity helpers.
Layer 1: assembly.jl — Pure Mathematics
Contains all mathematical operations for building stencils. No I/O, no parallelism — fully testable in isolation.
Key functions:
_build_stencil!(A, b, ℒrbf, ℒmon, data, eval_point, basis, mon, k)— Assembles and solves local system_build_collocation_matrix!(A, data, basis, mon, k)— Fills the collocation matrix_build_rhs!(b, ℒrbf, ℒmon, data, eval_point, basis, k)— Builds RHS vector
When to modify: Changing stencil mathematics or adding new RBF formulations.
Layer 2: execution.jl — Parallel Execution
Handles memory management and parallel kernel execution via KernelAbstractions.jl.
Key functions:
build_weights_kernel(...)— Main orchestrator: allocates, launches kernel, returnsStencilWeightsallocate_ell(TD, k, N_eval, num_ops)— Allocates the densek × N_evalvalue matrices and the sharedInt32index matrixlaunch_kernel!(...)— Dispatches parallel kernel over batches@kernel weight_kernel(...)— Per-batch kernel that builds weights for each evaluation point, writing each stencil's column viafill_entries!(orfill_dirichlet_column!for Dirichlet identity rows)_assemble_weights(vals_list, idx, N_data, num_ops)— Wraps the filled buffers intoStencilWeights
When to modify: Improving parallelism, changing batch strategy, or optimizing memory allocation.
Layer 3: api.jl — Entry Points
Public-facing entry points that operators call to build weights.
Key functions:
_build_weights(ℒ, op)— Entry from operator, extracts configuration_build_weights(ℒ, data, eval_points, adjl, basis)— Applies operator to basis_build_weights(...; batch_size, device)— Routes to kernel execution
When to modify: Adding new operator entry points or changing routing logic.
Key Concepts
Stencils
A stencil approximates a differential operator at a point using its k nearest neighbors:
The weights
Collocation Matrix Structure
The collocation matrix has a block structure:
where
Weight Storage: ELL Format (two layers)
The built weights are stored stencil-wise in StencilWeights — since v0.8 a thin RBF-policy wrapper over the self-contained generic module RadialBasisFunctions.EllSparse (src/ellsparse/, the seed of a future standalone package). The storage is a dense k × N_eval value matrix — column Int32 index structure recording which data point each weight multiplies. The logical size stays (N_eval, N_data), and parent(W) returns the dense value matrix — the supported handle for in-place mutation and for AD losses over built weights (e.g. sum(parent(W) .^ 2)). Gradient-family operators return a tuple of StencilWeights that all alias one EllSparse.SellStructure (one index matrix, one transpose map).
Format decision: ELL is single-slice SELL. EllSparse implements SELL-C (Sliced ELLpack) with the slice height C in the type domain; our stencil-major layout is exactly SellMatrix{T, 1}, and because RBF stencils have uniform length k there is zero padding — SELL and plain ELL are byte-identical here. C = 32 is the slot-major coalesced orientation cuSPARSE standardized on (EllSparse.reslice switches orientation explicitly; Adapt never changes layout). Revisit triggers for this design: ragged stencil support, real-GPU benchmark results from benchmark/sell_gpu.jl, and a CuSparseMatrixSELL wrapper landing in CUDA.jl (JuliaGPU/CUDA.jl#2782).
Evaluation is a gather, not a sparse matvec. Applying an operator runs Threads.@threads loop with @simd inner dot products; on GPU backends (an operator moved to device via Adapt, e.g. cu(op) — values and index structure adapt together) it is a KernelAbstractions kernel. Measured against the old SparseMatrixCSC matvec this is roughly 7.6× faster at N = 100k, k = 50 with 13 threads, and about 1.5× faster serially (issue #156). The adjoint apply W' * x (also the AD pullbacks' input-cotangent path) gathers through a precomputed transpose map (EllSparse.TransposeMap, a counting-sorted per-column adjacency built once at construction): one work item per column with a fixed summation order — deterministic under any thread count, no atomics, and the same kernel runs on GPU backends.
Dirichlet columns are padded. A Dirichlet boundary row is written by fill_dirichlet_column! as weight 1 at slot 1 and zero pads in the remaining slots, with every slot carrying the evaluation point's own index (in-bounds by construction). Converting with sparse(W) combines duplicate indices, so these padded columns collapse to single-entry identity rows — exactly what the previous sparse storage produced.
Sparse conversion is the assembly path. sparse(op) / SparseMatrixCSC(op) convert an operator's weights to SparseMatrixCSC (a tuple of them for gradient-family operators); this is the supported path for global system assembly and implicit solves (sparse(op) \ rhs). weights(op) returns the StencilWeights themselves. One exception: VirtualPartial operators still use SparseMatrixCSC weights internally, since they combine the union sparsity of two stencil sets.
Performance Notes
Memory allocation:
Fixed-size ELL allocation up front (
k × N_evalper component) — no non-zero counting pass, no COO/CSC assemblyupdate_weights!rewrites the dense value matrix in place; the index matrix is frozen at construction
Parallelization:
Batch processing to control memory usage
Work arrays reused within each batch
KernelAbstractions.jl structures the batch kernel; weight building is currently CPU-only (issue #88), while evaluation runs on GPU backends via the ELL apply kernel
Advanced: Hermite Interpolation
For problems with boundary conditions, the package supports Hermite interpolation. This is triggered by providing boundary information (is_boundary, boundary_conditions, normals) to the operator constructor.
Hermite interpolation modifies the collocation matrix to incorporate boundary operators (Dirichlet, Neumann, Robin) at boundary nodes. This is an advanced feature primarily used for solving PDEs with explicit boundary conditions. Hermite operators require eval_points to equal data — the boundary classification is defined on the data points, so a differing evaluation set throws an error rather than silently producing wrong stencils.
Navigation Guide
| Want to... | Look in... |
|---|---|
| Modify stencil mathematics | assembly.jl |
| Improve parallelism or batching | execution.jl |
| Add a new operator entry point | api.jl |
| Understand ELL allocation | execution.jl → allocate_ell |
| Change weight storage or the apply kernels | src/ellsparse/ → EllSparse.SellMatrix, mul! |
| Change RBF-specific weight policy | src/stencil_weights.jl → StencilWeights wrapper |
| Debug a specific stencil | assembly.jl → _build_stencil! |