Counting tropical

[xuanzhaogao]

No description provided.

[codecov]

Codecov Report

❌ Patch coverage is 86.40000% with 34 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.19%. Comparing base (6d22701) to head (64f1b08).

Files with missing lines	Patch %	Lines
crates/tropical-gemm/src/crt.rs	79.46%	23 Missing ⚠️
crates/tropical-gemm/src/types/modp.rs	83.33%	4 Missing ⚠️
...rates/tropical-gemm/src/testing/bigint_semiring.rs	89.65%	3 Missing ⚠️
crates/tropical-gemm/src/api.rs	93.10%	2 Missing ⚠️
crates/tropical-gemm/src/mat/ref_.rs	75.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #42      +/-   ##
==========================================
- Coverage   96.41%   94.19%   -2.22%     
==========================================
  Files          18       22       +4     
  Lines         892     1086     +194     
==========================================
+ Hits          860     1023     +163     
- Misses         32       63      +31     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

This PR adds full-stack support for “counting tropical” GEMM: new semiring direction tags (Max/Min), CRT-based exact counting on CPU and GPU (including a u64 fast-path), plus Python and Julia bindings and accompanying design specs/bench scripts.

Changes:

• Introduce TropicalDirection (Max/Min) and make CountingTropical direction-aware; add ReprTransparentTropical to safely support scalar-slice reinterpretation for transparent semiring wrappers.
• Add GPU counting CRT driver (BigInt + u64 fast-path), AoS pair layout helpers, counting CUDA kernels (including ones-specialized + warpk-transposed-B variants), and new CUDA benchmarks.
• Add Python bindings/tests for CPU+GPU count-ground-states and a Julia package (CountingTropicalGEMM.jl) with tests/benchmarks; add multiple design spec documents.

Reviewed changes

