diff --git a/.github/workflows/oracle-tests.yml b/.github/workflows/oracle-tests.yml index 3953bf65..f31d3eba 100644 --- a/.github/workflows/oracle-tests.yml +++ b/.github/workflows/oracle-tests.yml @@ -44,9 +44,11 @@ jobs: python numjvm-tests-oracle/generate_cases.py --check python numjvm-tests-oracle/generate_matmul_cases.py --check python numjvm-tests-oracle/generate_contraction_cases.py --check + python numjvm-tests-oracle/generate_parallel_cases.py --check python numjvm-tests-oracle/check_m3_1_contract.py python -m unittest discover -s numjvm-bench -p test_compare_expression.py python -m unittest discover -s numjvm-bench -p test_run_050_reference.py + python -m unittest discover -s numjvm-bench -p test_run_060_reference.py - name: Java tests, including the NumPy oracle replay run: ./gradlew test --stacktrace diff --git a/README.md b/README.md index 09c2c248..8160540e 100644 --- a/README.md +++ b/README.md @@ -14,10 +14,9 @@ Numerical computing for the Java Virtual Machine: N-dimensional arrays with NumPy semantics, built on the Java 25 Foreign Function & Memory API. No Python, NumPy or native libraries are needed at run time. -> **Status: 0.5.0-alpha implementation candidate (M3.1).** Experimental expression -> execution and real contractions extend the CPU APIs. Reference measurements show -> substantial AUTO regressions in some cases; see [0.5.0 status](docs/milestones/M3-1-status.md). -> Performance gates remain open and the API remains experimental. +> **Status: 0.6.0-alpha implementation candidate.** Optional CPU workers and +> generated fusion improvements are experimental. Reference performance and +> scaling remain unverified; see [0.6.0 status](docs/milestones/M4-0.6-status.md). ## Quick start @@ -144,6 +143,15 @@ Reference results show substantial AUTO slowdowns and allocation in some cases; see the [results](docs/benchmarks/M3-1-reference-results.md) and [reference protocol](docs/benchmarks/M3-1-protocol.md). +## Opt-in parallel CPU execution + +`NJ.parallel(workers)` provides a reusable bounded worker context for selected +contiguous float32/64 ufuncs, expressions and full-array sum. Existing calls +remain single-threaded. Confined scopes and unsupported layouts fall back to +the calling thread. Parallel floating-point sum uses a deterministic chunk +order, which can produce different low bits from sequential sum. See the +[parallel execution guide](docs/design/parallel-cpu.md). + ## Building Requires JDK 25. The Gradle wrapper downloads Gradle, and the build can provision a JDK 25 toolchain if none is diff --git a/ROADMAP.md b/ROADMAP.md index 57c1f2c8..48c5785a 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -8,6 +8,7 @@ | M2 | Optional CPU acceleration and portable matmul — **0.3.0-alpha** | Implementation candidate; release evidence pending | | M3 | JIT and fusion: tracing, IR, ClassFile API code generation — **0.4.0-alpha** | Merged; reference-performance evidence pending | | M3.1 | Fusion performance and real dot/vecdot — **0.5.0-alpha** | Implementation candidate; reference issues and open gates | +| M3.2 | Fusion allocation and opt-in CPU workers — **0.6.0-alpha** | Implementation candidate; reference evidence pending | | M4 | Ecosystem: Kotlin and Scala APIs, Arrow, DLPack, notebooks, Maven Central | Planned | | M5 | Advanced computing: GPU and experimental backends | Planned | | M6 | Production readiness: API freeze, compatibility, 1.0 | Planned | @@ -32,9 +33,9 @@ No multiworker scaling target applies to M2. ## Proposed post-M3 sequence 0.5.0 adds real dot/vecdot and fusion improvements; see its -[status and remaining gates](docs/milestones/M3-1-status.md). The proposed [0.6.0 plan](docs/milestones/M4-0.6-plan.md) prioritizes the -observed fusion regressions, then worker pools and deterministic parallel -reductions after an ownership design. +[status and remaining gates](docs/milestones/M3-1-status.md). The [0.6.0 plan](docs/milestones/M4-0.6-plan.md) addresses the +observed fusion regressions and introduces opt-in workers and deterministic +reductions. Its reference scaling and performance gates remain open. Consider a separately scoped 0.7.0 for portable solve/Cholesky and optional LAPACK before M4 ecosystem integration. Outstanding M2 and M3 reference-performance evidence remains open; no publication is implied by merging implementation code. diff --git a/docs/README.md b/docs/README.md index 46914d3a..f1e0ae2d 100644 --- a/docs/README.md +++ b/docs/README.md @@ -29,3 +29,5 @@ - [0.5.0 reference results](benchmarks/M3-1-reference-results.md) - [Proposed 0.6.0 milestone](milestones/M4-0.6-plan.md) +- [Opt-in parallel CPU execution](design/parallel-cpu.md) +- [0.6.0 reference protocol](benchmarks/M4-reference-protocol.md) diff --git a/docs/adr/ADR-010-parallel-execution.md b/docs/adr/ADR-010-parallel-execution.md new file mode 100644 index 00000000..d636c4d4 --- /dev/null +++ b/docs/adr/ADR-010-parallel-execution.md @@ -0,0 +1,40 @@ +# ADR-010: opt-in parallel CPU execution + +Status: proposed for 0.6.0. + +A reusable, explicitly closed execution context owns a bounded pool; no global +pool or implicit change to existing APIs is introduced. Calls remain synchronous. +The caller selects the worker limit; fixed, bounded chunking and deterministic +combination order do not depend on completion order. Small, unsupported or +worker-inaccessible arrays use the existing sequential implementation. + +## Ownership and safety + +Validate scope liveness and access on the calling thread before scheduling. +Only submit work for segments accessible to all target workers. A confined +`NJ.scope()` cannot be accessed by a pool worker: execute the original operation +on the calling thread, without copying or changing result ownership. Shared and +automatic memory may be eligible. Never assume caller accessibility implies +worker accessibility. Follow the existing output validation and overlap rules; +use parallel writes only for provably disjoint regions, with identical views +allowed only when each element is read before that element is written. + +A concurrent shared-scope close may cause an FFM access failure. It is not +transactional: output may be partially written. Await every submitted worker on +success, error or interruption before returning; restore interruption status +and propagate a documented exception. No worker may access arrays after return. +Closing the context must prevent new operations, wait for active operations, +and terminate the pool. Do not hold locks in a way that deadlocks nested calls; +worker-originated nested work must run sequentially or be rejected explicitly. +The caller must not mutate inputs or close their scope concurrently unless it +accepts the documented failure/partial-output behavior. + +## Numeric and native interaction + +Parallel sum uses fixed logical chunks and a fixed merge order independent of +worker count and scheduling. Reproducibility is scoped to the same algorithm, +dtype, backend and Java environment; it does not promise identical results to +sequential accumulation. Document NaNs, infinities and signed zero according to +the implemented arithmetic. Do not launch BLAS from NumJVM worker tasks or +modify process-wide native thread settings. Performance thresholds are +provisional until the reference campaign. No multiworker speedup is guaranteed. diff --git a/docs/adr/README.md b/docs/adr/README.md index 45c2a23d..c6271ef1 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -15,3 +15,4 @@ marked "Superseded by ADR-NNN" rather than deleted. - [ADR-007: explicit expression execution](ADR-007-expression-execution.md) (proposed) - [ADR-008: expression strategy and optional fused vectors](ADR-008-expression-strategy.md) (proposed) - [ADR-009: real dot and vecdot contractions](ADR-009-contractions.md) (proposed) +- [ADR-010: opt-in parallel CPU execution](ADR-010-parallel-execution.md) (proposed) diff --git a/docs/benchmarks/M3-1-protocol.md b/docs/benchmarks/M3-1-protocol.md index 66da5fac..07008bed 100644 --- a/docs/benchmarks/M3-1-protocol.md +++ b/docs/benchmarks/M3-1-protocol.md @@ -43,3 +43,10 @@ allocating methods only against other allocating methods. The proposed 1.5× large fused-output target, 1.8× allocating regression guard, and small/medium AUTO budgets remain open pending valid measurements. If a gate fails, report it and retain the failure; do not silently lower the target. + +The 0.6.0 benchmark module adds `ContractionBenchmark` for float32/64 vector +`dot` and `vecdot`, allocating and preallocated outputs. Preserve its JMH logs +as well as JSON: setup reports the observed backend and SIMD provider. Run the +existing `ExpressionLifecycleBenchmark` and `expressionCompileProbe` separately +for trace, first-call, compilation and warm costs; compilation-probe results +are single-process diagnostics, not a substitute for JMH distributions. diff --git a/docs/benchmarks/M3-1-reference-results.md b/docs/benchmarks/M3-1-reference-results.md index 2fa06c11..485257db 100644 --- a/docs/benchmarks/M3-1-reference-results.md +++ b/docs/benchmarks/M3-1-reference-results.md @@ -43,7 +43,7 @@ that run the recorded status also listed two untracked JSON files. For 1,048,576 elements, F-order CHAIN eager versus AUTO was 391.303 versus 3,664.833 µs for float32, and 3,002.260 versus 7,548.619 µs for float64. The earlier apparent F-order AUTO advantage therefore does not hold in these corrected cases. The -benchmark patch has not yet been applied to this repository. +benchmark correction is applied on the 0.6.0 branch. ## JFR and compilation @@ -81,3 +81,12 @@ compilation. The separate lifecycle JMH suite also ran. See the [protocol](M3-1-protocol.md) and the proposed [0.6.0 plan](../milestones/M4-0.6-plan.md). + +## 0.6.0 starting baseline + +The 0.6.0 work starts from the merged 0.5.0 implementation and oracle cache +fix. The reference timings above were measured at `4aff717` on the user's +machine, not at the later merge: source provenance is part of the measurement. +The clean follow-up was a focused investigation, not a complete replacement for +the earlier dirty full campaign. Keep those datasets distinct when comparing +future work. No reference-machine measurements have yet been made for 0.6.0. diff --git a/docs/benchmarks/M4-fusion-investigation.md b/docs/benchmarks/M4-fusion-investigation.md new file mode 100644 index 00000000..b6efcb5a --- /dev/null +++ b/docs/benchmarks/M4-fusion-investigation.md @@ -0,0 +1,51 @@ +# Fused vector investigation + +The supplied JDK 25 recordings reproduce the 0.5.0 allocation issue on the +reference host: float32, 1,048,576 elements, JMH GC profiler: AUTO ADD +16,785,621 B/op and CHAIN 33,568,252 B/op versus eager approximately 1,089 +and 3,267 B/op. JFR allocation samples point to `FloatFusedKernel.run` vector +loads, including `float[]` and `Float256Vector` objects. Samples are not counts; +see [reference results](M3-1-reference-results.md) for provenance and limits. + +The loop loads vectors into a `FloatVector[]` and sends that array through the +`FloatExpression.evaluate` interface at every vector-sized iteration. This +can cause objects to escape scalar replacement, but the profile does not prove +why HotSpot allocated them. Generate or otherwise specialize a whole traversal +loop and measure both allocation and speed before concluding that a change +fixes the cause. Confirm float64 and mixed independently. Small AUTO overhead +occurs even when eager is selected and needs a separate investigation. + +`FusedAllocationBenchmark` provides a focused JMH regression matrix. Run with +`-prof gc` and inspect `gc.alloc.rate.norm` as well as elapsed time, recording +provider and `NUMJVM_AUTO_DECISION` from its setup. For example, pass +`-PjmhArgs='FusedAllocationBenchmark -p size=1048576 -p layout=C -p types=FLOAT32 -p chain=ADD,CHAIN -prof gc'` +to `:numjvm-bench:jmh` with `-PsimdBench` and the required vector JVM options. + +An exploratory shared-machine JDK 25 run (one fork, three one-second warmups +and measurements, 1,048,576 float32 C-order elements) compared the old array +based path with a generated whole-run loop. AUTO ADD allocation fell from about +12.6 MB/op to 8.1 kB/op and CHAIN from 25.2 MB/op to 13.7 kB/op. These are +separate runs on a shared machine, not reference-machine acceptance results. +The generated loop keeps vector intermediates in locals, preserving a separate +scalar-tail path; mixed graphs are handled by a separately generated typed loop. Repeat the full matrix on the reference +machine before closing performance gates. + +A short exploratory JDK 25 smoke run also exercised C- and F-order float64 +and mixed CHAIN cases with the generated loops and GC profiler. Mixed +float32-to-float64 transitions remain explicit so float32 nodes round before +promotion. These smoke tests do not replace the reference campaign. + +For a single unary or binary ufunc with array inputs, the eager scheduler needs +no temporary arrays. AUTO now checks the small-input eager path before building +a generated traversal, and that single-ufunc case avoids opening a temporary +scope. More complex small graphs still incur scheduler and scope costs; the +reference-machine small-input budget has not yet been remeasured. + +AUTO selects eager for a single direct array ufunc with an available SIMD +provider. Such an operation has no intermediate array for fusion to eliminate; +the clean 0.5.0 ADD comparison also favored eager. A decision is reported as +`single ufunc`. Scalar-provider dispatch and multi-operation thresholds remain +unchanged. Since the kernels changed after the reference run, numerical +threshold calibration requires new reference measurements; do not infer it +from the older measurements. The allocation benchmark's ADD case now measures +the selected eager path, while CHAIN exercises eligible fused execution. diff --git a/docs/benchmarks/M4-performance-status.md b/docs/benchmarks/M4-performance-status.md new file mode 100644 index 00000000..2fa5b8ae --- /dev/null +++ b/docs/benchmarks/M4-performance-status.md @@ -0,0 +1,31 @@ +# 0.6.0 fusion performance status + +The original reference measurements remain documented in +[0.5.0 results](M3-1-reference-results.md). They preceded the generated-loop +changes and cannot establish the performance of this revision. The user's +Ryzen 9 5950X reference machine is not available in this environment. + +A one-fork JDK 25 shared-machine smoke comparison of C-order float32 ADD and +CHAIN with 1,048,576 elements found substantially lower AUTO allocation after +generating the whole vector traversal loop; see the +[investigation](M4-fusion-investigation.md). Float64, mixed, C and F smoke +runs also completed. These short, separate runs can identify gross regressions, +but do not satisfy three-fork reference gates or establish stable speedups. +AUTO now selects eager for direct single-ufunc expressions with a SIMD provider; +it must be evaluated as selected, not relabeled as a fused ADD measurement. + +| Gate | Status | +|---|---| +| 1.5× large fused-output target | Open: rerun eligible multi-operation cases on the reference host. | +| 1.8× allocating regression guard | Open: rerun the specified allocating cases on the reference host. | +| Small and medium AUTO budgets | Open: the implementation changed after the old reference measurements. | +| Real dot/vecdot performance | Open: run the new contraction benchmark with backend provenance. | +| 0.6.0 worker scaling | Open: requires the parallel implementation and reference campaign. | + +Do not lower or waive the earlier targets. Record clean revision, patch status, +JDK, SIMD/native provider and thread settings with raw JSON and logs. + +The [0.6.0 scaling protocol](M4-reference-protocol.md) is ready, but the +reference campaign is outstanding. Crossover thresholds and multiworker +speedups are not established. A short shared-machine benchmark verified the +harness executes; it does not provide reference performance evidence. diff --git a/docs/benchmarks/M4-reference-protocol.md b/docs/benchmarks/M4-reference-protocol.md new file mode 100644 index 00000000..904abf7c --- /dev/null +++ b/docs/benchmarks/M4-reference-protocol.md @@ -0,0 +1,38 @@ +# 0.6.0 reference performance protocol + +The reference campaign requires the user's Ryzen 9 5950X and has not been run +in this environment. Do not treat shared-machine smoke checks as scaling or +release evidence. Run scalar and SIMD campaigns separately on an idle machine, +with JDK 25, Python 3 and a clean checkout of the actual candidate revision. +Record the output of `git status --short` and retain the entire result directory. +Each command requires a new empty output directory without spaces. + +```sh +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python3 numjvm-bench/run_060_reference.py \ + --mode scalar --workers 1,2,4,8,16 --output /tmp/numjvm-060-scalar +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python3 numjvm-bench/run_060_reference.py \ + --mode simd --workers 1,2,4,8,16 --output /tmp/numjvm-060-simd +``` + +The runner records revision, checkout status, JVM, CPU, memory, governor, SMT, +thread settings and the exact JMH command. It uses one JMH caller thread, +three forks, three one-second warmups and five one-second measurements. The +benchmark's own worker count is the `workers` parameter. Each mode runs 120 +parameter/method combinations and can take an hour or more. Use `--dry-run` +to inspect the command. A separate optional 32-worker campaign requires a +third output directory and `--workers 32`. + +Preserve `manifest.json`, `jmh.log` and `results.json`. The log includes +`NUMJVM_PARALLEL_DECISION`, provider and worker count from setup; verify the +intended paths were actually selected. `JAVA_TOOL_OPTIONS` is recorded, so +inspect it for private information before sharing. Native BLAS thread settings +are not changed by NumJVM; report any external overrides. Compare the parallel +and sequential methods only at matching parameters and report error bars, +small-array overhead, crossover and memory-bandwidth saturation. A single +JMH caller is distinct from the requested NumJVM worker count. + +The current shared-machine smoke run used one fork, no warmup and 100 ms +measurements solely to check that ADD, CHAIN and SUM benchmarks execute at one +and two workers. Its timings are not suitable for performance conclusions. +The 0.5.0 fusion, allocating and AUTO targets remain open pending the +reference reruns listed in [performance status](M4-performance-status.md). diff --git a/docs/design/parallel-cpu.md b/docs/design/parallel-cpu.md new file mode 100644 index 00000000..9c9b97ce --- /dev/null +++ b/docs/design/parallel-cpu.md @@ -0,0 +1,76 @@ +# Opt-in parallel CPU execution + +`NJ.parallel(workers)` creates an experimental, reusable CPU context. Existing +NumJVM calls remain single-worker. Always close the context: + +```java +try (var pool = NJ.parallel(4)) { + var a = NJ.ones(DType.FLOAT32, 65536); + var b = NJ.ones(DType.FLOAT32, 65536); + var out = NJ.zeros(DType.FLOAT32, 65536); + pool.apply(Ufuncs.ADD, a, b, out); + var total = pool.sum(out); + System.out.println(pool.lastDecision()); +} +``` + +The worker limit is 1–256. At most that many worker tasks are submitted per +operation. Calls on one context are serialized; separate contexts can run +concurrently. Workers are reused until `close()`, which waits for active work. +Closing from a context's worker is rejected. Nested work on the same context +runs on that worker without submitting further tasks. A worker cannot start an +operation on a different context or close any context; those calls are rejected +to avoid cross-pool deadlocks. Worker failures propagate +after all submitted tasks finish. An interrupted caller waits for workers, +restores its interrupt flag and receives an `IllegalStateException`. Unexpected +failures can leave preallocated output partially written; operations are not +transactional. Do not mutate inputs or outputs concurrently. + +## Supported operations and fallbacks + +The `apply` overloads accept a unary or binary ufunc and a preallocated output. +Parallel execution is currently limited to float32/64 C-contiguous arrays with +matching shapes and dtypes, at least 32,768 elements and more than one worker. +The supported unary operations are negative, absolute, square and sqrt; binary +operations are add, subtract, multiply and divide. Safe identical-view in-place +outputs can run in parallel. Shifted overlap, broadcasting, mixed ufunc dtypes, +other layouts and other operations use the existing sequential ufunc, including +its validation and temporary-copy behavior. The threshold is provisional. + +For expressions, call `expression.evalInto(pool, out, inputs...)`. Parallel +execution requires C-contiguous inputs of the same shape as a C-contiguous +output, float32/64 dtypes, safe output overlap and at least 32,768 elements. +Mixed float32/64 expressions are eligible. Otherwise the existing expression +executor handles the call. The JIT context and the pool must both remain open +until it returns. + +Arrays backed by a confined `NJ.scope()` cannot be handed to worker threads. +When any input or output is confined to the caller, the operation falls back +on that caller, without copying or changing the existing ownership rules. +Automatic memory and arrays in `NJ.scopeShared()` can be worker-accessible. +The caller must keep all scopes open until the call returns. Concurrently +closing a shared scope can make a worker fail and can leave output partial; +all submitted workers finish before the failure reaches the caller. + +`lastDecision()` reports the calling thread's latest selected strategy and +reason. `parallelRuns()` and `fallbackRuns()` count completed dispatches across +the context. A nested worker selection reports `sequential`. These diagnostics +are observational and should not be used to synchronize application logic. + +## Sum and reproducibility + +`pool.sum(input)` returns a scalar using the input's existing scope ownership +rules. Large C-contiguous float32/64 inputs use fixed 32,768-element chunks and +combine partials in increasing chunk order, regardless of worker count or +scheduling. Even one worker uses that same chunked algorithm for an eligible +input. The same data, dtype, algorithm and runtime therefore produce the same +result across worker counts. Floating-point grouping differs from the existing +sequential sum, so their low bits can differ. Each float32 partial and merge is +rounded to float32. The accumulation starts from positive zero, and IEEE NaN +and infinity behavior follows those additions. Small, non-C-contiguous, other +dtype and confined inputs use the existing sequential sum. + +NumJVM does not change OpenBLAS thread settings and does not call BLAS from the +parallel paths above. Performance depends on array size, memory bandwidth and +provider; the [reference campaign](../benchmarks/M4-reference-protocol.md) is +still needed to calibrate thresholds and quantify scaling. diff --git a/docs/milestones/M4-0.6-plan.md b/docs/milestones/M4-0.6-plan.md index 7d24d0e0..73babc1d 100644 --- a/docs/milestones/M4-0.6-plan.md +++ b/docs/milestones/M4-0.6-plan.md @@ -1,6 +1,7 @@ # Proposed NumJVM 0.6.0-alpha: fusion and opt-in parallel CPU execution -This is a proposal, not implemented functionality or a release commitment. +This is the implementation plan, not a release commitment. See the +[status](M4-0.6-status.md) for implemented scope and open evidence. The roadmap calls for worker pools and deterministic parallel reductions after an ownership design. The 0.5.0 reference investigation adds an earlier priority: correct the large fused-vector allocation and performance problem diff --git a/docs/milestones/M4-0.6-release-checklist.md b/docs/milestones/M4-0.6-release-checklist.md new file mode 100644 index 00000000..96fc31ad --- /dev/null +++ b/docs/milestones/M4-0.6-release-checklist.md @@ -0,0 +1,26 @@ +# 0.6.0-alpha candidate checklist + +This is an implementation candidate, not an instruction to tag or publish. + +- [x] One implementation commit per A01–A05, B01–B05, C01–C05, + D01–D05, E01–E04 and F01–F03; task dates and author are recorded in Git. +- [x] JDK 25 local `build`, fused-provider tests and the fused-expression + example completed before the final nested-context and interruption changes. + Targeted parallel context, stress and expression tests passed afterward. +- [x] NumPy 2.5.4 read all 4,245 generated oracle files; the new parallel + NumPy fixture check and reference-runner unit tests passed locally. +- [x] Separate classpath and module-path consumers exercised the new expression + and sum API using staged local artifacts, with and without the vector provider. +- [x] A JMH smoke run exercised parallel ADD, CHAIN and SUM. A separate + shared-machine experiment investigated fused allocation; neither is a + reference-machine performance claim. +- [ ] Verify build, oracle, CPU backend and quality checks on the actual PR head. +- [ ] Run the [0.6.0 reference campaign](../benchmarks/M4-reference-protocol.md) + on the designated Ryzen machine; report scaling and crossover. +- [ ] Re-run and assess the 0.5.0 fusion, allocating and AUTO gates against the + modified kernels. Complete dot/vecdot reference measurements. +- [ ] Review supported layouts, scope lifetime, nested contexts, native provider + settings and reproducibility limitations before deciding release readiness. + +Do not mark unmeasured targets as passed or silently lower them. In particular, +the E04 commit prepares the campaign; the external reference run remains open. diff --git a/docs/milestones/M4-0.6-status.md b/docs/milestones/M4-0.6-status.md new file mode 100644 index 00000000..a1c76048 --- /dev/null +++ b/docs/milestones/M4-0.6-status.md @@ -0,0 +1,24 @@ +# 0.6.0-alpha implementation candidate + +The candidate introduces generated whole-run homogeneous and mixed vector +loops, reduces small eager setup, and adds an opt-in bounded CPU context for +selected ufuncs, expressions and deterministic full-array sum. The supported +scope, fallback and numerical semantics are described in the +[parallel guide](../design/parallel-cpu.md). Unsupported layouts and confined +memory use the caller thread. Existing APIs remain single-threaded. + +Correctness coverage includes provider-backed fused graphs, layouts, float32 +rounding, overlap, worker interruption, shared-scope closure, and NumPy-generated +parallel ufunc and sum references. Verify CI at the actual PR head independently; +local test results do not establish remote checks. + +The reference Ryzen 9 5950X campaign has not been run on this revision. The +0.5.0 fusion, allocating and AUTO performance targets remain open, as do 0.6.0 +scaling and crossover calibration. A short shared-machine experiment observed +lower fused allocation, but it is not acceptance evidence. See the +[performance status](../benchmarks/M4-performance-status.md) and +[reference protocol](../benchmarks/M4-reference-protocol.md). This status does +not imply a tag, publication or release. + +See the [candidate checklist](M4-0.6-release-checklist.md) for validation and +outstanding release evidence. diff --git a/examples/fused-expressions/src/main/java/examples/FusedExpressions.java b/examples/fused-expressions/src/main/java/examples/FusedExpressions.java index 470cb0e6..16a4e401 100644 --- a/examples/fused-expressions/src/main/java/examples/FusedExpressions.java +++ b/examples/fused-expressions/src/main/java/examples/FusedExpressions.java @@ -26,5 +26,16 @@ public static void main(String[] args) { System.out.println(mode + " " + context.statistics()); } } + // Automatic memory is accessible to the context's workers. + try (var pool = NJ.parallel(4); var context = Jit.context(Jit.Mode.AUTO)) { + var expression = context.trace(2, i -> i.input(0).add(i.input(1)).square()); + var a = NJ.ones(DType.FLOAT32, 65536); + var b = NJ.ones(DType.FLOAT32, 65536); + var out = NJ.zeros(DType.FLOAT32, 65536); + expression.evalInto(pool, out, a, b); + if (out.getDouble(0) != 4) throw new AssertionError("parallel expression result"); + System.out.println("parallel sum=" + pool.sum(out).getDouble()); + System.out.println("execution=" + pool.lastDecision()); + } } } diff --git a/gradle.properties b/gradle.properties index 03af7a31..8966672b 100644 --- a/gradle.properties +++ b/gradle.properties @@ -1,5 +1,5 @@ group=io.github.stamich -version=0.5.0-alpha +version=0.6.0-alpha org.gradle.parallel=true org.gradle.caching=true org.gradle.jvmargs=-Xmx1g -Dfile.encoding=UTF-8 diff --git a/numjvm-bench/run_060_reference.py b/numjvm-bench/run_060_reference.py new file mode 100644 index 00000000..d6c34988 --- /dev/null +++ b/numjvm-bench/run_060_reference.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +"""Record and run a standalone 0.6.0 worker-scaling campaign.""" +import argparse +import datetime +import json +import os +from pathlib import Path +import platform +import shlex +import subprocess +import sys + +ROOT = Path(__file__).resolve().parent.parent + + +def command_output(command): + try: + return subprocess.check_output(command, cwd=ROOT, text=True, stderr=subprocess.STDOUT).strip() + except (OSError, subprocess.CalledProcessError) as error: + return f'unavailable: {error}' + + +def optional_file(path): + try: + return Path(path).read_text().strip() + except OSError: + return 'unavailable' + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--mode', required=True, choices=('scalar', 'simd')) + parser.add_argument('--output', required=True) + parser.add_argument('--workers', default='1,2,4,8,16') + parser.add_argument('--dry-run', action='store_true') + args = parser.parse_args() + output = Path(args.output).resolve() + if any(c.isspace() for c in str(output)): + parser.error('output path cannot contain whitespace') + try: + workers = [int(value) for value in args.workers.split(',')] + if not workers or any(value < 1 or value > 256 for value in workers): raise ValueError() + except ValueError: + parser.error('workers must be comma-separated integers from 1 to 256') + jmh = ('ParallelScalingBenchmark -p workers=' + ','.join(map(str, workers)) + + ' -t 1 -f 3 -wi 3 -w 1s -i 5 -r 1s -prof gc -rf json -rff ' + str(output / 'results.json')) + command = ['./gradlew', ':numjvm-bench:jmh'] + (['-PsimdBench'] if args.mode == 'simd' else []) + ['-PjmhArgs=' + jmh] + flags = '--add-modules=jdk.incubator.vector -Dnumjvm.simd=required' if args.mode == 'simd' else '-Dnumjvm.simd=off' + print('Command:', shlex.join(command), flush=True) + print('JAVA_TOOL_OPTIONS additions:', flags, flush=True) + if args.dry_run: return 0 + if output.exists() and any(output.iterdir()): parser.error('output directory must be empty') + output.mkdir(parents=True, exist_ok=True) + env = os.environ.copy() + env['JAVA_TOOL_OPTIONS'] = (env.get('JAVA_TOOL_OPTIONS', '') + ' ' + flags).strip() + env.setdefault('OMP_NUM_THREADS', '1') + env.setdefault('OPENBLAS_NUM_THREADS', '1') + status = command_output(['git', 'status', '--porcelain']) + manifest = {'recorded_at_utc': datetime.datetime.now(datetime.timezone.utc).isoformat(), + 'revision': command_output(['git', 'rev-parse', 'HEAD']), 'git_status': status, + 'system': platform.platform(), 'cpu': command_output(['lscpu']), + 'memory': optional_file('/proc/meminfo'), + 'cpu0_governor': optional_file('/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor'), + 'smt_active': optional_file('/sys/devices/system/cpu/smt/active'), + 'java': command_output(['java', '-version']), 'mode': args.mode, 'workers': workers, + 'jmh_threads': 1, 'java_tool_options': env['JAVA_TOOL_OPTIONS'], + 'thread_settings': {key: env.get(key) for key in ('OMP_NUM_THREADS', 'OPENBLAS_NUM_THREADS')}, + 'command': command, 'exit_code': None} + with (output / 'jmh.log').open('w') as log: + run = subprocess.run(command, cwd=ROOT, env=env, stdout=log, stderr=subprocess.STDOUT, check=False) + failed = run.returncode != 0 or not (output / 'results.json').is_file() or '' in (output / 'jmh.log').read_text() + manifest['exit_code'] = run.returncode or int(failed) + (output / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n') + if failed: print('Benchmark failed; inspect', output / 'jmh.log', file=sys.stderr) + return manifest['exit_code'] + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/numjvm-bench/src/jmh/java/org/numjvm/bench/ContractionBenchmark.java b/numjvm-bench/src/jmh/java/org/numjvm/bench/ContractionBenchmark.java new file mode 100644 index 00000000..ec050f61 --- /dev/null +++ b/numjvm-bench/src/jmh/java/org/numjvm/bench/ContractionBenchmark.java @@ -0,0 +1,39 @@ +package org.numjvm.bench; + +import java.util.concurrent.TimeUnit; +import org.numjvm.NJ; +import org.numjvm.BackendInfo; +import org.numjvm.dtype.DType; +import org.numjvm.linalg.Linalg; +import org.numjvm.ndarray.NDArray; +import org.openjdk.jmh.annotations.*; + +/** End-to-end vector contractions; setup reports the observed backend. */ +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MICROSECONDS) +@Warmup(iterations = 3, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(3) +@State(Scope.Thread) +public class ContractionBenchmark { + @Param({"256", "65536", "1048576"}) public int size; + @Param({"FLOAT32", "FLOAT64"}) public DType dtype; + private NDArray a, b, out; + + @Setup public void setup() { + a = NJ.ones(dtype, size); + b = NJ.ones(dtype, size); + out = NJ.zeros(dtype); + Linalg.dot(a, b, out); + System.out.println("NUMJVM_DOT_BACKEND=" + Linalg.lastBackend()); + Linalg.vecdot(a, b, out); + System.out.println("NUMJVM_VECDOT_BACKEND=" + Linalg.lastBackend()); + System.out.println("NUMJVM_PROVIDER=" + BackendInfo.simdProvider()); + if (out.getDouble() != size) throw new AssertionError("invalid contraction result"); + } + + @Benchmark public NDArray dotAllocating() { return Linalg.dot(a, b); } + @Benchmark public NDArray dotOut() { return Linalg.dot(a, b, out); } + @Benchmark public NDArray vecdotAllocating() { return Linalg.vecdot(a, b); } + @Benchmark public NDArray vecdotOut() { return Linalg.vecdot(a, b, out); } +} diff --git a/numjvm-bench/src/jmh/java/org/numjvm/bench/ExpressionMatrixBenchmark.java b/numjvm-bench/src/jmh/java/org/numjvm/bench/ExpressionMatrixBenchmark.java index c8b6eb54..35502cb8 100644 --- a/numjvm-bench/src/jmh/java/org/numjvm/bench/ExpressionMatrixBenchmark.java +++ b/numjvm-bench/src/jmh/java/org/numjvm/bench/ExpressionMatrixBenchmark.java @@ -42,6 +42,13 @@ public class ExpressionMatrixBenchmark { if (layout == Layouts.Layout.F) out = out.copy(Order.F); temp1 = NJ.zeros(resultType, Layouts.rows(size), Layouts.cols(size)); temp2 = NJ.zeros(resultType, Layouts.rows(size), Layouts.cols(size)); + // Match intermediate traversal to the Fortran-ordered inputs and output. + if (layout == Layouts.Layout.F) { + temp1 = temp1.copy(Order.F); + temp2 = temp2.copy(Order.F); + if (!out.isFContiguous() || !temp1.isFContiguous() || !temp2.isFContiguous()) + throw new IllegalStateException("Fortran benchmark buffers must be F-contiguous"); + } auto = Jit.context(Jit.Mode.AUTO); block = Jit.context(Jit.Mode.BLOCK); generated = Jit.context(Jit.Mode.REQUIRED); scheduled = Jit.context(Jit.Mode.EAGER, 64, 65536, 64L * 1024 * 1024); autoPlan = trace(auto); blockPlan = trace(block); generatedPlan = trace(generated); scheduledPlan = trace(scheduled); diff --git a/numjvm-bench/src/jmh/java/org/numjvm/bench/FusedAllocationBenchmark.java b/numjvm-bench/src/jmh/java/org/numjvm/bench/FusedAllocationBenchmark.java new file mode 100644 index 00000000..934e4633 --- /dev/null +++ b/numjvm-bench/src/jmh/java/org/numjvm/bench/FusedAllocationBenchmark.java @@ -0,0 +1,29 @@ +package org.numjvm.bench; + +import java.util.concurrent.TimeUnit; +import org.numjvm.ndarray.NDArray; +import org.openjdk.jmh.annotations.*; + +/** Use JMH's gc profiler to compare warm fused and eager output allocation. */ +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MICROSECONDS) +@Warmup(iterations = 3, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(3) +@State(Scope.Thread) +public class FusedAllocationBenchmark { + @Param({"65536", "1048576"}) public int size; + @Param({"FLOAT32", "FLOAT64", "MIXED"}) public String types; + @Param({"ADD", "CHAIN"}) public String chain; + @Param({"C", "F"}) public Layouts.Layout layout; + private ExpressionMatrixBenchmark matrix; + + @Setup public void setup() { + matrix = new ExpressionMatrixBenchmark(); + matrix.size = size; matrix.types = types; matrix.chain = chain; matrix.layout = layout; + matrix.setup(); + } + @Benchmark public NDArray autoOut() { return matrix.autoOut(); } + @Benchmark public NDArray eagerOut() { return matrix.eagerOut(); } + @TearDown public void close() { matrix.close(); } +} diff --git a/numjvm-bench/src/jmh/java/org/numjvm/bench/ParallelScalingBenchmark.java b/numjvm-bench/src/jmh/java/org/numjvm/bench/ParallelScalingBenchmark.java new file mode 100644 index 00000000..273ae9de --- /dev/null +++ b/numjvm-bench/src/jmh/java/org/numjvm/bench/ParallelScalingBenchmark.java @@ -0,0 +1,64 @@ +package org.numjvm.bench; + +import java.util.concurrent.TimeUnit; +import org.numjvm.BackendInfo; +import org.numjvm.NJ; +import org.numjvm.ParallelContext; +import org.numjvm.dtype.DType; +import org.numjvm.jit.Expression; +import org.numjvm.jit.Jit; +import org.numjvm.ndarray.NDArray; +import org.numjvm.ufunc.Reductions; +import org.numjvm.ufunc.Ufuncs; +import org.openjdk.jmh.annotations.*; + +/** End-to-end warm execution, with a single JMH thread and explicit worker count. */ +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MICROSECONDS) +@Warmup(iterations = 3, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(3) +@State(Scope.Thread) +public class ParallelScalingBenchmark { + @Param({"1", "2", "4", "8", "16"}) public int workers; + @Param({"65536", "1048576"}) public int size; + @Param({"FLOAT32", "FLOAT64"}) public DType dtype; + @Param({"ADD", "CHAIN", "SUM"}) public String operation; + private ParallelContext pool; + private Jit.Context jit; + private Expression expression; + private NDArray a, b, out, temporary; + + @Setup public void setup() { + pool = NJ.parallel(workers); + jit = Jit.context(Jit.Mode.AUTO); + expression = jit.trace(2, i -> i.input(0).add(i.input(1)).square()); + a = NJ.ones(dtype, size); b = NJ.ones(dtype, size); + out = NJ.zeros(dtype, size); temporary = NJ.zeros(dtype, size); + NDArray actual = parallel(); + double expected = operation.equals("SUM") ? size : operation.equals("CHAIN") ? 4 : 2; + if (actual.getDouble(operation.equals("SUM") ? new long[]{} : new long[]{size - 1}) != expected) + throw new AssertionError("incorrect benchmark result"); + System.out.println("NUMJVM_PARALLEL_DECISION=" + pool.lastDecision()); + System.out.println("NUMJVM_PROVIDER=" + BackendInfo.simdProvider()); + System.out.println("NUMJVM_WORKERS=" + workers); + } + + @Benchmark public NDArray parallel() { + return switch (operation) { + case "ADD" -> pool.apply(Ufuncs.ADD, a, b, out); + case "CHAIN" -> expression.evalInto(pool, out, a, b); + case "SUM" -> pool.sum(a); + default -> throw new IllegalArgumentException(operation); + }; + } + @Benchmark public NDArray sequential() { + return switch (operation) { + case "ADD" -> Ufuncs.ADD.apply(a, b, out); + case "CHAIN" -> Ufuncs.SQUARE.apply(Ufuncs.ADD.apply(a, b, temporary), out); + case "SUM" -> Reductions.sum(a); + default -> throw new IllegalArgumentException(operation); + }; + } + @TearDown public void close() { pool.close(); jit.close(); } +} diff --git a/numjvm-bench/test_run_060_reference.py b/numjvm-bench/test_run_060_reference.py new file mode 100644 index 00000000..a0c13dab --- /dev/null +++ b/numjvm-bench/test_run_060_reference.py @@ -0,0 +1,32 @@ +import subprocess +import sys +from pathlib import Path +import tempfile +import unittest + +SCRIPT = Path(__file__).with_name('run_060_reference.py') + + +class ParallelRunnerTest(unittest.TestCase): + def test_dry_run_has_explicit_workers_and_does_not_create_output(self): + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) / 'results' + run = subprocess.run([sys.executable, str(SCRIPT), '--mode', 'simd', '--workers', '1,2', + '--output', str(output), '--dry-run'], text=True, capture_output=True) + self.assertEqual(0, run.returncode, run.stderr) + self.assertIn('-p workers=1,2', run.stdout) + self.assertIn('-PsimdBench', run.stdout) + self.assertFalse(output.exists()) + + def test_rejects_invalid_worker_counts_and_paths(self): + for workers in ('0', '257', 'two'): + run = subprocess.run([sys.executable, str(SCRIPT), '--mode', 'scalar', '--workers', workers, + '--output', '/tmp/example', '--dry-run'], text=True, capture_output=True) + self.assertNotEqual(0, run.returncode) + run = subprocess.run([sys.executable, str(SCRIPT), '--mode', 'scalar', '--output', '/tmp/a space', + '--dry-run'], text=True, capture_output=True) + self.assertNotEqual(0, run.returncode) + + +if __name__ == '__main__': + unittest.main() diff --git a/numjvm-core/src/main/java/org/numjvm/ChunkPlan.java b/numjvm-core/src/main/java/org/numjvm/ChunkPlan.java new file mode 100644 index 00000000..9b589d48 --- /dev/null +++ b/numjvm-core/src/main/java/org/numjvm/ChunkPlan.java @@ -0,0 +1,20 @@ +package org.numjvm; + +/** Fixed logical chunks, independent of worker count and scheduling. */ +final class ChunkPlan { + final long size, chunkSize, count; + ChunkPlan(long size, long chunkSize) { + if (size < 0 || chunkSize <= 0) throw new IllegalArgumentException("invalid chunk dimensions"); + this.size = size; + this.chunkSize = chunkSize; + this.count = size == 0 ? 0 : 1 + (size - 1) / chunkSize; + } + long start(long index) { + if (index < 0 || index >= count) throw new IndexOutOfBoundsException(); + return index * chunkSize; + } + long end(long index) { + long start = start(index); + return start + Math.min(chunkSize, size - start); + } +} diff --git a/numjvm-core/src/main/java/org/numjvm/NJ.java b/numjvm-core/src/main/java/org/numjvm/NJ.java index 8afeebc8..4394627d 100644 --- a/numjvm-core/src/main/java/org/numjvm/NJ.java +++ b/numjvm-core/src/main/java/org/numjvm/NJ.java @@ -52,6 +52,16 @@ public static String version() { return VERSION; } + /** + * Creates an opt-in, reusable CPU execution context. + * @param workers maximum worker count, from 1 through 256 + * @return a context that must be closed + */ + @Experimental + public static ParallelContext parallel(int workers) { + return new ParallelContext(workers); + } + private static HybridAllocator automatic() { return HybridAllocator.defaultInstance(); } diff --git a/numjvm-core/src/main/java/org/numjvm/ParallelContext.java b/numjvm-core/src/main/java/org/numjvm/ParallelContext.java new file mode 100644 index 00000000..64729d35 --- /dev/null +++ b/numjvm-core/src/main/java/org/numjvm/ParallelContext.java @@ -0,0 +1,353 @@ +package org.numjvm; + +import java.util.ArrayList; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.LongAdder; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.numjvm.internal.ExpressionSemantics; +import org.numjvm.internal.SharedSecrets; +import org.numjvm.ndarray.NDArray; +import org.numjvm.ndarray.Slice; +import org.numjvm.ndarray.Shape; +import org.numjvm.ndarray.Order; +import java.lang.foreign.ValueLayout; +import org.numjvm.dtype.DType; +import org.numjvm.ufunc.BinaryUfunc; +import org.numjvm.ufunc.UnaryUfunc; +import org.numjvm.ufunc.Ufuncs; +import org.numjvm.ufunc.Reductions; + +/** + * An opt-in, bounded CPU execution context. Close it after its operations finish. + * Existing NumJVM operations continue to execute on the calling thread. + */ +@Experimental +public final class ParallelContext implements AutoCloseable { + private static final long CHUNK_SIZE = 32768; + static { + SharedSecrets.setParallelAccess((context, arrays, size, eligible, task, fallback) -> { + if (!eligible || context.workers == 1 || size < CHUNK_SIZE) { + context.sequential(() -> { fallback.run(); return null; }, "small or unsupported expression"); + } else { + context.runOrElse(arrays, new ChunkPlan(size, CHUNK_SIZE), + (index, start, end) -> task.run(start, end), fallback); + } + }); + } + + + private static final AtomicInteger CONTEXT_IDS = new AtomicInteger(); + // A distinct thread identifies confined FFM segments without starting a worker. + private static final Thread ACCESS_PROBE = new Thread(() -> { }, "numjvm-access-probe"); + private static final ThreadLocal CURRENT = new ThreadLocal<>(); + private final Object operationLock = new Object(); + private final LongAdder parallelRuns = new LongAdder(); + private final LongAdder fallbackRuns = new LongAdder(); + private final ThreadLocal decision = ThreadLocal.withInitial(() -> new Decision("none", "no operation")); + private final int workers; + private final ThreadPoolExecutor executor; + private volatile boolean closed; + + ParallelContext(int workers) { + if (workers < 1 || workers > 256) throw new IllegalArgumentException("workers must be between 1 and 256"); + this.workers = workers; + int id = CONTEXT_IDS.incrementAndGet(); + var sequence = new AtomicInteger(); + executor = new ThreadPoolExecutor(workers, workers, 0, TimeUnit.MILLISECONDS, + new ArrayBlockingQueue<>(workers), task -> { + var thread = new Thread(task, "numjvm-" + id + "-worker-" + sequence.incrementAndGet()); + thread.setDaemon(true); + return thread; + }, new ThreadPoolExecutor.AbortPolicy()); + } + + /** Returns the configured maximum number of workers. + * @return worker limit + */ + public int workers() { return workers; } + + /** Whether this context has been closed. + * @return true after closure starts + */ + public boolean isClosed() { return closed; } + + /** The strategy most recently selected on the calling thread. + * @param strategy execution strategy + * @param reason selection reason + */ + public record Decision(String strategy, String reason) { } + + /** Returns the calling thread's last selection. + * @return last decision, initially {@code none} + */ + public Decision lastDecision() { return decision.get(); } + + /** Returns the number of completed parallel dispatches. + * @return completed parallel dispatches + */ + public long parallelRuns() { return parallelRuns.sum(); } + + /** Returns the number of completed single-worker fallbacks. + * @return completed fallbacks + */ + public long fallbackRuns() { return fallbackRuns.sum(); } + + private void dispatch(NDArray[] arrays, ChunkPlan plan, ChunkAction action, Runnable fallback) { + if (workerAccessible(arrays)) { + boolean nested = CURRENT.get() == this; + runChunks(plan, action); + if (nested) { + fallbackRuns.increment(); + decision.set(new Decision("sequential", "nested worker")); + } else { + parallelRuns.increment(); + decision.set(new Decision("parallel", "worker-accessible memory")); + } + } else { + fallback.run(); + fallbackRuns.increment(); + decision.set(new Decision("sequential", "confined memory")); + } + } + + private static boolean supported(DType dtype) { + return dtype == DType.FLOAT32 || dtype == DType.FLOAT64; + } + + private boolean compatible(NDArray input, NDArray output) { + return workers > 1 && output.size() >= CHUNK_SIZE && input.isCContiguous() && output.isCContiguous() + && input.shape().equals(output.shape()) && input.dtype() == output.dtype() + && supported(input.dtype()) && output.isWriteable() && safeOverlap(input, output); + } + + private static boolean safeOverlap(NDArray input, NDArray output) { + if (!input.sharesMemoryWith(output)) return true; + var access = SharedSecrets.arrayAccess(); + return access.buffer(input).segment() == access.buffer(output).segment() + && access.offset(input) == access.offset(output) + && java.util.Arrays.equals(input.strides().values(), output.strides().values()); + } + + /** Applies a supported unary ufunc into an output, using workers when eligible. + * @param op unary ufunc + * @param input input array + * @param output destination + * @return output + */ + public NDArray apply(UnaryUfunc op, NDArray input, NDArray output) { + java.util.Objects.requireNonNull(op, "op"); + ExpressionSemantics.accessible(input, output); + boolean eligible = (op == Ufuncs.NEGATIVE || op == Ufuncs.ABSOLUTE || op == Ufuncs.SQUARE || op == Ufuncs.SQRT) + && compatible(input, output); + if (!eligible) return sequential(() -> op.apply(input, output), "small or unsupported operation"); + var source = input.reshapeView(input.size()); + var destination = output.reshapeView(output.size()); + runOrElse(new NDArray[]{input, output}, new ChunkPlan(output.size(), CHUNK_SIZE), + (i, start, end) -> op.apply(source.slice(Slice.range(start, end)), destination.slice(Slice.range(start, end))), + () -> op.apply(input, output)); + return output; + } + + /** Applies a supported binary ufunc into an output, using workers when eligible. + * @param op binary ufunc + * @param left first input + * @param right second input + * @param output destination + * @return output + */ + public NDArray apply(BinaryUfunc op, NDArray left, NDArray right, NDArray output) { + java.util.Objects.requireNonNull(op, "op"); + ExpressionSemantics.accessible(left, right, output); + boolean eligible = (op == Ufuncs.ADD || op == Ufuncs.SUBTRACT || op == Ufuncs.MULTIPLY || op == Ufuncs.DIVIDE) + && compatible(left, output) && compatible(right, output); + if (!eligible) return sequential(() -> op.apply(left, right, output), "small or unsupported operation"); + var a = left.reshapeView(left.size()); + var b = right.reshapeView(right.size()); + var destination = output.reshapeView(output.size()); + runOrElse(new NDArray[]{left, right, output}, new ChunkPlan(output.size(), CHUNK_SIZE), + (i, start, end) -> op.apply(a.slice(Slice.range(start, end)), b.slice(Slice.range(start, end)), + destination.slice(Slice.range(start, end))), + () -> op.apply(left, right, output)); + return output; + } + + /** + * Sums a contiguous real array using fixed logical chunks and a fixed merge order. + * The parallel result is reproducible across worker counts for the same inputs, + * dtype and runtime, but can differ from a sequential floating-point sum. + * Unsupported or confined arrays use the existing sequential reduction. + * @param input array to reduce + * @return scalar result with the existing sum dtype and ownership rules + */ + public NDArray sum(NDArray input) { + rejectOtherContext(); + if (CURRENT.get() == this) return sumNow(input); + synchronized (operationLock) { return sumNow(input); } + } + + private NDArray sumNow(NDArray input) { + if (closed) throw new IllegalStateException("parallel context is closed"); + ExpressionSemantics.accessible(input); + var plan = new ChunkPlan(input.size(), CHUNK_SIZE); + if (!input.isCContiguous() || !supported(input.dtype()) || input.size() < CHUNK_SIZE + || plan.count > Integer.MAX_VALUE) + return sequential(() -> Reductions.sum(input), "small or unsupported reduction"); + var result = SharedSecrets.arrayAccess().allocate(ExpressionSemantics.allocator(input), + input.dtype(), Shape.of(), Order.C); + var segment = SharedSecrets.arrayAccess().buffer(input).segment(); + long base = SharedSecrets.arrayAccess().offset(input); + // Each index identifies a fixed input interval regardless of worker count. + var partials = new double[(int) plan.count]; + var computed = new java.util.concurrent.atomic.AtomicBoolean(); + runOrElse(new NDArray[]{input, result}, plan, (chunk, start, end) -> { + if (input.dtype() == DType.FLOAT32) { + float value = 0; + for (long i = start; i < end; i++) value += segment.get(ValueLayout.JAVA_FLOAT_UNALIGNED, base + i * Float.BYTES); + partials[(int) chunk] = value; + } else { + double value = 0; + for (long i = start; i < end; i++) value += segment.get(ValueLayout.JAVA_DOUBLE_UNALIGNED, base + i * Double.BYTES); + partials[(int) chunk] = value; + } + computed.set(true); + }, () -> { + var sequentialResult = Reductions.sum(input); + result.setDouble(sequentialResult.getDouble()); + }); + if (computed.get()) { + if (input.dtype() == DType.FLOAT32) { + float value = 0; + for (double part : partials) value += (float) part; + result.setDouble(value); + } else { + double value = 0; + for (double part : partials) value += part; + result.setDouble(value); + } + } + return result; + } + + private NDArray sequential(java.util.function.Supplier action, String reason) { + rejectOtherContext(); + if (CURRENT.get() == this) return sequentialNow(action, reason); + synchronized (operationLock) { return sequentialNow(action, reason); } + } + + private NDArray sequentialNow(java.util.function.Supplier action, String reason) { + if (closed) throw new IllegalStateException("parallel context is closed"); + NDArray result = action.get(); + fallbackRuns.increment(); + decision.set(new Decision("sequential", reason)); + return result; + } + + @FunctionalInterface + interface ChunkAction { void run(long index, long start, long end); } + + boolean workerAccessible(NDArray... arrays) { + ExpressionSemantics.accessible(arrays); + for (var array : arrays) { + var buffer = SharedSecrets.arrayAccess().buffer(array); + buffer.checkAlive(); + if (!buffer.segment().isAccessibleBy(ACCESS_PROBE)) return false; + } + return true; + } + + void runOrElse(NDArray[] arrays, ChunkPlan plan, ChunkAction action, Runnable fallback) { + rejectOtherContext(); + if (CURRENT.get() == this) { + if (closed) throw new IllegalStateException("parallel context is closed"); + dispatch(arrays, plan, action, fallback); + return; + } + synchronized (operationLock) { + if (closed) throw new IllegalStateException("parallel context is closed"); + dispatch(arrays, plan, action, fallback); + } + } + + /** Submits at most one task per worker; tasks claim fixed chunks dynamically. */ + void runChunks(ChunkPlan plan, ChunkAction action) { + rejectOtherContext(); + if (CURRENT.get() == this) { + if (closed) throw new IllegalStateException("parallel context is closed"); + for (long i = 0; i < plan.count; i++) action.run(i, plan.start(i), plan.end(i)); + return; + } + synchronized (operationLock) { + var futures = new ArrayList>(); + var next = new AtomicLong(); + Throwable submissionFailure = null; + synchronized (this) { + if (closed) throw new IllegalStateException("parallel context is closed"); + int tasks = (int) Math.min(plan.count, workers); + for (int worker = 0; worker < tasks; worker++) try { + futures.add(executor.submit(() -> { + CURRENT.set(this); + try { + for (;;) { + long index = next.getAndIncrement(); + if (index < 0 || index >= plan.count) break; + action.run(index, plan.start(index), plan.end(index)); + } + } finally { decision.remove(); CURRENT.remove(); } + })); + } catch (RuntimeException | Error ex) { + submissionFailure = ex; + break; + } + } + boolean interrupted = false; + Throwable failure = submissionFailure; + for (var future : futures) { + for (;;) { + // Keep the interrupt recorded while allowing every worker to be joined. + if (interrupted) Thread.interrupted(); + try { future.get(); break; } + catch (InterruptedException ex) { interrupted = true; Thread.currentThread().interrupt(); } + catch (ExecutionException ex) { if (failure == null) failure = ex.getCause(); break; } + } + } + if (interrupted) Thread.currentThread().interrupt(); + if (failure instanceof Error error) throw error; + if (failure instanceof RuntimeException error) throw error; + if (failure != null) throw new IllegalStateException("parallel task failed", failure); + if (interrupted) throw new IllegalStateException("parallel operation interrupted"); + } + } + + private void rejectOtherContext() { + var current = CURRENT.get(); + if (current != null && current != this) + throw new IllegalStateException("nested execution across parallel contexts is unsupported"); + } + + /** Stops accepting work and waits for all submitted workers to terminate. */ + @Override public void close() { + if (CURRENT.get() != null) throw new IllegalStateException("cannot close a context from a NumJVM worker"); + synchronized (operationLock) { + synchronized (this) { + if (!closed) { + closed = true; + executor.shutdown(); + } + } + } + boolean interrupted = false; + for (;;) { + if (interrupted) Thread.interrupted(); + try { + if (executor.awaitTermination(1, TimeUnit.DAYS)) break; + } catch (InterruptedException ex) { interrupted = true; Thread.currentThread().interrupt(); } + } + decision.remove(); + if (interrupted) Thread.currentThread().interrupt(); + } +} diff --git a/numjvm-core/src/main/java/org/numjvm/internal/SharedSecrets.java b/numjvm-core/src/main/java/org/numjvm/internal/SharedSecrets.java index 40ed3dde..353be1b5 100644 --- a/numjvm-core/src/main/java/org/numjvm/internal/SharedSecrets.java +++ b/numjvm-core/src/main/java/org/numjvm/internal/SharedSecrets.java @@ -1,5 +1,6 @@ package org.numjvm.internal; +import org.numjvm.ParallelContext; import org.numjvm.dtype.DType; import org.numjvm.internal.memory.Buffer; import org.numjvm.internal.memory.MemoryAllocator; @@ -111,6 +112,35 @@ public interface ArrayAccess { long offset(NDArray array); } + /** Narrow bridge for the optional expression module's validated chunk execution. */ + public interface ParallelAccess { + /** Executes disjoint chunks or the caller-thread fallback. */ + void execute(ParallelContext context, NDArray[] arrays, long size, boolean eligible, + RangeTask task, Runnable fallback); + } + + /** An internal logical element range. */ + @FunctionalInterface public interface RangeTask { + /** Executes a half-open range. */ + void run(long start, long end); + } + + private static volatile ParallelAccess parallelAccess; + + /** Registers the parallel context bridge. */ + public static void setParallelAccess(ParallelAccess access) { parallelAccess = access; } + + /** Returns the initialized bridge. */ + public static ParallelAccess parallelAccess() { + ParallelAccess access = parallelAccess; + if (access == null) { + ensureInitialized(ParallelContext.class); + access = parallelAccess; + if (access == null) throw notRegistered(ParallelContext.class); + } + return access; + } + private static volatile ScopeAccess scopeAccess; private static volatile ArrayAccess arrayAccess; diff --git a/numjvm-core/src/test/java/org/numjvm/ParallelContextTest.java b/numjvm-core/src/test/java/org/numjvm/ParallelContextTest.java new file mode 100644 index 00000000..5da2d643 --- /dev/null +++ b/numjvm-core/src/test/java/org/numjvm/ParallelContextTest.java @@ -0,0 +1,236 @@ +package org.numjvm; + +import org.junit.jupiter.api.Test; +import static org.junit.jupiter.api.Assertions.*; + +class ParallelContextTest { + @Test void identicalViewsAreParallelAndShiftedOverlapUsesExistingFallback() { + try (var context = NJ.parallel(3)) { + var input = NJ.ones(org.numjvm.dtype.DType.FLOAT32, 65536); + var identical = input.reshapeView(65536); + context.apply(org.numjvm.ufunc.Ufuncs.ADD, input, identical, input); + assertEquals("parallel", context.lastDecision().strategy()); + assertEquals(2.0, input.getDouble(0)); + assertEquals(2.0, input.getDouble(65535)); + + var backing = NJ.ones(org.numjvm.dtype.DType.FLOAT32, 65537); + var left = backing.slice(org.numjvm.ndarray.Slice.range(0, 65536)); + var shifted = backing.slice(org.numjvm.ndarray.Slice.range(1, 65537)); + context.apply(org.numjvm.ufunc.Ufuncs.ADD, left, left, shifted); + assertEquals("sequential", context.lastDecision().strategy()); + assertEquals(1.0, backing.getDouble(0)); + assertEquals(2.0, backing.getDouble(100)); + } + } + + @Test void fixedChunksCoverEveryElementOnce() { + var counts = new java.util.concurrent.atomic.AtomicIntegerArray(103); + try (var context = NJ.parallel(3)) { + context.runChunks(new ChunkPlan(103, 16), (chunk, start, end) -> { + for (long i = start; i < end; i++) counts.incrementAndGet((int) i); + }); + } + for (int i = 0; i < 103; i++) assertEquals(1, counts.get(i)); + var plan = new ChunkPlan(Long.MAX_VALUE, 8192); + assertEquals(Long.MAX_VALUE, plan.end(plan.count - 1)); + assertEquals(0, new ChunkPlan(0, 3).count); + } + + @Test void confinedMemoryFallsBackToCallingThread() { + try (var context = NJ.parallel(2); var scope = NJ.scope(); var shared = NJ.scopeShared()) { + var confined = scope.ones(org.numjvm.dtype.DType.FLOAT64, 4); + var shareable = shared.ones(org.numjvm.dtype.DType.FLOAT64, 4); + assertFalse(context.workerAccessible(confined)); + assertTrue(context.workerAccessible(shareable)); + assertTrue(context.workerAccessible(NJ.ones(4))); + var caller = Thread.currentThread(); + context.runOrElse(new org.numjvm.ndarray.NDArray[]{confined}, new ChunkPlan(4, 2), + (i, start, end) -> fail("confined memory reached a worker"), + () -> { assertSame(caller, Thread.currentThread()); assertEquals(1.0, confined.getDouble(0)); }); + assertEquals("sequential", context.lastDecision().strategy()); + assertEquals(1, context.fallbackRuns()); + context.runOrElse(new org.numjvm.ndarray.NDArray[]{shareable}, new ChunkPlan(4, 2), + (i, start, end) -> assertEquals(1.0, shareable.getDouble(start)), () -> fail("unexpected fallback")); + assertEquals("parallel", context.lastDecision().strategy()); + assertEquals(1, context.parallelRuns()); + } + } + + @Test void nestedWorkDoesNotDeadlockAndWorkerCannotCloseContext() { + try (var context = NJ.parallel(2)) { + var seen = new java.util.concurrent.atomic.AtomicInteger(); + var input = NJ.ones(org.numjvm.dtype.DType.FLOAT32, 65536); + context.runChunks(new ChunkPlan(2, 1), (index, start, end) -> { + context.runChunks(new ChunkPlan(3, 1), (j, a, b) -> seen.incrementAndGet()); + assertEquals(65536.0, context.sum(input).getDouble()); + assertEquals("sequential", context.lastDecision().strategy()); + assertEquals("nested worker", context.lastDecision().reason()); + assertThrows(IllegalStateException.class, context::close); + }); + assertEquals(6, seen.get()); + } + } + + @Test void failureWaitsForOtherWorkersBeforeReturning() throws Exception { + try (var context = NJ.parallel(2)) { + var started = new java.util.concurrent.CountDownLatch(1); + var release = new java.util.concurrent.CountDownLatch(1); + var result = new java.util.concurrent.atomic.AtomicReference(); + Thread caller = new Thread(() -> { + try { + context.runChunks(new ChunkPlan(2, 1), (index, start, end) -> { + if (index == 0) { + try { started.await(); } catch (InterruptedException e) { throw new AssertionError(e); } + throw new IllegalStateException("expected"); + } + started.countDown(); + try { release.await(); } catch (InterruptedException e) { throw new AssertionError(e); } + }); + } catch (Throwable ex) { result.set(ex); } + }); + caller.start(); + assertTrue(started.await(5, java.util.concurrent.TimeUnit.SECONDS)); + caller.join(100); + assertTrue(caller.isAlive()); + release.countDown(); + caller.join(5000); + assertFalse(caller.isAlive()); + assertEquals("expected", result.get().getMessage()); + } + } + + @Test void interruptedCallerWaitsForWorkerAndRestoresInterrupt() throws Exception { + try (var context = NJ.parallel(1)) { + var started = new java.util.concurrent.CountDownLatch(1); + var release = new java.util.concurrent.CountDownLatch(1); + var interrupted = new java.util.concurrent.atomic.AtomicBoolean(); + var result = new java.util.concurrent.atomic.AtomicReference(); + Thread caller = new Thread(() -> { + try { + context.runChunks(new ChunkPlan(1, 1), (index, start, end) -> { + started.countDown(); + try { release.await(); } catch (InterruptedException ex) { throw new AssertionError(ex); } + }); + } catch (Throwable ex) { + result.set(ex); + interrupted.set(Thread.currentThread().isInterrupted()); + } + }); + caller.start(); + assertTrue(started.await(5, java.util.concurrent.TimeUnit.SECONDS)); + caller.interrupt(); + caller.join(100); + assertTrue(caller.isAlive()); + release.countDown(); + caller.join(5000); + assertFalse(caller.isAlive()); + assertInstanceOf(IllegalStateException.class, result.get()); + assertTrue(interrupted.get()); + } + } + + @Test void appliesContiguousFloatUfuncsInParallel() { + for (var dtype : new org.numjvm.dtype.DType[]{org.numjvm.dtype.DType.FLOAT32, org.numjvm.dtype.DType.FLOAT64}) { + try (var context = NJ.parallel(3)) { + var left = NJ.ones(dtype, 256, 256); + var right = NJ.ones(dtype, 256, 256); + var output = NJ.zeros(dtype, 256, 256); + assertSame(output, context.apply(org.numjvm.ufunc.Ufuncs.ADD, left, right, output)); + assertEquals("parallel", context.lastDecision().strategy()); + assertEquals(2.0, output.getDouble(255, 255)); + context.apply(org.numjvm.ufunc.Ufuncs.SQUARE, output, left); + assertEquals("parallel", context.lastDecision().strategy()); + assertEquals(4.0, left.getDouble(0, 0)); + } + } + } + + @Test void sumUsesIdenticalFixedChunksAcrossWorkerCounts() { + float[] data = new float[100003]; + for (int i = 0; i < data.length; i++) data[i] = (i % 17 - 8) / 7.0f; + var input = NJ.array(data); + int expected = 0; + for (int workers : new int[]{1, 2, 3, 7}) { + try (var context = NJ.parallel(workers)) { + int actual = Float.floatToRawIntBits((float) context.sum(input).getDouble()); + if (workers == 1) expected = actual; + else assertEquals(expected, actual); + assertEquals("parallel", context.lastDecision().strategy()); + } + } + var doubleInput = input.astype(org.numjvm.dtype.DType.FLOAT64); + long expectedDouble = 0; + for (int workers : new int[]{1, 2, 4}) { + try (var context = NJ.parallel(workers)) { + long actual = Double.doubleToRawLongBits(context.sum(doubleInput).getDouble()); + if (workers == 1) expectedDouble = actual; + else assertEquals(expectedDouble, actual); + } + } + } + + @Test void sumPreservesEmptyAndSpecialValueBehavior() { + try (var context = NJ.parallel(2); var scope = NJ.scope()) { + assertEquals(0.0, context.sum(NJ.zeros(org.numjvm.dtype.DType.FLOAT32, 0)).getDouble()); + var confined = scope.ones(org.numjvm.dtype.DType.FLOAT64, 65536); + assertEquals(65536.0, context.sum(confined).getDouble()); + assertEquals("sequential", context.lastDecision().strategy()); + var values = new float[65536]; values[7] = Float.NaN; + assertTrue(Double.isNaN(context.sum(NJ.array(values)).getDouble())); + } + } + + @Test void fallsBackForStridesSmallSizesAndRejectsReadOnlyOutput() { + try (var context = NJ.parallel(2)) { + var small = NJ.ones(org.numjvm.dtype.DType.FLOAT32, 10); + var out = NJ.zeros(org.numjvm.dtype.DType.FLOAT32, 10); + context.apply(org.numjvm.ufunc.Ufuncs.NEGATIVE, small, out); + assertEquals("sequential", context.lastDecision().strategy()); + assertEquals(-1.0, out.getDouble(9)); + var source = NJ.ones(org.numjvm.dtype.DType.FLOAT32, 131072) + .slice(org.numjvm.ndarray.Slice.step(2)); + var dest = NJ.zeros(org.numjvm.dtype.DType.FLOAT32, 65536); + context.apply(org.numjvm.ufunc.Ufuncs.SQUARE, source, dest); + assertEquals("sequential", context.lastDecision().strategy()); + assertEquals(1.0, dest.getDouble(65535)); + var readonly = NJ.ones(org.numjvm.dtype.DType.FLOAT32, 1).broadcastTo(65536); + assertThrows(org.numjvm.exception.ReadOnlyArrayException.class, + () -> context.apply(org.numjvm.ufunc.Ufuncs.SQUARE, dest, readonly)); + } + } + + @Test void sumHandlesInfinitiesAndSignedZero() { + try (var context = NJ.parallel(3)) { + float[] values = new float[65536]; + java.util.Arrays.fill(values, -0.0f); + assertEquals(Float.floatToRawIntBits(0.0f), Float.floatToRawIntBits((float) context.sum(NJ.array(values)).getDouble())); + values[0] = Float.POSITIVE_INFINITY; values[values.length - 1] = Float.NEGATIVE_INFINITY; + assertTrue(Double.isNaN(context.sum(NJ.array(values)).getDouble())); + } + } + + @Test void validatesWorkerLimitAndClosesIdempotently() { + assertThrows(IllegalArgumentException.class, () -> NJ.parallel(0)); + assertThrows(IllegalArgumentException.class, () -> NJ.parallel(257)); + var context = NJ.parallel(2); + assertEquals(2, context.workers()); + assertFalse(context.isClosed()); + context.close(); + context.close(); + assertTrue(context.isClosed()); + } + + @Test void closePreservesAnExistingInterrupt() throws Exception { + var observed = new java.util.concurrent.atomic.AtomicBoolean(); + Thread thread = new Thread(() -> { + var context = NJ.parallel(1); + Thread.currentThread().interrupt(); + context.close(); + observed.set(Thread.currentThread().isInterrupted()); + }); + thread.start(); + thread.join(5000); + assertFalse(thread.isAlive()); + assertTrue(observed.get()); + } +} diff --git a/numjvm-core/src/test/java/org/numjvm/ParallelStressTest.java b/numjvm-core/src/test/java/org/numjvm/ParallelStressTest.java new file mode 100644 index 00000000..4227a1e4 --- /dev/null +++ b/numjvm-core/src/test/java/org/numjvm/ParallelStressTest.java @@ -0,0 +1,68 @@ +package org.numjvm; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import org.numjvm.dtype.DType; +import org.numjvm.ndarray.NDArray; +import org.numjvm.ufunc.Ufuncs; +import static org.junit.jupiter.api.Assertions.*; + +class ParallelStressTest { + @Test void concurrentContextsAndRepeatedShutdown() throws Exception { + var callers = Executors.newFixedThreadPool(4); + try { + var shared = NJ.ones(DType.FLOAT32, 65536); + var futures = new java.util.ArrayList>(); + for (int t = 0; t < 4; t++) futures.add(callers.submit(() -> { + for (int repeat = 0; repeat < 4; repeat++) { + try (var context = NJ.parallel(2)) { + var output = NJ.zeros(DType.FLOAT32, 65536); + context.apply(Ufuncs.ADD, shared, shared, output); + assertEquals("parallel", context.lastDecision().strategy()); + assertEquals(2.0, output.getDouble(65535)); + } + } + })); + for (var future : futures) future.get(30, TimeUnit.SECONDS); + } finally { callers.shutdownNow(); } + } + + @Test void nestedCallsAcrossContextsAreRejected() { + try (var first = NJ.parallel(2); var second = NJ.parallel(2)) { + var value = NJ.ones(DType.FLOAT32, 65536); + first.runChunks(new ChunkPlan(1, 1), (i, start, end) -> { + assertThrows(IllegalStateException.class, () -> second.sum(value)); + assertThrows(IllegalStateException.class, second::close); + }); + } + } + + @Test void scopeClosureDuringWorkFailsBeforeCallerReturns() throws Exception { + var scope = NJ.scopeShared(); + try (var context = NJ.parallel(2)) { + var input = scope.ones(DType.FLOAT64, 65536); + var started = new CountDownLatch(1); + var release = new CountDownLatch(1); + var failure = new AtomicReference(); + var caller = new Thread(() -> { + try { + context.runOrElse(new NDArray[]{input}, new ChunkPlan(1, 1), (i, a, b) -> { + started.countDown(); + try { release.await(); } catch (InterruptedException e) { throw new AssertionError(e); } + input.getDouble(0); + }, () -> fail("unexpected fallback")); + } catch (Throwable error) { failure.set(error); } + }); + caller.start(); + assertTrue(started.await(5, TimeUnit.SECONDS)); + scope.close(); + release.countDown(); + caller.join(5000); + assertFalse(caller.isAlive()); + assertNotNull(failure.get()); + } finally { if (scope.isAlive()) scope.close(); } + } +} diff --git a/numjvm-io/build.gradle.kts b/numjvm-io/build.gradle.kts index 55b83dc2..a9e81036 100644 --- a/numjvm-io/build.gradle.kts +++ b/numjvm-io/build.gradle.kts @@ -14,6 +14,7 @@ val oracleFixtures = rootProject.file("numjvm-tests-oracle/fixtures") val oracleResults = layout.buildDirectory.dir("oracle-out") tasks.test { inputs.dir(oracleFixtures).withPathSensitivity(PathSensitivity.RELATIVE) + inputs.file(rootProject.file("numjvm-tests-oracle/parallel-reference.json")) // NumPy verifies these files in a later CI step, including on a cache hit. outputs.dir(oracleResults) systemProperty("numjvm.oracle.fixtures", oracleFixtures.absolutePath) diff --git a/numjvm-io/src/test/java/org/numjvm/io/oracle/ParallelOracleTest.java b/numjvm-io/src/test/java/org/numjvm/io/oracle/ParallelOracleTest.java new file mode 100644 index 00000000..03e583d6 --- /dev/null +++ b/numjvm-io/src/test/java/org/numjvm/io/oracle/ParallelOracleTest.java @@ -0,0 +1,55 @@ +package org.numjvm.io.oracle; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.numjvm.NJ; +import org.numjvm.dtype.DType; +import org.numjvm.ufunc.BinaryUfunc; +import org.numjvm.ufunc.UnaryUfunc; +import org.numjvm.ufunc.Ufuncs; +import static org.junit.jupiter.api.Assertions.*; + +/** Replays NumPy-generated periodic values over arrays large enough to use workers. */ +class ParallelOracleTest { + @SuppressWarnings("unchecked") + @Test void parallelUfuncsAndSumsMatchNumpy() throws Exception { + var path = OracleTest.fixtures().getParent().resolve("parallel-reference.json"); + var root = (Map) Json.parse(Files.readString(path)); + int size = ((Number) root.get("size")).intValue(), period = ((Number) root.get("period")).intValue(); + var cases = (Map) root.get("cases"); + for (var dtype : new DType[]{DType.FLOAT32, DType.FLOAT64}) { + var values = (Map) cases.get(dtype == DType.FLOAT32 ? "float32" : "float64"); + var x = (List) values.get("x"); var y = (List) values.get("y"); + double[] a = new double[size], b = new double[size]; + for (int i = 0; i < size; i++) { a[i] = ((Number) x.get(i % period)).doubleValue(); b[i] = ((Number) y.get(i % period)).doubleValue(); } + var left = NJ.array(a).astype(dtype); var right = NJ.array(b).astype(dtype); + try (var pool = NJ.parallel(3)) { + for (var item : Map.of("add", Ufuncs.ADD, "subtract", Ufuncs.SUBTRACT, + "multiply", Ufuncs.MULTIPLY, "divide", Ufuncs.DIVIDE).entrySet()) { + var out = NJ.zeros(dtype, size); + pool.apply(item.getValue(), left, right, out); + assertEquals("parallel", pool.lastDecision().strategy()); + check(out, (List) values.get(item.getKey()), period, dtype); + } + for (var item : Map.of("negative", Ufuncs.NEGATIVE, "absolute", Ufuncs.ABSOLUTE, + "square", Ufuncs.SQUARE, "sqrt", Ufuncs.SQRT).entrySet()) { + var out = NJ.zeros(dtype, size); + pool.apply(item.getValue(), left, out); + assertEquals("parallel", pool.lastDecision().strategy()); + check(out, (List) values.get(item.getKey()), period, dtype); + } + double expected = ((Number) values.get("sum")).doubleValue(); + assertEquals(expected, pool.sum(left).getDouble(), dtype == DType.FLOAT32 ? 0.1 : 1e-8); + } + } + } + private static void check(org.numjvm.ndarray.NDArray actual, List expected, int period, DType type) { + for (long i = 0; i < actual.size(); i++) { + double value = ((Number) expected.get((int) (i % period))).doubleValue(); + assertEquals(value, actual.getDouble(i), type == DType.FLOAT32 ? 1e-6 : 1e-14, "index " + i); + } + } +} diff --git a/numjvm-jit/src/main/java/org/numjvm/jit/EagerExecutor.java b/numjvm-jit/src/main/java/org/numjvm/jit/EagerExecutor.java index 5d9a9372..0ec816f8 100644 --- a/numjvm-jit/src/main/java/org/numjvm/jit/EagerExecutor.java +++ b/numjvm-jit/src/main/java/org/numjvm/jit/EagerExecutor.java @@ -16,9 +16,23 @@ final class EagerExecutor { private EagerExecutor() { } + static boolean isDirect(Expression expression) { + var root = expression.root(); + return root.left != null && root.left.op == Expr.Op.INPUT + && (root.right == null || root.right.op == Expr.Op.INPUT); + } + /** Returns false without writing if the conservative intermediate budget is exceeded. */ static boolean tryRun(Expression expression, NDArray[] inputs, NDArray output, long budget) { if (budget < 0) throw new IllegalArgumentException("negative temporary budget"); + var root = expression.root(); + // A single ufunc needs no intermediate scope or scheduling metadata. + if (isDirect(expression)) { + NDArray left = inputs[root.left.slot]; + if (root.right == null) unary(root.op).apply(left, output); + else binary(root.op).apply(left, inputs[root.right.slot], output); + return true; + } DType[] inputTypes = new DType[inputs.length]; for (int i = 0; i < inputs.length; i++) inputTypes[i] = inputs[i].dtype(); var types = new TypeResolver(inputTypes); diff --git a/numjvm-jit/src/main/java/org/numjvm/jit/Expression.java b/numjvm-jit/src/main/java/org/numjvm/jit/Expression.java index 6f6e04c4..5b5cee07 100644 --- a/numjvm-jit/src/main/java/org/numjvm/jit/Expression.java +++ b/numjvm-jit/src/main/java/org/numjvm/jit/Expression.java @@ -1,6 +1,9 @@ package org.numjvm.jit; import org.numjvm.NJ; +import org.numjvm.ParallelContext; +import org.numjvm.dtype.DType; +import org.numjvm.ndarray.Slice; import org.numjvm.Experimental; import org.numjvm.internal.SharedSecrets; import org.numjvm.ndarray.NDArray; @@ -75,6 +78,43 @@ public NDArray evalInto(NDArray output, NDArray... inputs) { return output; } + /** Evaluates into an output using an opt-in parallel context when eligible. + * Unsupported layouts, small arrays and confined memory use the caller thread. + * @param parallel reusable CPU execution context + * @param output destination + * @param inputs expression inputs + * @return output + */ + public NDArray evalInto(ParallelContext parallel, NDArray output, NDArray... inputs) { + java.util.Objects.requireNonNull(parallel, "parallel"); + context.admit(); + var plan = new InvocationPlan(this, inputs); + plan.checkOutput(output); + boolean eligible = output.isCContiguous() && !plan.needsTemporary(output) + && (output.dtype() == DType.FLOAT32 || output.dtype() == DType.FLOAT64); + for (var input : inputs) { + eligible &= input.isCContiguous() && input.shape().equals(output.shape()) + && (input.dtype() == DType.FLOAT32 || input.dtype() == DType.FLOAT64); + } + NDArray[] arrays = java.util.Arrays.copyOf(inputs, inputs.length + 1); + arrays[inputs.length] = output; + NDArray[] flat = null; + NDArray flatOutput = null; + if (eligible) { + flat = new NDArray[inputs.length]; + for (int i = 0; i < inputs.length; i++) flat[i] = inputs[i].reshapeView(inputs[i].size()); + flatOutput = output.reshapeView(output.size()); + } + NDArray[] finalFlat = flat; + NDArray finalOutput = flatOutput; + SharedSecrets.parallelAccess().execute(parallel, arrays, output.size(), eligible, (start, end) -> { + NDArray[] part = new NDArray[finalFlat.length]; + for (int i = 0; i < part.length; i++) part[i] = finalFlat[i].slice(Slice.range(start, end)); + evalInto(finalOutput.slice(Slice.range(start, end)), part); + }, () -> evalInto(output, inputs)); + return output; + } + private void execute(NDArray[] inputs, NDArray output) { if (context.mode() == Jit.Mode.EAGER) { if (!EagerExecutor.tryRun(this, inputs, output, context.eagerBytes())) @@ -90,24 +130,31 @@ private void execute(NDArray[] inputs, NDArray output) { context.selected("block", "explicit mode"); return; } + long size = output.shape().size(); + if (size == 0) { context.selected("empty", "zero elements"); return; } + boolean small = size <= 512; + if (context.mode() == Jit.Mode.AUTO && small && EagerExecutor.tryRun(this, inputs, output, context.eagerBytes())) { + context.eagerRun(); + context.selected("eager", "small expression"); + return; + } boolean eligible = KernelLoader.supports(this) && GeneratedExecutor.eligible(inputs, output); if (context.mode() == Jit.Mode.REQUIRED && !eligible) throw new IllegalArgumentException("generated execution does not support this layout or graph size"); - long size = output.shape().size(); - if (size == 0) { context.selected("empty", "zero elements"); return; } var types = new org.numjvm.dtype.DType[inputs.length]; for (int i = 0; i < inputs.length; i++) types[i] = inputs[i].dtype(); boolean cached = eligible && context.cached(this, types); boolean homogeneous = true; for (int i = 1; i < types.length; i++) homogeneous &= types[i] == types[0]; boolean vectorProvider = !org.numjvm.BackendInfo.simdProvider().equals("scalar"); - boolean small = size <= 512; - boolean eagerCandidate = small || (size <= 65536 && homogeneous && vectorProvider && graph.nodes().size() >= 3); - if (context.mode() == Jit.Mode.AUTO && small && EagerExecutor.tryRun(this, inputs, output, context.eagerBytes())) { + // There is no intermediate traffic to save for one direct ufunc. + if (context.mode() == Jit.Mode.AUTO && vectorProvider && EagerExecutor.isDirect(this) + && EagerExecutor.tryRun(this, inputs, output, context.eagerBytes())) { context.eagerRun(); - context.selected("eager", "small expression"); + context.selected("eager", "single ufunc"); return; } + boolean eagerCandidate = small || (size <= 65536 && homogeneous && vectorProvider && graph.nodes().size() >= 3); String fusedFallback = null; if (context.mode() == Jit.Mode.AUTO && eligible && vectorProvider && FusedProviders.provider() != null) { boolean fusedCached = context.fusedCached(this, types); diff --git a/numjvm-jit/src/test/java/org/numjvm/jit/EagerExecutorTest.java b/numjvm-jit/src/test/java/org/numjvm/jit/EagerExecutorTest.java index 4763ac48..b7ac4596 100644 --- a/numjvm-jit/src/test/java/org/numjvm/jit/EagerExecutorTest.java +++ b/numjvm-jit/src/test/java/org/numjvm/jit/EagerExecutorTest.java @@ -8,6 +8,19 @@ import static org.junit.jupiter.api.Assertions.*; class EagerExecutorTest { + @Test void directSingleUfuncSupportsMixedInputsAndInPlaceOutputWithoutScratch() { + try (var context = Jit.context(Jit.Mode.EAGER, 4, 2048, 0)) { + var left = NJ.ones(DType.FLOAT64, 17); + var right = NJ.ones(DType.FLOAT32, 17); + var add = context.trace(2, i -> i.input(0).add(i.input(1))); + assertSame(left, add.evalInto(left, left, right)); + assertEquals(2.0, left.getDouble(16)); + var square = context.trace(1, i -> i.input(0).square()); + square.evalInto(left, left); + assertEquals(4.0, left.getDouble(16)); + } + } + @Test void sharedDagWeakScalarsAndMixedTypesMatchEager() { try (var context = Jit.context(Jit.Mode.BLOCK)) { var expression = context.trace(2, i -> { diff --git a/numjvm-jit/src/test/java/org/numjvm/jit/FusedDispatchTest.java b/numjvm-jit/src/test/java/org/numjvm/jit/FusedDispatchTest.java index a01b1ec1..0a9b59e7 100644 --- a/numjvm-jit/src/test/java/org/numjvm/jit/FusedDispatchTest.java +++ b/numjvm-jit/src/test/java/org/numjvm/jit/FusedDispatchTest.java @@ -9,6 +9,18 @@ import static org.junit.jupiter.api.Assertions.*; class FusedDispatchTest { + @Test void singleVectorUfuncUsesEagerWithoutCompilation() { + Assumptions.assumeTrue(FusedProviders.provider() != null, "optional vector provider unavailable"); + try (var context = Jit.context(Jit.Mode.AUTO)) { + var graph = context.trace(2, i -> i.input(0).add(i.input(1))); + var a = NJ.ones(DType.FLOAT32, 65536); + assertEquals(2.0, graph.eval(a, a).getDouble(0)); + assertEquals("eager", context.lastDecision().strategy()); + assertEquals("single ufunc", context.lastDecision().reason()); + assertEquals(0, context.fusedStatistics().compilations()); + } + } + @Test void autoRunsAndCachesFusedGraphWhileExplicitModesRemainStable() { Assumptions.assumeTrue(FusedProviders.provider() != null, "optional vector provider unavailable"); try (var context = Jit.context(Jit.Mode.AUTO, 1)) { diff --git a/numjvm-jit/src/test/java/org/numjvm/jit/FusedFloatTest.java b/numjvm-jit/src/test/java/org/numjvm/jit/FusedFloatTest.java index f87b43d2..1c2f298b 100644 --- a/numjvm-jit/src/test/java/org/numjvm/jit/FusedFloatTest.java +++ b/numjvm-jit/src/test/java/org/numjvm/jit/FusedFloatTest.java @@ -8,6 +8,25 @@ import static org.junit.jupiter.api.Assertions.*; class FusedFloatTest { + @Test void generatedLoopHandlesFiveInputsBroadcastAndTail() { + Assumptions.assumeTrue(FusedProviders.provider() != null, "optional vector provider unavailable"); + try (var context = Jit.context(Jit.Mode.REQUIRED)) { + var graph = context.trace(5, i -> i.input(0).multiply(i.input(1)).add(i.input(2)) + .divide(i.input(3)).subtract(i.input(4))); + NDArray[] inputs = {NJ.ones(DType.FLOAT32, 513), NJ.ones(DType.FLOAT32, 513), + NJ.ones(DType.FLOAT32, 1), NJ.ones(DType.FLOAT32, 513), NJ.ones(DType.FLOAT32, 513)}; + var kernel = FusedProviders.compile(graph, new DType[]{DType.FLOAT32, DType.FLOAT32, + DType.FLOAT32, DType.FLOAT32, DType.FLOAT32}); + assertNotNull(kernel); + var actual = new InvocationPlan(graph, inputs).allocate(); + GeneratedExecutor.run(kernel, inputs, actual); + var expected = graph.eval(inputs); + for (int i = 0; i < 513; i++) + assertEquals(Float.floatToRawIntBits((float) expected.getDouble(i)), + Float.floatToRawIntBits((float) actual.getDouble(i)), "index " + i); + } + } + @Test void allOperationsMatchScalar() { Assumptions.assumeTrue(FusedProviders.provider() != null, "optional vector provider unavailable"); var input = NJ.array(new float[]{-0.0f, 0, -1, 1, Float.NaN, Float.POSITIVE_INFINITY, diff --git a/numjvm-jit/src/test/java/org/numjvm/jit/FusedMixedTest.java b/numjvm-jit/src/test/java/org/numjvm/jit/FusedMixedTest.java index 952958a6..a29ef10a 100644 --- a/numjvm-jit/src/test/java/org/numjvm/jit/FusedMixedTest.java +++ b/numjvm-jit/src/test/java/org/numjvm/jit/FusedMixedTest.java @@ -5,9 +5,25 @@ import org.numjvm.NJ; import org.numjvm.dtype.DType; import org.numjvm.ndarray.NDArray; +import org.numjvm.ndarray.Order; import static org.junit.jupiter.api.Assertions.*; class FusedMixedTest { + @Test void fortranMixedRunsKeepIntermediateRounding() { + Assumptions.assumeTrue(FusedProviders.provider() != null, "optional vector provider unavailable"); + try (var context = Jit.context(Jit.Mode.REQUIRED)) { + var graph = context.trace(2, i -> i.input(0).add(Math.scalb(1.0, -24)).add(i.input(1))); + NDArray[] inputs = {NJ.ones(DType.FLOAT32, 33, 17).copy(Order.F), + NJ.multiply(NJ.ones(DType.FLOAT64, 33, 17), Math.scalb(1.0, -25)).copy(Order.F)}; + var output = NJ.zeros(DType.FLOAT64, 33, 17).copy(Order.F); + var kernel = FusedProviders.compile(graph, new DType[]{DType.FLOAT32, DType.FLOAT64}); + assertNotNull(kernel); + GeneratedExecutor.run(kernel, inputs, output); + for (int i = 0; i < 33; i++) for (int j = 0; j < 17; j++) + assertEquals(1 + Math.scalb(1.0, -25), output.getDouble(i, j)); + } + } + @Test void mixedSpecialValuesAndBroadcastFollowScalarGraph() { Assumptions.assumeTrue(FusedProviders.provider() != null, "optional vector provider unavailable"); float[] data = new float[35]; diff --git a/numjvm-jit/src/test/java/org/numjvm/jit/ParallelExpressionTest.java b/numjvm-jit/src/test/java/org/numjvm/jit/ParallelExpressionTest.java new file mode 100644 index 00000000..c336866b --- /dev/null +++ b/numjvm-jit/src/test/java/org/numjvm/jit/ParallelExpressionTest.java @@ -0,0 +1,41 @@ +package org.numjvm.jit; + +import org.junit.jupiter.api.Test; +import org.numjvm.NJ; +import org.numjvm.dtype.DType; +import static org.junit.jupiter.api.Assertions.*; + +class ParallelExpressionTest { + @Test void handlesIdenticalOutputAndHazardousOverlap() { + try (var pool = NJ.parallel(3); var jit = Jit.context(Jit.Mode.AUTO)) { + var expression = jit.trace(1, i -> i.input(0).add(2).square()); + var input = NJ.ones(DType.FLOAT32, 65536); + expression.evalInto(pool, input, input); + assertEquals("parallel", pool.lastDecision().strategy()); + assertEquals(9.0, input.getDouble(65535)); + var backing = NJ.ones(DType.FLOAT32, 65537); + var source = backing.slice(org.numjvm.ndarray.Slice.range(0, 65536)); + var destination = backing.slice(org.numjvm.ndarray.Slice.range(1, 65537)); + expression.evalInto(pool, destination, source); + assertEquals("sequential", pool.lastDecision().strategy()); + assertEquals(1.0, backing.getDouble(0)); + assertEquals(9.0, backing.getDouble(65536)); + } + } + + @Test void evaluatesContiguousExpressionsAndConfinedFallback() { + try (var pool = NJ.parallel(3); var jit = Jit.context(Jit.Mode.AUTO); var scope = NJ.scope()) { + var expression = jit.trace(2, i -> i.input(0).add(i.input(1)).square()); + var left = NJ.ones(DType.FLOAT32, 256, 256); + var right = NJ.ones(DType.FLOAT32, 256, 256); + var output = NJ.zeros(DType.FLOAT32, 256, 256); + assertSame(output, expression.evalInto(pool, output, left, right)); + assertEquals("parallel", pool.lastDecision().strategy()); + assertEquals(4.0, output.getDouble(255, 255)); + var confined = scope.ones(DType.FLOAT32, 65536); + expression.evalInto(pool, confined, confined, confined); + assertEquals("sequential", pool.lastDecision().strategy()); + assertEquals(4.0, confined.getDouble(65535)); + } + } +} diff --git a/numjvm-tests-oracle/generate_parallel_cases.py b/numjvm-tests-oracle/generate_parallel_cases.py new file mode 100644 index 00000000..2a6446c4 --- /dev/null +++ b/numjvm-tests-oracle/generate_parallel_cases.py @@ -0,0 +1,38 @@ +"""Generate compact periodic NumPy reference values for parallel operations.""" +import argparse +import json +from pathlib import Path +import numpy as np + +PATH = Path(__file__).with_name('parallel-reference.json') +N = 65537 + +def render(): + cases = {} + for dtype in (np.float32, np.float64): + x = (np.arange(37, dtype=dtype) + dtype(1)) / dtype(8) + y = ((np.arange(37, dtype=dtype) * dtype(7)) % dtype(37) + dtype(1)) / dtype(8) + cases[np.dtype(dtype).name] = { + 'x': x.tolist(), 'y': y.tolist(), + 'add': np.add(x, y).tolist(), 'subtract': np.subtract(x, y).tolist(), + 'multiply': np.multiply(x, y).tolist(), 'divide': np.divide(x, y).tolist(), + 'negative': np.negative(x).tolist(), 'absolute': np.absolute(x).tolist(), + 'square': np.square(x).tolist(), 'sqrt': np.sqrt(x).tolist(), + 'sum': float(np.sum(np.resize(x, N), dtype=dtype)), + } + return json.dumps({'size': N, 'period': 37, 'cases': cases}, indent=2, sort_keys=True) + '\n' + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--check', action='store_true') + args = parser.parse_args() + expected = render() + if args.check: + if not PATH.exists() or PATH.read_text() != expected: + raise SystemExit('parallel NumPy fixtures differ; regenerate with pinned NumPy') + print('parallel NumPy fixtures match') + else: + PATH.write_text(expected) + +if __name__ == '__main__': + main() diff --git a/numjvm-tests-oracle/parallel-reference.json b/numjvm-tests-oracle/parallel-reference.json new file mode 100644 index 00000000..c4e59517 --- /dev/null +++ b/numjvm-tests-oracle/parallel-reference.json @@ -0,0 +1,792 @@ +{ + "cases": { + "float32": { + "absolute": [ + 0.125, + 0.25, + 0.375, + 0.5, + 0.625, + 0.75, + 0.875, + 1.0, + 1.125, + 1.25, + 1.375, + 1.5, + 1.625, + 1.75, + 1.875, + 2.0, + 2.125, + 2.25, + 2.375, + 2.5, + 2.625, + 2.75, + 2.875, + 3.0, + 3.125, + 3.25, + 3.375, + 3.5, + 3.625, + 3.75, + 3.875, + 4.0, + 4.125, + 4.25, + 4.375, + 4.5, + 4.625 + ], + "add": [ + 0.25, + 1.25, + 2.25, + 3.25, + 4.25, + 5.25, + 1.625, + 2.625, + 3.625, + 4.625, + 5.625, + 2.0, + 3.0, + 4.0, + 5.0, + 6.0, + 2.375, + 3.375, + 4.375, + 5.375, + 6.375, + 7.375, + 3.75, + 4.75, + 5.75, + 6.75, + 7.75, + 4.125, + 5.125, + 6.125, + 7.125, + 8.125, + 4.5, + 5.5, + 6.5, + 7.5, + 8.5 + ], + "divide": [ + 1.0, + 0.25, + 0.20000000298023224, + 0.1818181872367859, + 0.17241379618644714, + 0.1666666716337204, + 1.1666666269302368, + 0.6153846383094788, + 0.44999998807907104, + 0.37037035822868347, + 0.3235294222831726, + 3.0, + 1.1818181276321411, + 0.7777777910232544, + 0.6000000238418579, + 0.5, + 8.5, + 2.0, + 1.1875, + 0.8695651888847351, + 0.699999988079071, + 0.5945945978164673, + 3.2857143878936768, + 1.7142857313156128, + 1.1904761791229248, + 0.9285714030265808, + 0.7714285850524902, + 5.599999904632568, + 2.4166667461395264, + 1.5789474248886108, + 1.192307710647583, + 0.9696969985961914, + 11.0, + 3.4000000953674316, + 2.058823585510254, + 1.5, + 1.1935484409332275 + ], + "multiply": [ + 0.015625, + 0.25, + 0.703125, + 1.375, + 2.265625, + 3.375, + 0.65625, + 1.625, + 2.8125, + 4.21875, + 5.84375, + 0.75, + 2.234375, + 3.9375, + 5.859375, + 8.0, + 0.53125, + 2.53125, + 4.75, + 7.1875, + 9.84375, + 12.71875, + 2.515625, + 5.25, + 8.203125, + 11.375, + 14.765625, + 2.1875, + 5.4375, + 8.90625, + 12.59375, + 16.5, + 1.546875, + 5.3125, + 9.296875, + 13.5, + 17.921875 + ], + "negative": [ + -0.125, + -0.25, + -0.375, + -0.5, + -0.625, + -0.75, + -0.875, + -1.0, + -1.125, + -1.25, + -1.375, + -1.5, + -1.625, + -1.75, + -1.875, + -2.0, + -2.125, + -2.25, + -2.375, + -2.5, + -2.625, + -2.75, + -2.875, + -3.0, + -3.125, + -3.25, + -3.375, + -3.5, + -3.625, + -3.75, + -3.875, + -4.0, + -4.125, + -4.25, + -4.375, + -4.5, + -4.625 + ], + "sqrt": [ + 0.3535533845424652, + 0.5, + 0.6123724579811096, + 0.7071067690849304, + 0.7905694246292114, + 0.8660253882408142, + 0.9354143738746643, + 1.0, + 1.0606601238250732, + 1.1180340051651, + 1.172603964805603, + 1.2247449159622192, + 1.2747548818588257, + 1.3228756189346313, + 1.3693064451217651, + 1.4142135381698608, + 1.457737922668457, + 1.5, + 1.541103482246399, + 1.5811388492584229, + 1.620185136795044, + 1.658312439918518, + 1.6955825090408325, + 1.7320507764816284, + 1.7677669525146484, + 1.8027756214141846, + 1.837117314338684, + 1.8708287477493286, + 1.9039433002471924, + 1.9364917278289795, + 1.9685019254684448, + 2.0, + 2.0310096740722656, + 2.0615527629852295, + 2.0916500091552734, + 2.1213202476501465, + 2.1505813598632812 + ], + "square": [ + 0.015625, + 0.0625, + 0.140625, + 0.25, + 0.390625, + 0.5625, + 0.765625, + 1.0, + 1.265625, + 1.5625, + 1.890625, + 2.25, + 2.640625, + 3.0625, + 3.515625, + 4.0, + 4.515625, + 5.0625, + 5.640625, + 6.25, + 6.890625, + 7.5625, + 8.265625, + 9.0, + 9.765625, + 10.5625, + 11.390625, + 12.25, + 13.140625, + 14.0625, + 15.015625, + 16.0, + 17.015625, + 18.0625, + 19.140625, + 20.25, + 21.390625 + ], + "subtract": [ + 0.0, + -0.75, + -1.5, + -2.25, + -3.0, + -3.75, + 0.125, + -0.625, + -1.375, + -2.125, + -2.875, + 1.0, + 0.25, + -0.5, + -1.25, + -2.0, + 1.875, + 1.125, + 0.375, + -0.375, + -1.125, + -1.875, + 2.0, + 1.25, + 0.5, + -0.25, + -1.0, + 2.875, + 2.125, + 1.375, + 0.625, + -0.125, + 3.75, + 3.0, + 2.25, + 1.5, + 0.75 + ], + "sum": 155633.5, + "x": [ + 0.125, + 0.25, + 0.375, + 0.5, + 0.625, + 0.75, + 0.875, + 1.0, + 1.125, + 1.25, + 1.375, + 1.5, + 1.625, + 1.75, + 1.875, + 2.0, + 2.125, + 2.25, + 2.375, + 2.5, + 2.625, + 2.75, + 2.875, + 3.0, + 3.125, + 3.25, + 3.375, + 3.5, + 3.625, + 3.75, + 3.875, + 4.0, + 4.125, + 4.25, + 4.375, + 4.5, + 4.625 + ], + "y": [ + 0.125, + 1.0, + 1.875, + 2.75, + 3.625, + 4.5, + 0.75, + 1.625, + 2.5, + 3.375, + 4.25, + 0.5, + 1.375, + 2.25, + 3.125, + 4.0, + 0.25, + 1.125, + 2.0, + 2.875, + 3.75, + 4.625, + 0.875, + 1.75, + 2.625, + 3.5, + 4.375, + 0.625, + 1.5, + 2.375, + 3.25, + 4.125, + 0.375, + 1.25, + 2.125, + 3.0, + 3.875 + ] + }, + "float64": { + "absolute": [ + 0.125, + 0.25, + 0.375, + 0.5, + 0.625, + 0.75, + 0.875, + 1.0, + 1.125, + 1.25, + 1.375, + 1.5, + 1.625, + 1.75, + 1.875, + 2.0, + 2.125, + 2.25, + 2.375, + 2.5, + 2.625, + 2.75, + 2.875, + 3.0, + 3.125, + 3.25, + 3.375, + 3.5, + 3.625, + 3.75, + 3.875, + 4.0, + 4.125, + 4.25, + 4.375, + 4.5, + 4.625 + ], + "add": [ + 0.25, + 1.25, + 2.25, + 3.25, + 4.25, + 5.25, + 1.625, + 2.625, + 3.625, + 4.625, + 5.625, + 2.0, + 3.0, + 4.0, + 5.0, + 6.0, + 2.375, + 3.375, + 4.375, + 5.375, + 6.375, + 7.375, + 3.75, + 4.75, + 5.75, + 6.75, + 7.75, + 4.125, + 5.125, + 6.125, + 7.125, + 8.125, + 4.5, + 5.5, + 6.5, + 7.5, + 8.5 + ], + "divide": [ + 1.0, + 0.25, + 0.2, + 0.18181818181818182, + 0.1724137931034483, + 0.16666666666666666, + 1.1666666666666667, + 0.6153846153846154, + 0.45, + 0.37037037037037035, + 0.3235294117647059, + 3.0, + 1.1818181818181819, + 0.7777777777777778, + 0.6, + 0.5, + 8.5, + 2.0, + 1.1875, + 0.8695652173913043, + 0.7, + 0.5945945945945946, + 3.2857142857142856, + 1.7142857142857142, + 1.1904761904761905, + 0.9285714285714286, + 0.7714285714285715, + 5.6, + 2.4166666666666665, + 1.5789473684210527, + 1.1923076923076923, + 0.9696969696969697, + 11.0, + 3.4, + 2.0588235294117645, + 1.5, + 1.1935483870967742 + ], + "multiply": [ + 0.015625, + 0.25, + 0.703125, + 1.375, + 2.265625, + 3.375, + 0.65625, + 1.625, + 2.8125, + 4.21875, + 5.84375, + 0.75, + 2.234375, + 3.9375, + 5.859375, + 8.0, + 0.53125, + 2.53125, + 4.75, + 7.1875, + 9.84375, + 12.71875, + 2.515625, + 5.25, + 8.203125, + 11.375, + 14.765625, + 2.1875, + 5.4375, + 8.90625, + 12.59375, + 16.5, + 1.546875, + 5.3125, + 9.296875, + 13.5, + 17.921875 + ], + "negative": [ + -0.125, + -0.25, + -0.375, + -0.5, + -0.625, + -0.75, + -0.875, + -1.0, + -1.125, + -1.25, + -1.375, + -1.5, + -1.625, + -1.75, + -1.875, + -2.0, + -2.125, + -2.25, + -2.375, + -2.5, + -2.625, + -2.75, + -2.875, + -3.0, + -3.125, + -3.25, + -3.375, + -3.5, + -3.625, + -3.75, + -3.875, + -4.0, + -4.125, + -4.25, + -4.375, + -4.5, + -4.625 + ], + "sqrt": [ + 0.3535533905932738, + 0.5, + 0.6123724356957945, + 0.7071067811865476, + 0.7905694150420949, + 0.8660254037844386, + 0.9354143466934853, + 1.0, + 1.0606601717798212, + 1.118033988749895, + 1.1726039399558574, + 1.224744871391589, + 1.2747548783981961, + 1.3228756555322954, + 1.3693063937629153, + 1.4142135623730951, + 1.4577379737113252, + 1.5, + 1.541103500742244, + 1.5811388300841898, + 1.620185174601965, + 1.6583123951777, + 1.695582495781317, + 1.7320508075688772, + 1.7677669529663689, + 1.8027756377319946, + 1.8371173070873836, + 1.8708286933869707, + 1.9039432764659772, + 1.9364916731037085, + 1.9685019685029528, + 2.0, + 2.03100960115899, + 2.0615528128088303, + 2.091650066335189, + 2.1213203435596424, + 2.1505813167606567 + ], + "square": [ + 0.015625, + 0.0625, + 0.140625, + 0.25, + 0.390625, + 0.5625, + 0.765625, + 1.0, + 1.265625, + 1.5625, + 1.890625, + 2.25, + 2.640625, + 3.0625, + 3.515625, + 4.0, + 4.515625, + 5.0625, + 5.640625, + 6.25, + 6.890625, + 7.5625, + 8.265625, + 9.0, + 9.765625, + 10.5625, + 11.390625, + 12.25, + 13.140625, + 14.0625, + 15.015625, + 16.0, + 17.015625, + 18.0625, + 19.140625, + 20.25, + 21.390625 + ], + "subtract": [ + 0.0, + -0.75, + -1.5, + -2.25, + -3.0, + -3.75, + 0.125, + -0.625, + -1.375, + -2.125, + -2.875, + 1.0, + 0.25, + -0.5, + -1.25, + -2.0, + 1.875, + 1.125, + 0.375, + -0.375, + -1.125, + -1.875, + 2.0, + 1.25, + 0.5, + -0.25, + -1.0, + 2.875, + 2.125, + 1.375, + 0.625, + -0.125, + 3.75, + 3.0, + 2.25, + 1.5, + 0.75 + ], + "sum": 155633.5, + "x": [ + 0.125, + 0.25, + 0.375, + 0.5, + 0.625, + 0.75, + 0.875, + 1.0, + 1.125, + 1.25, + 1.375, + 1.5, + 1.625, + 1.75, + 1.875, + 2.0, + 2.125, + 2.25, + 2.375, + 2.5, + 2.625, + 2.75, + 2.875, + 3.0, + 3.125, + 3.25, + 3.375, + 3.5, + 3.625, + 3.75, + 3.875, + 4.0, + 4.125, + 4.25, + 4.375, + 4.5, + 4.625 + ], + "y": [ + 0.125, + 1.0, + 1.875, + 2.75, + 3.625, + 4.5, + 0.75, + 1.625, + 2.5, + 3.375, + 4.25, + 0.5, + 1.375, + 2.25, + 3.125, + 4.0, + 0.25, + 1.125, + 2.0, + 2.875, + 3.75, + 4.625, + 0.875, + 1.75, + 2.625, + 3.5, + 4.375, + 0.625, + 1.5, + 2.375, + 3.25, + 4.125, + 0.375, + 1.25, + 2.125, + 3.0, + 3.875 + ] + } + }, + "period": 37, + "size": 65537 +} diff --git a/numjvm-vector/src/main/java/org/numjvm/vector/DoubleFusedKernel.java b/numjvm-vector/src/main/java/org/numjvm/vector/DoubleFusedKernel.java index 87a52b54..9e1fa129 100644 --- a/numjvm-vector/src/main/java/org/numjvm/vector/DoubleFusedKernel.java +++ b/numjvm-vector/src/main/java/org/numjvm/vector/DoubleFusedKernel.java @@ -17,12 +17,9 @@ final class DoubleFusedKernel implements FusedKernel { var species = DoubleVector.SPECIES_PREFERRED; var values = new DoubleVector[segments.length - 1]; long bound = species.loopBound(count), i = 0; - for (; i < bound; i += species.length()) { - for (int slot = 0; slot < values.length; slot++) - values[slot] = strides[slot] == 0 - ? DoubleVector.broadcast(species, segments[slot].get(ValueLayout.JAVA_DOUBLE_UNALIGNED, offsets[slot])) - : DoubleVector.fromMemorySegment(species, segments[slot], offsets[slot] + i * strides[slot], ByteOrder.nativeOrder()); - expression.evaluate(values).intoMemorySegment(segments[values.length], offsets[values.length] + i * strides[values.length], ByteOrder.nativeOrder()); + if (bound > 0) { + ((VectorLoop) expression).runVectors(segments, offsets, strides, bound, species.length()); + i = bound; } // Scalar loads/stores keep short tails within their memory segment bounds. for (; i < count; i++) { diff --git a/numjvm-vector/src/main/java/org/numjvm/vector/FloatFusedKernel.java b/numjvm-vector/src/main/java/org/numjvm/vector/FloatFusedKernel.java index 5d668f3d..02e7f264 100644 --- a/numjvm-vector/src/main/java/org/numjvm/vector/FloatFusedKernel.java +++ b/numjvm-vector/src/main/java/org/numjvm/vector/FloatFusedKernel.java @@ -17,12 +17,9 @@ final class FloatFusedKernel implements FusedKernel { var species = FloatVector.SPECIES_PREFERRED; var values = new FloatVector[segments.length - 1]; long bound = species.loopBound(count), i = 0; - for (; i < bound; i += species.length()) { - for (int slot = 0; slot < values.length; slot++) - values[slot] = strides[slot] == 0 - ? FloatVector.broadcast(species, segments[slot].get(ValueLayout.JAVA_FLOAT_UNALIGNED, offsets[slot])) - : FloatVector.fromMemorySegment(species, segments[slot], offsets[slot] + i * strides[slot], ByteOrder.nativeOrder()); - expression.evaluate(values).intoMemorySegment(segments[values.length], offsets[values.length] + i * strides[values.length], ByteOrder.nativeOrder()); + if (bound > 0) { + ((VectorLoop) expression).runVectors(segments, offsets, strides, bound, species.length()); + i = bound; } // Scalar loads/stores keep short tails within their memory segment bounds. for (; i < count; i++) { diff --git a/numjvm-vector/src/main/java/org/numjvm/vector/MixedExpressionCompiler.java b/numjvm-vector/src/main/java/org/numjvm/vector/MixedExpressionCompiler.java index 4c3aeabe..408cddfe 100644 --- a/numjvm-vector/src/main/java/org/numjvm/vector/MixedExpressionCompiler.java +++ b/numjvm-vector/src/main/java/org/numjvm/vector/MixedExpressionCompiler.java @@ -34,11 +34,16 @@ static MixedExpression compile(FusedGraph graph) { } ClassDesc name = ClassDesc.of("org.numjvm.vector.GeneratedMixedExpression"); byte[] bytes = ClassFile.of().build(name, builder -> builder.withFlags(ClassFile.ACC_FINAL | ClassFile.ACC_SUPER) - .withInterfaceSymbols(ClassDesc.of("org.numjvm.vector.MixedExpression")) + .withInterfaceSymbols(ClassDesc.of("org.numjvm.vector.MixedExpression"), ClassDesc.of("org.numjvm.vector.VectorLoop")) .withMethodBody("", MethodTypeDesc.of(ConstantDescs.CD_void), 0, code -> code.aload(0).invokespecial(ConstantDescs.CD_Object, "", MethodTypeDesc.of(ConstantDescs.CD_void)).return_()) .withMethodBody("evaluate", MethodTypeDesc.of(DOUBLE, ConstantDescs.CD_Object.arrayType()), ClassFile.ACC_PUBLIC, - code -> emit(code, graph))); + code -> { emit(code, graph, 2, null, null, null); code.areturn(); }) + .withMethodBody("runVectors", MethodTypeDesc.of(ConstantDescs.CD_void, + ClassDesc.of("java.lang.foreign.MemorySegment").arrayType(), + ConstantDescs.CD_long.arrayType(), ConstantDescs.CD_long.arrayType(), + ConstantDescs.CD_long, ConstantDescs.CD_int), ClassFile.ACC_PUBLIC, + code -> emitLoop(code, graph))); try { var lookup = MethodHandles.lookup().defineHiddenClass(bytes, true); return (MixedExpression) lookup.findConstructor(lookup.lookupClass(), MethodType.methodType(void.class)).invoke(); @@ -46,21 +51,51 @@ static MixedExpression compile(FusedGraph graph) { catch (Throwable error) { throw new IllegalStateException("cannot define mixed fused expression", error); } } - private static void operand(CodeBuilder code, FusedGraph graph, int index, DType target) { + private static void operand(CodeBuilder code, FusedGraph graph, int index, DType target, int first) { DType source = graph.nodes().get(index).type(); - code.aload(2 + index); + code.aload(first + index); if (source == DType.FLOAT32 && target == DType.FLOAT64) code.invokestatic(CONVERT, "toDouble", MethodTypeDesc.of(DOUBLE, FLOAT)); else if (source == DType.FLOAT64 && target == DType.FLOAT32) code.invokestatic(CONVERT, "toFloat", MethodTypeDesc.of(FLOAT, DOUBLE)); } - private static void emit(CodeBuilder code, FusedGraph graph) { + private static void emitLoop(CodeBuilder code, FusedGraph graph) { + int count = graph.inputTypes().size() + 1; + int[] segments = new int[count], offsets = new int[count], strides = new int[count]; + int local = 9; + for (int i = 0; i < count; i++) { + segments[i] = local++; offsets[i] = local; local += 2; strides[i] = local; local += 2; + code.aload(1).loadConstant(i).aaload().astore(segments[i]); + code.aload(2).loadConstant(i).laload().lstore(offsets[i]); + code.aload(3).loadConstant(i).laload().lstore(strides[i]); + } + var loop = code.newLabel(); + var done = code.newLabel(); + code.lconst_0().lstore(7).labelBinding(loop).lload(7).lload(4).lcmp().ifge(done); + code.aload(segments[count - 1]).lload(offsets[count - 1]); + emit(code, graph, local, segments, offsets, strides); + code.invokestatic(ClassDesc.of("org.numjvm.vector.VectorMemory"), "storeDouble", + MethodTypeDesc.of(ConstantDescs.CD_void, ClassDesc.of("java.lang.foreign.MemorySegment"), ConstantDescs.CD_long, DOUBLE)); + for (int i = 0; i < count; i++) + code.lload(offsets[i]).lload(strides[i]).iload(6).i2l().lmul().ladd().lstore(offsets[i]); + code.lload(7).iload(6).i2l().ladd().lstore(7).goto_(loop).labelBinding(done).return_(); + } + + private static void emit(CodeBuilder code, FusedGraph graph, int first, + int[] segments, int[] offsets, int[] strides) { for (int i = 0; i < graph.nodes().size(); i++) { var node = graph.nodes().get(i); ClassDesc type = node.type() == DType.FLOAT32 ? FLOAT : DOUBLE; switch (node.op()) { - case INPUT -> code.aload(1).loadConstant(node.slot()).aaload().checkcast(type); + case INPUT -> { + if (segments == null) code.aload(1).loadConstant(node.slot()).aaload().checkcast(type); + else code.aload(segments[node.slot()]).lload(offsets[node.slot()]).lload(strides[node.slot()]) + .invokestatic(ClassDesc.of("org.numjvm.vector.VectorMemory"), + node.type() == DType.FLOAT32 ? "loadMixedFloat" : "loadDouble", + MethodTypeDesc.of(type, ClassDesc.of("java.lang.foreign.MemorySegment"), + ConstantDescs.CD_long, ConstantDescs.CD_long)); + } case CONSTANT -> { code.getstatic(CONVERT, "DOUBLE", SPECIES).loadConstant(node.bits()) .invokestatic(ClassDesc.of("java.lang.Double"), "longBitsToDouble", @@ -71,21 +106,21 @@ private static void emit(CodeBuilder code, FusedGraph graph) { case ADD, SUBTRACT, MULTIPLY, DIVIDE -> { String method = switch (node.op()) { case ADD -> "add"; case SUBTRACT -> "sub"; case MULTIPLY -> "mul"; default -> "div"; }; - operand(code, graph, node.left(), node.type()); operand(code, graph, node.right(), node.type()); + operand(code, graph, node.left(), node.type(), first); operand(code, graph, node.right(), node.type(), first); code.invokevirtual(type, method, MethodTypeDesc.of(type, VECTOR)); } case NEGATIVE, ABSOLUTE, SQRT -> { String method = node.op() == FusedGraph.Op.NEGATIVE ? "neg" : node.op() == FusedGraph.Op.ABSOLUTE ? "abs" : "sqrt"; - operand(code, graph, node.left(), node.type()); + operand(code, graph, node.left(), node.type(), first); code.invokevirtual(type, method, MethodTypeDesc.of(type)); } case SQUARE -> { - operand(code, graph, node.left(), node.type()); operand(code, graph, node.left(), node.type()); + operand(code, graph, node.left(), node.type(), first); operand(code, graph, node.left(), node.type(), first); code.invokevirtual(type, "mul", MethodTypeDesc.of(type, VECTOR)); } } - code.astore(2 + i); + code.astore(first + i); } - code.aload(1 + graph.nodes().size()).areturn(); + code.aload(first + graph.nodes().size() - 1); } } diff --git a/numjvm-vector/src/main/java/org/numjvm/vector/MixedFusedKernel.java b/numjvm-vector/src/main/java/org/numjvm/vector/MixedFusedKernel.java index dbcd4a1c..2d7548a5 100644 --- a/numjvm-vector/src/main/java/org/numjvm/vector/MixedFusedKernel.java +++ b/numjvm-vector/src/main/java/org/numjvm/vector/MixedFusedKernel.java @@ -23,19 +23,9 @@ final class MixedFusedKernel implements FusedKernel { var floats = MixedConversions.FLOAT; Object[] values = new Object[inputTypes.size()]; long bound = doubles.loopBound(count), i = 0; - for (; i < bound; i += doubles.length()) { - for (int slot = 0; slot < values.length; slot++) { - if (inputTypes.get(slot) == DType.FLOAT32) { - values[slot] = strides[slot] == 0 - ? FloatVector.broadcast(floats, segments[slot].get(ValueLayout.JAVA_FLOAT_UNALIGNED, offsets[slot])) - : FloatVector.fromMemorySegment(floats, segments[slot], offsets[slot] + i * strides[slot], ByteOrder.nativeOrder()); - } else { - values[slot] = strides[slot] == 0 - ? DoubleVector.broadcast(doubles, segments[slot].get(ValueLayout.JAVA_DOUBLE_UNALIGNED, offsets[slot])) - : DoubleVector.fromMemorySegment(doubles, segments[slot], offsets[slot] + i * strides[slot], ByteOrder.nativeOrder()); - } - } - expression.evaluate(values).intoMemorySegment(segments[values.length], offsets[values.length] + i * strides[values.length], ByteOrder.nativeOrder()); + if (bound > 0) { + ((VectorLoop) expression).runVectors(segments, offsets, strides, bound, doubles.length()); + i = bound; } for (; i < count; i++) { for (int slot = 0; slot < values.length; slot++) { diff --git a/numjvm-vector/src/main/java/org/numjvm/vector/VectorExpressionCompiler.java b/numjvm-vector/src/main/java/org/numjvm/vector/VectorExpressionCompiler.java index 2e691d5f..9bcb4c5d 100644 --- a/numjvm-vector/src/main/java/org/numjvm/vector/VectorExpressionCompiler.java +++ b/numjvm-vector/src/main/java/org/numjvm/vector/VectorExpressionCompiler.java @@ -34,11 +34,16 @@ static Object compile(FusedGraph graph, DType type) { ClassDesc name = ClassDesc.of("org.numjvm.vector.Generated" + kind + "Expression"); ClassDesc contract = ClassDesc.of("org.numjvm.vector." + kind + "Expression"); byte[] bytes = ClassFile.of().build(name, builder -> builder.withFlags(ClassFile.ACC_FINAL | ClassFile.ACC_SUPER) - .withInterfaceSymbols(contract) + .withInterfaceSymbols(contract, ClassDesc.of("org.numjvm.vector.VectorLoop")) .withMethodBody("", MethodTypeDesc.of(ConstantDescs.CD_void), 0, code -> code.aload(0).invokespecial(ConstantDescs.CD_Object, "", MethodTypeDesc.of(ConstantDescs.CD_void)).return_()) .withMethodBody("evaluate", MethodTypeDesc.of(vector, vector.arrayType()), ClassFile.ACC_PUBLIC, - code -> emit(code, graph, vector, single))); + code -> { emit(code, graph, vector, single, 2, null, null, null); code.areturn(); }) + .withMethodBody("runVectors", MethodTypeDesc.of(ConstantDescs.CD_void, + ClassDesc.of("java.lang.foreign.MemorySegment").arrayType(), + ConstantDescs.CD_long.arrayType(), ConstantDescs.CD_long.arrayType(), + ConstantDescs.CD_long, ConstantDescs.CD_int), ClassFile.ACC_PUBLIC, + code -> emitLoop(code, graph, vector, single))); try { var lookup = MethodHandles.lookup().defineHiddenClass(bytes, true); return lookup.findConstructor(lookup.lookupClass(), MethodType.methodType(void.class)).invoke(); @@ -46,11 +51,40 @@ static Object compile(FusedGraph graph, DType type) { catch (Throwable error) { throw new IllegalStateException("cannot define fused vector expression", error); } } - private static void emit(CodeBuilder code, FusedGraph graph, ClassDesc vector, boolean single) { + private static void emitLoop(CodeBuilder code, FusedGraph graph, ClassDesc vector, boolean single) { + int count = graph.inputTypes().size() + 1; + int[] segments = new int[count], offsets = new int[count], strides = new int[count]; + int local = 9; + for (int i = 0; i < count; i++) { + segments[i] = local++; offsets[i] = local; local += 2; strides[i] = local; local += 2; + code.aload(1).loadConstant(i).aaload().astore(segments[i]); + code.aload(2).loadConstant(i).laload().lstore(offsets[i]); + code.aload(3).loadConstant(i).laload().lstore(strides[i]); + } + var loop = code.newLabel(); + var done = code.newLabel(); + code.lconst_0().lstore(7).labelBinding(loop).lload(7).lload(4).lcmp().ifge(done); + code.aload(segments[count - 1]).lload(offsets[count - 1]); + emit(code, graph, vector, single, local, segments, offsets, strides); + code.invokestatic(ClassDesc.of("org.numjvm.vector.VectorMemory"), single ? "storeFloat" : "storeDouble", + MethodTypeDesc.of(ConstantDescs.CD_void, ClassDesc.of("java.lang.foreign.MemorySegment"), ConstantDescs.CD_long, vector)); + for (int i = 0; i < count; i++) + code.lload(offsets[i]).lload(strides[i]).iload(6).i2l().lmul().ladd().lstore(offsets[i]); + code.lload(7).iload(6).i2l().ladd().lstore(7).goto_(loop).labelBinding(done).return_(); + } + + private static void emit(CodeBuilder code, FusedGraph graph, ClassDesc vector, boolean single, + int first, int[] segments, int[] offsets, int[] strides) { for (int i = 0; i < graph.nodes().size(); i++) { var node = graph.nodes().get(i); switch (node.op()) { - case INPUT -> code.aload(1).loadConstant(node.slot()).aaload(); + case INPUT -> { + if (segments == null) code.aload(1).loadConstant(node.slot()).aaload(); + else code.aload(segments[node.slot()]).lload(offsets[node.slot()]).lload(strides[node.slot()]) + .invokestatic(ClassDesc.of("org.numjvm.vector.VectorMemory"), single ? "loadFloat" : "loadDouble", + MethodTypeDesc.of(vector, ClassDesc.of("java.lang.foreign.MemorySegment"), + ConstantDescs.CD_long, ConstantDescs.CD_long)); + } case CONSTANT -> { code.getstatic(vector, "SPECIES_PREFERRED", SPECIES); if (single) { @@ -67,18 +101,18 @@ private static void emit(CodeBuilder code, FusedGraph graph, ClassDesc vector, b case ADD, SUBTRACT, MULTIPLY, DIVIDE -> { String method = switch (node.op()) { case ADD -> "add"; case SUBTRACT -> "sub"; case MULTIPLY -> "mul"; default -> "div"; }; - code.aload(2 + node.left()).aload(2 + node.right()) + code.aload(first + node.left()).aload(first + node.right()) .invokevirtual(vector, method, MethodTypeDesc.of(vector, BASE_VECTOR)); } case NEGATIVE, ABSOLUTE, SQRT -> { String method = node.op() == FusedGraph.Op.NEGATIVE ? "neg" : node.op() == FusedGraph.Op.ABSOLUTE ? "abs" : "sqrt"; - code.aload(2 + node.left()).invokevirtual(vector, method, MethodTypeDesc.of(vector)); + code.aload(first + node.left()).invokevirtual(vector, method, MethodTypeDesc.of(vector)); } - case SQUARE -> code.aload(2 + node.left()).aload(2 + node.left()) + case SQUARE -> code.aload(first + node.left()).aload(first + node.left()) .invokevirtual(vector, "mul", MethodTypeDesc.of(vector, BASE_VECTOR)); } - code.astore(2 + i); + code.astore(first + i); } - code.aload(1 + graph.nodes().size()).areturn(); + code.aload(first + graph.nodes().size() - 1); } } diff --git a/numjvm-vector/src/main/java/org/numjvm/vector/VectorLoop.java b/numjvm-vector/src/main/java/org/numjvm/vector/VectorLoop.java new file mode 100644 index 00000000..99321868 --- /dev/null +++ b/numjvm-vector/src/main/java/org/numjvm/vector/VectorLoop.java @@ -0,0 +1,8 @@ +package org.numjvm.vector; + +import java.lang.foreign.MemorySegment; + +/** A generated loop for the whole vector portion of a contiguous run. */ +interface VectorLoop { + void runVectors(MemorySegment[] segments, long[] offsets, long[] strides, long count, int lanes); +} diff --git a/numjvm-vector/src/main/java/org/numjvm/vector/VectorMemory.java b/numjvm-vector/src/main/java/org/numjvm/vector/VectorMemory.java new file mode 100644 index 00000000..a18b3dc9 --- /dev/null +++ b/numjvm-vector/src/main/java/org/numjvm/vector/VectorMemory.java @@ -0,0 +1,33 @@ +package org.numjvm.vector; + +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.nio.ByteOrder; +import jdk.incubator.vector.DoubleVector; +import jdk.incubator.vector.FloatVector; + +/** Small helpers used by generated loops; vector values never enter an array. */ +final class VectorMemory { + private VectorMemory() { } + static FloatVector loadFloat(MemorySegment segment, long offset, long stride) { + var species = FloatVector.SPECIES_PREFERRED; + return stride == 0 ? FloatVector.broadcast(species, segment.get(ValueLayout.JAVA_FLOAT_UNALIGNED, offset)) + : FloatVector.fromMemorySegment(species, segment, offset, ByteOrder.nativeOrder()); + } + static FloatVector loadMixedFloat(MemorySegment segment, long offset, long stride) { + var species = MixedConversions.FLOAT; + return stride == 0 ? FloatVector.broadcast(species, segment.get(ValueLayout.JAVA_FLOAT_UNALIGNED, offset)) + : FloatVector.fromMemorySegment(species, segment, offset, ByteOrder.nativeOrder()); + } + static DoubleVector loadDouble(MemorySegment segment, long offset, long stride) { + var species = DoubleVector.SPECIES_PREFERRED; + return stride == 0 ? DoubleVector.broadcast(species, segment.get(ValueLayout.JAVA_DOUBLE_UNALIGNED, offset)) + : DoubleVector.fromMemorySegment(species, segment, offset, ByteOrder.nativeOrder()); + } + static void storeFloat(MemorySegment segment, long offset, FloatVector value) { + value.intoMemorySegment(segment, offset, ByteOrder.nativeOrder()); + } + static void storeDouble(MemorySegment segment, long offset, DoubleVector value) { + value.intoMemorySegment(segment, offset, ByteOrder.nativeOrder()); + } +} diff --git a/release-check-modular/build.gradle.kts b/release-check-modular/build.gradle.kts index 9767cffe..a8167de2 100644 --- a/release-check-modular/build.gradle.kts +++ b/release-check-modular/build.gradle.kts @@ -2,11 +2,12 @@ plugins { application } -val numjvmVersion = providers.gradleProperty("numjvmVersion").orElse("0.5.0-alpha") +val numjvmVersion = providers.gradleProperty("numjvmVersion").orElse("0.6.0-alpha") dependencies { implementation(numjvmVersion.map { "io.github.stamich:numjvm-io:$it" }) implementation(numjvmVersion.map { "io.github.stamich:numjvm-linalg:$it" }) + implementation(numjvmVersion.map { "io.github.stamich:numjvm-jit:$it" }) if (providers.gradleProperty("withVector").isPresent || providers.gradleProperty("unresolvedVector").isPresent) { runtimeOnly(numjvmVersion.map { "io.github.stamich:numjvm-vector:$it" }) } diff --git a/release-check-modular/src/main/java/module-info.java b/release-check-modular/src/main/java/module-info.java index 5461bc4b..87ee1c11 100644 --- a/release-check-modular/src/main/java/module-info.java +++ b/release-check-modular/src/main/java/module-info.java @@ -1,4 +1,5 @@ module numjvm.releasecheck { requires org.numjvm.io; requires org.numjvm.linalg; + requires org.numjvm.jit; } diff --git a/release-check/build.gradle.kts b/release-check/build.gradle.kts index 4f96ff95..9fc7640b 100644 --- a/release-check/build.gradle.kts +++ b/release-check/build.gradle.kts @@ -2,11 +2,12 @@ plugins { application } -val numjvmVersion = providers.gradleProperty("numjvmVersion").orElse("0.5.0-alpha") +val numjvmVersion = providers.gradleProperty("numjvmVersion").orElse("0.6.0-alpha") dependencies { implementation(numjvmVersion.map { "io.github.stamich:numjvm-io:$it" }) implementation(numjvmVersion.map { "io.github.stamich:numjvm-linalg:$it" }) + implementation(numjvmVersion.map { "io.github.stamich:numjvm-jit:$it" }) if (providers.gradleProperty("withVector").isPresent || providers.gradleProperty("unresolvedVector").isPresent) { runtimeOnly(numjvmVersion.map { "io.github.stamich:numjvm-vector:$it" }) } diff --git a/release-check/src/main/java/smoke/ReleaseCheck.java b/release-check/src/main/java/smoke/ReleaseCheck.java index 4b35e9fb..6fe063f8 100644 --- a/release-check/src/main/java/smoke/ReleaseCheck.java +++ b/release-check/src/main/java/smoke/ReleaseCheck.java @@ -32,6 +32,14 @@ public static void main(String[] args) throws Exception { } NDArray product = org.numjvm.linalg.Linalg.matmul(a, a.transpose()); if (product.getDouble(0, 1) != 32.0) { throw new AssertionError("matmul failed"); } + try (var pool = NJ.parallel(2); var jit = org.numjvm.jit.Jit.context(org.numjvm.jit.Jit.Mode.AUTO)) { + var input = NJ.ones(org.numjvm.dtype.DType.FLOAT32, 65536); + var output = NJ.zeros(org.numjvm.dtype.DType.FLOAT32, 65536); + var expression = jit.trace(1, i -> i.input(0).add(2).square()); + expression.evalInto(pool, output, input); + if (output.getDouble(65535) != 9 || pool.sum(output).getDouble() != 589824) + throw new AssertionError("parallel expression or sum failed"); + } if (Boolean.getBoolean("expectVector")) { long before = org.numjvm.BackendInfo.vectorRuns(); NDArray vectorInput = NJ.ones(257);