Conversation
There was a problem hiding this comment.
Pull request overview
This PR expands TorchMorph’s morphology layer with new CUDA-accelerated primitives and Python APIs, aiming to more closely mirror SciPy ndimage behavior while keeping (B, C, Spatial...) batch-channel semantics explicit.
Changes:
- Adds fused CUDA grey morphology (erosion/dilation) and composes higher-level grey operators (opening/closing/gradient/laplace/top-hats) in Python.
- Refactors binary morphology to dispatch CUDA erosion/dilation through a new generalized fused CUDA primitive (with CPU convolution fallback).
- Introduces
iterate_structure(SciPy-compatible) plus accompanying tests, benchmarks, and documentation updates.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| torchmorph/morphology/structure.py | Adds iterate_structure and helper routines to detect/fast-path generated connectivity structures. |
| torchmorph/morphology/grey.py | New grey morphology Python API and argument normalization that calls CUDA kernels in _C. |
| torchmorph/morphology/binary.py | Routes CUDA inputs through new fused CUDA primitive; keeps CPU fallback via conv/pad. |
| torchmorph/morphology/init.py | Exports new grey operators and iterate_structure. |
| torchmorph/csrc/torchmorph.cpp | Extends pybind module with new _C entry points for grey/binary morphology. |
| torchmorph/csrc/grey_kernel.cu | Implements fused grey morphology CUDA kernel with SciPy-like border modes. |
| torchmorph/csrc/binary_kernel.cu | Implements fused binary morphology CUDA kernel (erosion/dilation). |
| torchmorph/init.py | Re-exports new morphology APIs at the top-level package. |
| test/test_structure.py | Adds SciPy parity tests for iterate_structure and origin handling. |
| test/test_grey.py | Adds CUDA SciPy parity tests for grey morphology operators, modes, and output semantics. |
| test/test_binary.py | Updates/expands CUDA SciPy parity tests for binary morphology, including fallback checks. |
| README.md | Updates feature bullets and adds a NaN-behavior note. |
| pytest.ini | Adds a cuda marker entry. |
| docs/adr/0001-generalized-binary-morphology-cuda-primitive.md | ADR documenting the new generalized binary CUDA primitive decision. |
| CONTEXT.md | Adds glossary/terminology guidance aligned with SciPy ndimage language. |
| benchmark/structure.py | New benchmark script for generate/iterate structure helpers. |
| benchmark/grey.py | New benchmark script for grey morphology operators. |
| benchmark/binary.py | Refactors binary benchmarks into a CLI with operation selection. |
| benchmark/generate.py | Removes legacy generate benchmark script (superseded by benchmark/structure.py). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
1
to
2
| #include <torch/extension.h> | ||
|
|
Comment on lines
+301
to
+305
| auto d_struct_vals = torch::from_blob( | ||
| h_struct_vals.data(), {num_struct}, opts_f).to(input.device()); | ||
| auto d_struct_meta = torch::from_blob( | ||
| h_struct_meta.data(), {num_struct * (ndim_spatial + 1)}, opts_i | ||
| ).to(input.device()); |
Comment on lines
+227
to
+229
| auto d_struct_meta = torch::from_blob( | ||
| h_struct_meta.data(), {num_struct * (ndim_spatial + 1)}, opts_i | ||
| ).to(input.device()); |
Comment on lines
+75
to
+87
| if input.is_cuda: | ||
| x = input != 0 | ||
| input_bool = x | ||
| mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None | ||
|
|
||
| def step(value: Tensor) -> Tensor: | ||
| return _binary_morphology_cuda_step( | ||
| value, | ||
| structure, | ||
| border_value, | ||
| origin, | ||
| mode=mode, | ||
| ) |
Comment on lines
+83
to
+88
| for origin_value, structure_size in zip(origin_list, struct.shape): | ||
| min_origin = -(structure_size // 2) | ||
| max_origin = (structure_size - 1) // 2 | ||
| if not min_origin <= origin_value <= max_origin: | ||
| raise ValueError("invalid origin") | ||
|
|
Comment on lines
+37
to
+41
| pytest.param(scipy_generate_binary_structure(2, 1), 2, id="2d_cross_2"), | ||
| pytest.param(scipy_generate_binary_structure(2, 2), 3, id="2d_full_3"), | ||
| pytest.param(scipy_generate_binary_structure(3, 1), 2, id="3d_cross_2"), | ||
| pytest.param(np.array([[False, True], [True, True]], dtype=bool), 3, id="asymmetric"), | ||
| pytest.param(np.zeros((3, 3), dtype=bool), 2, id="empty"), |
* optimal transport * debug ot function and add test * add basic benchmark * decouple the iteration from the derivation * Implement Sinkhorn iterations in CUDA * Add Sinkhorn CUDA benchmark and correctness tests * Add CUDA support for Sinkhorn log-domain iterations and related tests. * Refactor Sinkhorn functions and test gradient accuracy * more test * try test * test on cuda kernal * tighten numeric gradient test and support more dimensions * cancel unnecessary change in test.yml Removed TORCH_CUDA_ARCH_LIST environment variable from the workflow. * Rework Sinkhorn OT: standard epsilon convention, log-space end-to-end, batched (n, d) API - Use the standard convention K = exp(-C / epsilon); POT baselines need no parameter inversion anymore. - Keep log potentials in log space end-to-end: the log-domain CUDA solver returns (log_u, log_v), gradients are computed as f = eps * log_u, and the transport plan is reconstructed as exp(log_u - C/eps + log_v) for all backends so it cannot underflow at small epsilon. - Every backend takes strictly (n, d) batched histograms sharing one (d, d) cost matrix; the CUDA kernels are batched over a 2-D launch grid, removing the previous B=C=1 restriction. - Merge the four kernels into two (scaling_update, log_scaling_update), fixing the uncoalesced K^T reads; collapse entry points into solve(); delete the dead sinkhorn__nobatch and loss_test. - CUDA hygiene: current-stream launches, CUDAGuard, TORCH_CHECK input validation, kernel-launch checks, and a NaN guard for all -inf LSE rows. - Export SinkhornSolver / build_cost_matrix from the package; tests import it directly, run the torch backend on CPU, and add a small-epsilon stability test against POT's sinkhorn_log. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Apply black formatting to optimal transport files Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Make SinkhornSolver a differentiable nn.Module with automatic device dispatch - SinkhornSolver now inherits torch.nn.Module: forward(source, target, cost_matrix) returns the (n,) transport distances, and a custom autograd.Function implements backward via the centered dual potentials (envelope theorem), chained through the differentiable normalization. - Replace the confused torch/cuda/cuda_log backend strings with one orthogonal option: log_space=True/False selects the algorithm, while the implementation is picked from the input automatically (fused CUDA kernels for float32 CUDA tensors, pure torch ops elsewhere - including a logsumexp-based log-space path that also runs on CPU and float64). - Convenience methods plan() and potentials() expose the transport plan and dual potentials; extra_repr() reports the module configuration. - Tests cover both log_space modes across forward shapes, marginals, POT baselines, fused-vs-CPU consistency, small-epsilon stability, and autograd-vs-numeric gradients; benchmark updated to the module API. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Expand Sinkhorn tests to a full device x log_space x dtype matrix Previously most tests picked a single device (CUDA when available), so the CPU paths went untested on GPU machines, POT baselines only covered the plain iteration, and backward never ran through the fused kernels. Now: - Shared DEVICES parametrization runs every applicable test on both CPU and CUDA (skipping cleanly when CUDA is absent). - POT distance/plan baselines and the small-epsilon stability test cover both log_space modes and both devices (torch logsumexp path on CPU, fused log kernel on CUDA). - New tests: dtype/device consistency (float64 CUDA exercises the torch fallback on GPU), backward through the fused float32 kernels vs CPU, fused-vs-CPU potentials, and threshold early stopping reaching the same solution on both algorithms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Optimize Sinkhorn CUDA path: batch tiling, online logsumexp, CUDA graphs Kernel changes (both memory-bandwidth-bound, so all wins target matrix traffic and launch latency): - Batch tiling: one block now streams a matrix row once and applies it to up to 8 batch items held in registers, cutting d^2 matrix reads by up to 8x for batched inputs. A TILE=1 template instantiation keeps the n=1 fast path free of predication overhead. - Online logsumexp: the log-domain update now computes the LSE in a single pass over the row (running max with rescaled running sum, NaN-free -inf merging), halving matrix traffic vs the two-pass max-then-sum version. - Warp-shuffle reductions replace the shared-memory tree (one barrier instead of eight per reduction). Python-side: - Long fused runs (max_iter >= 100) capture a 25-iteration chunk into a torch.cuda.CUDAGraph and replay it, hiding the 2-launches-per-iteration overhead that dominates small problems; falls back to plain launches when capture is unavailable (extra fixed-point iterations are harmless). - Skip the (d, d) transpose copy when the cost/kernel matrix is symmetric, which build_cost_matrix outputs always are. Tests: fused-vs-CPU now covers single-item, partial-tile, and multi-tile batches (n = 1, 3, 20) plus a short-run case below the CUDA-graph threshold to exercise the plain-launch fused path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Test gradients w.r.t. both source and target The backward pass already returned gradients for both marginals (the centered dual potentials f and g), but the tests only verified the source side. The numeric-derivative test now perturbs source and target independently, and the fused-kernel backward test compares both gradients against CPU. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: maokovski <1498743918@qq.com> Co-authored-by: maokovski <“1498743918@qq.com”> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
- Add torchmorph/_validation.py with shared validate_bcs_input / validate_output; remove three near-identical copies from binary, grey, and distance_transform (error messages preserved). - binary: drop per-iteration tensor .clone() in the stabilize loop (step/where already return fresh tensors) and skip the redundant bool conversion when the value is already bool + contiguous. - Move _validate_origin to structure.py; grey now reuses _normalize_origin/_validate_origin instead of hand-rolling it. - grey: pack the shared structuring-element args via _element_kwargs across the composite operators. Net ~-108 lines; behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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
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.
No description provided.