Copilot reviewed 69 out of 71 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
sanity_julia.jl	Local Julia sanity script for CountingTropical GPU matmul (currently contains machine-specific paths).
bench_julia_vs_rust.jl	Julia benchmark to compare CUDA.jl generic CountingTropical matmul vs Rust (currently contains machine-specific paths).
docs/superpowers/specs/2026-04-27-julia-binding-design.md	Design spec for Julia binding architecture and FFI surface.
docs/superpowers/specs/2026-04-27-counting-u64-fastpath-design.md	Design spec for u64 fast-path in GPU driver + CRT helpers.
docs/superpowers/specs/2026-04-27-counting-kernel-warpk-transposed-design.md	Design spec for warpk kernel reading transposed B for coalesced loads.
docs/superpowers/specs/2026-04-27-counting-kernel-ones-specialized-design.md	Design spec for ones-specialized counting kernels.
docs/superpowers/specs/2026-04-27-counting-kernel-aos-design.md	Design spec for AoS (value,count) layout for counting kernels.
docs/superpowers/specs/2026-04-21-counting-tropical-cuda-design.md	Prior design spec for GPU counting kernel + CRT driver.
docs/superpowers/specs/2026-04-21-counting-tropical-crt-design.md	Prior design spec for CRT/BigInt CPU driver and bound contract.
docs/superpowers/specs/2026-04-21-counting-tropical-compose-design.md	Prior design spec for making CountingTropical compose with GEMM pipeline.
docs/superpowers/specs/2026-04-21-counting-kernel-tiled-design.md	Prior design spec for tiled CUDA counting kernel.
crates/tropical-gemm/tests/counting_crt.rs	New/updated CRT correctness tests vs BigInt oracle + edge cases.
crates/tropical-gemm/tests/counting_compose.rs	Tests for CountingTropical composing through tropical_matmul_t (Max/Min, ties, multiplication).
crates/tropical-gemm/src/types/traits.rs	Adds ReprTransparentTropical marker trait for safe scalar↔element reinterpretation.
crates/tropical-gemm/src/types/modp.rs	Adds Mod<P> scalar for CRT residue arithmetic.
crates/tropical-gemm/src/types/mod.rs	Wires new direction + Mod modules and re-exports Max/Min/Mod/ReprTransparentTropical.
crates/tropical-gemm/src/types/min_plus.rs	Implements ReprTransparentTropical for TropicalMinPlus.
crates/tropical-gemm/src/types/max_plus.rs	Implements ReprTransparentTropical for TropicalMaxPlus.
crates/tropical-gemm/src/types/max_mul.rs	Implements ReprTransparentTropical for TropicalMaxMul.
crates/tropical-gemm/src/types/direction.rs	Introduces TropicalDirection with Max and Min tags.
crates/tropical-gemm/src/types/counting.rs	Makes CountingTropical direction-aware (Max/Min), disables SIMD claim, adds Min tests.
crates/tropical-gemm/src/types/and_or.rs	Implements ReprTransparentTropical for TropicalAndOr.
crates/tropical-gemm/src/testing/mod.rs	Adds test-only module wiring (feature-gated).
crates/tropical-gemm/src/testing/bigint_semiring.rs	Adds BigInt-based reference implementation for CRT tests.
crates/tropical-gemm/src/simd/kernels/portable.rs	Updates microkernel signatures to take *const T and fixes tests accordingly.
crates/tropical-gemm/src/simd/kernels/neon.rs	Updates NEON kernels to accept *const T and casts internally to scalar pointers.
crates/tropical-gemm/src/simd/kernels/avx2.rs	Updates AVX2 kernels to accept *const T and casts internally to scalar pointers; updates tests.
crates/tropical-gemm/src/mat/ref_.rs	Refactors MatRef to store element slices, adds safe scalar conversion via ReprTransparentTropical.
crates/tropical-gemm/src/mat/owned.rs	Refactors Mat APIs to use element slices; adds from_elements convenience alias.
crates/tropical-gemm/src/mat/ops.rs	Updates multiplication trait bounds to include Default.
crates/tropical-gemm/src/lib.rs	Exposes new CRT module and Max/Min exports; wires test helpers.
crates/tropical-gemm/src/core/packing.rs	Packing now operates on element type T (requires Copy + Default) instead of TropicalScalar.
crates/tropical-gemm/src/core/kernel.rs	Microkernel interface now uses *const T rather than *const T::Scalar.
crates/tropical-gemm/src/api.rs	Adds tropical_matmul_t for compound elements; updates scalar APIs to require ReprTransparentTropical.
crates/tropical-gemm/Cargo.toml	Adds testing feature and BigInt/Integer deps.
crates/tropical-gemm-python/tests/test_count_ground_states_gpu.py	Adds GPU round-trip tests for counting.
crates/tropical-gemm-python/tests/test_count_ground_states.py	Adds/updates CPU round-trip tests for counting + typed errors.
crates/tropical-gemm-python/src/lib.rs	Adds CPU + GPU counting bindings (CRT BigInt counts).
crates/tropical-gemm-python/Cargo.toml	Adds num-bigint dependency.
crates/tropical-gemm-cuda/src/pair.rs	Adds AoS pair types + packing helpers + unit tests.
crates/tropical-gemm-cuda/src/lib.rs	Exposes counting CRT APIs + CountedMatU64; wires new modules.
crates/tropical-gemm-cuda/src/gpu_mat.rs	Updates MatRef upload path to require ReprTransparentTropical.
crates/tropical-gemm-cuda/src/error.rs	Adds InvalidState error variant.
crates/tropical-gemm-cuda/src/crt.rs	Implements GPU CRT driver (BigInt + u64 fast-path) and shared per-prime runner.
crates/tropical-gemm-cuda/src/counting_kernel.rs	Adds counting kernel launcher and dispatch, including ones-specialized path.
crates/tropical-gemm-cuda/src/context.rs	Adds compilation/registration of counting kernels and warpk dispatch constants/dims.
crates/tropical-gemm-cuda/examples/bench_warpk_crossover.rs	Bench for naive vs warpk (transposed-B) crossover.
crates/tropical-gemm-cuda/examples/bench_maxplus.rs	Bench for MaxPlus GPU kernel throughput reference.
crates/tropical-gemm-cuda/examples/bench_kernel_only.rs	Kernel-only bench for AoS vs ones-specialized kernels.
crates/tropical-gemm-cuda/examples/bench_e2e_u64.rs	End-to-end bench: BigInt vs u64 fast-path.
crates/tropical-gemm-cuda/examples/bench_counting.rs	End-to-end bench: GPU vs CPU counting.
crates/tropical-gemm-cuda/Cargo.toml	Builds as cdylib + adds BigInt-related deps.
CountingTropicalGEMM.jl/test/runtests.jl	Julia package tests incl. reference parity, bounds, and dimension mismatch.
CountingTropicalGEMM.jl/src/CountingTropicalGEMM.jl	Julia FFI wrapper to C ABI with library resolution and error mapping.
CountingTropicalGEMM.jl/bench/bench_huge.jl	Julia kernel-only huge-shape benchmark.
CountingTropicalGEMM.jl/bench/bench.jl	Julia kernel-only benchmark driver.
CountingTropicalGEMM.jl/README.md	Julia package build/install/usage docs.
CountingTropicalGEMM.jl/Project.toml	Julia package metadata.
Cargo.lock	Locks new Rust dependencies (num-bigint/num-integer).
.gitignore	Ignores Julia Manifest for the new Julia package.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Same as the BigInt path: avoid assert_eq! on user-controlled input lengths in this Result-returning API; return a CudaError instead of panicking.

Suggested change

	assert_eq!(a_values.len(), m * k);
	assert_eq!(b_values.len(), k * n);
	if a_values.len() != m * k {
	return Err(CudaError::InvalidState(
	format!(
	"invalid a_values length: got {}, expected {} (m * k)",
	a_values.len(),
	m * k
	)
	.into(),
	));
	}
	if b_values.len() != k * n {
	return Err(CudaError::InvalidState(
	format!(
	"invalid b_values length: got {}, expected {} (k * n)",
	b_values.len(),
	k * n
	)
	.into(),
	));
	}

[Copilot]

This script hard-codes a machine-specific local path (Pkg.add(path = "/mnt/home/.../GenericTensorNetworks.jl")), which will fail for other developers/CI. Consider removing this file from the repo, or make the path configurable (e.g., via ENV) / use a relative path / documented Pkg.develop instructions instead.

[Copilot]

This benchmark script hard-codes a local filesystem path for GenericTensorNetworks.jl, which makes it non-portable and likely to break in other environments. Prefer documenting setup steps and using Pkg.develop/relative paths or environment configuration rather than committing an absolute path.

Suggested change

	Pkg.add(path = "/mnt/home/xgao1/work/better_gpu_gemm/GenericTensorNetworks.jl")
	if !haskey(ENV, "GENERIC_TENSOR_NETWORKS_PATH")
	error("Set GENERIC_TENSOR_NETWORKS_PATH to a local GenericTensorNetworks.jl checkout before running this benchmark.")
	end
	gtn_path = abspath(expanduser(ENV["GENERIC_TENSOR_NETWORKS_PATH"]))
	isdir(gtn_path) || error("GENERIC_TENSOR_NETWORKS_PATH does not exist or is not a directory: " * gtn_path)
	Pkg.develop(path = gtn_path)

[Copilot]

Doc comment is misleading: raw() is an accessor, not a way to reconstruct/create a Mod<P> from an existing raw representative. Consider rewording to avoid implying there is a from_raw constructor, or add an explicit from_raw_unchecked/from_raw API if that capability is intended.

[Copilot]

count_ground_states_gpu_py constructs a fresh CudaContext::new() on every call, which will redo NVRTC compilation/context setup and defeats the crate’s existing global-context caching. Use the shared/global context (or a cached context per process/device) instead of creating a new one per invocation.

Suggested change

	let ctx = CudaContext::new().map_err(|e| format!("CUDA init: {}", e))?;
	match direction {
	"max" => count_ground_states_gpu::<f32, Max>(
	&ctx, &a_data, m, k, &b_data, n, &bound,
	)
	.map_err(|e| format!("GPU compute: {}", e)),
	"min" => count_ground_states_gpu::<f32, Min>(
	&ctx, &a_data, m, k, &b_data, n, &bound,
	// Reuse the process-global CUDA context so we do not pay
	// initialization / NVRTC compilation costs on every call.
	let ctx = CudaContext::global().map_err(|e| format!("CUDA init: {}", e))?;
	match direction {
	"max" => count_ground_states_gpu::<f32, Max>(
	ctx, &a_data, m, k, &b_data, n, &bound,
	)
	.map_err(|e| format!("GPU compute: {}", e)),
	"min" => count_ground_states_gpu::<f32, Min>(
	ctx, &a_data, m, k, &b_data, n, &bound,

[Copilot]

Converting BigInt to Python int via py.eval("int(<decimal>)") is unnecessarily slow (parses Python code per element) and relies on code-eval. Prefer constructing PyLong directly (e.g., via the CPython API PyLong_FromString / PyLong_FromUnicodeObject, or calling builtins.int with the decimal string without eval).

Suggested change

	let counts_py: Vec<PyObject> = result
	.counts
	.into_iter()
	.map(|bn| {
	let s = bn.to_string();
	// Use Python's built-in int() on the decimal string.
	py.eval(
	std::ffi::CString::new(format!("int({})", s))
	.unwrap()
	.as_c_str(),
	None,
	None,
	)
	.map(|b| b.unbind())
	let py_int = py.import("builtins")?.getattr("int")?;
	let counts_py: Vec<PyObject> = result
	.counts
	.into_iter()
	.map(|bn| {
	let s = bn.to_string();
	py_int.call1((s,)).map(|b| b.unbind())

[Copilot]

Same issue as the GPU binding: BigInt → Python int conversion uses py.eval("int(<decimal>)"), which is very expensive for large outputs and should be replaced with direct PyLong creation (no eval).

Suggested change

	// Convert BigInt → Python int via decimal string representation.
	let counts_py: Vec<PyObject> = result
	.counts
	.into_iter()
	.map(|bn| {
	let s = bn.to_string();
	// Use Python's built-in int() on the decimal string.
	py.eval(
	std::ffi::CString::new(format!("int({})", s))
	.unwrap()
	.as_c_str(),
	None,
	None,
	)
	.map(|b| b.unbind())
	})
	// Convert BigInt → Python int directly via CPython's PyLong API.
	fn bigint_to_pyobject(py: Python<'_>, bn: BigInt) -> PyResult<PyObject> {
	let bytes = bn.to_signed_bytes_le();
	let ptr = unsafe {
	pyo3::ffi::_PyLong_FromByteArray(
	bytes.as_ptr(),
	bytes.len(),
	1,
	1,
	)
	};
	unsafe { Bound::from_owned_ptr_or_err(py, ptr).map(|obj| obj.unbind()) }
	}
	let counts_py: Vec<PyObject> = result
	.counts
	.into_iter()
	.map(|bn| bigint_to_pyobject(py, bn))

[Copilot]

This public API uses assert_eq! for input length validation, which will panic on user error. For a Result-returning GPU entry point, prefer returning a structured error (e.g., CudaError::DimensionMismatch / CudaError::InvalidState) so callers (including Python/Julia bindings) can handle it without a panic.

Suggested change

	assert_eq!(a_values.len(), m * k);
	assert_eq!(b_values.len(), k * n);
	if a_values.len() != m * k {
	return Err(CudaError::DimensionMismatch);
	}
	if b_values.len() != k * n {
	return Err(CudaError::DimensionMismatch);
	}