Java based Numerical library (JNum-v0.1)

[u/CutGroundbreaking305]

previous post

And here I am, made a Java-based numerical library called JNum.

I used the new FFM API and Vector API (Project Panama) to make it 100% pure Java, unlike ND4J which relies heavily on JNI and massive C++ backends. Here is the repo: https://github.com/CH-Abhinav/JNum . It is currently in a v0.1 (PREVIEW).

Some of you may ask: Isn't the Vector API still in incubator? Yeah, even though it's still in incubation I preferred to continue building with it as it doesn't have any major API changes planned except the inclusion of value classes (hopium it is coming in Java 27 🙃).

The Performance so far: By avoiding the JNI crossover latency, the basic math tasks (add, mul) are actually faster compared to ND4J and NumPy on small/medium arrays.

The main wins are the reduction methods (sum, max, min) which are about 2x faster compared to ND4J.

Because there is no native C++ backend, the entire library is under 100KB, compared to the hundreds of megabytes required to bundle native binaries.

The Matmul Struggle: Obviously, the main talking point for tensor engines is matmul. Not gonna lie, this ate my brain while trying to figure out which memory settings and SIMD loops work best. Right now, a 1024x1024 float matrix multiplication takes about ~51ms. It's fast, but we still haven't reached the massive performance of ND4J or NumPy on huge matrices (I haven't implemented multi-threading or L1/L2 cache tiling yet).

Use case (potential): ND4J is bulky, and when making applications (web or Android) which require some sort of math and performance, Java devs need to bundle that bulky dependency. We can run JNum anywhere as it doesn't have any .dll or .so files, nor JNI—just pure Java.

I guess this project will become more like multik but better and javaish. And I'm expecting ML guys in Java can also use it (though ND4J/DJL is better for now).

I want the Java community to help me build this project! I am still learning the deeper JVM optimizations(stylish way of saying i am newbie), so if anyone has experience with SIMD loop unrolling, cache tiling or anything helpful I'd love some code reviews, advice, or PRs and help this fellow java guy.

[u/agibsonccc]

Hey! Nd4j maintainer here. There's a fairly large rewrite going on here attempting to address that. I actually agree with you! Not to dunk on you here but we tried your approach more than a decade ago.

Pure java is just not going to be a performant runtime for numercial software even *WITH* panama. You'll never have access to the low level gpu runtimes from the mobile vendors for android. You also won't be able to benefit from many of the low level optimizations that c++ compilers just innately offer without working around the runtime.

Broadly, GC runtimes are just NOT worth it.

I will be publishing a slimmer deployment focused binary to tackle this while also addressing the small matrices overhead. We mainly built nd4j for deep learning so small matrices were far and few between. The way the kernels are written it unfortunately means threading overhead among other things.

I won't try to sell you on cooperating, nor on discouraging you from trying this. User choice matters.

I get wanting to do your own thing and hope it succeeds.

I'll keep an eye on feedback. I hope you carve out a niche for yourself good luck!

[u/CutGroundbreaking305 → u/agibsonccc]

its nd4j dev himself 🙏

I way seeing how Nd4j/DLJ and were doing. I completely agree that c++ based lib will always be better than java based. But the better question would be calling c++ code into java via JNI/FFM is better than just running java based code? And some cases c++ is better but in other cases java is. At least that's what i learnt while i was making my project. I agree with GC runtimes issues but off heap memory via FFM and potential vector api being value classes could reduce that a bit.

I will be grateful to help in nd4j if I can. May be you can try out hybrid approach of pure java + c++ backed java in nd4j instead of entirely depending on c++ itself. This will make things slimmer and better. And deprecation of Unsafe and FFM/FFI introduction I guess you guys need to revamp things. In this cases, I can definitely help you in nd4j/dlj. But I will continue my journey on the pure java front(till i hit the wall i guess).

And instead of supporting just cuda based gpu frameworks you guys can use webgpu instead. idk about exacts but i guess it will cover every gpu architecture instead of single nvidia based cuda.

[u/agibsonccc → u/CutGroundbreaking305]

You don't need to help! You have your own opinion on how to solve the problems. You have your own goals. My point was that it's just not a focus of the framework. I just strongly don't believe java itself will keep up. I'd rather java be an interface language like how python does it.

You may disagree with me and that's fine! Give it a shot!

In my experience, the main bottleneck now a days is the following:

1. JNI calls are expensive. That is why small matrices are hard to do well. Batching normally needs to happen. I totally get why you'd want everything in pure java for that specific case.

2. Proper threading/simd. Small matrices just don't need that.

Between the overhead of the JNI calls plus the need to spin up openmp thread pools that is fast in most places but not for small matrices we just made that trade off.

It definitely wasn't perfect. It also just didn't matter for deep learning.

Binary size is also a big issue. Needing to include blas libraries inflates library cost a lot.

One thing I've done but haven't mastered how to make generally usable yet is made it so you can pick and choose which op kernels you want to include so we can thin down library size. I have a new minimal backend as a proof of concept that sort of works on that front but I haven't been able to quite get the details right for that. I had to table that for now.

For optimizing for different gpus/tpus and the like I'm actually tackling that! I'm introducing a new compiler framework that actually allows amd gpus, tpus and other things to be used there. There's no reason why web assembly also couldn't be supported.

There's a lot more I can elaborate on here but I'm excited for what the rewrite will be able to do!

[u/CutGroundbreaking305 → u/agibsonccc]

I guess having different opinion is what we need to make completely different architecture.

Regrading Binary size I guess you guys can make architecture based installing i.e. instead of openblas bindings in lib we can install openblas in users system instead or use openblas if user already has one. But I guess this will come with code portability issues.

[u/agibsonccc → u/CutGroundbreaking305]

No there's not really a trade off to make there. The c++ kernel we pick doesn't HAVE to use openblas or even be included at all. The user doesn't have to include matmul at all if they aren't using it in the library.

The point is to allow the user to dynamically slim down the library to only pick what they do/don't use. We also have default implementations of every op kernel there as well.

As I said you have an opinion on how it should be done go for it.

[u/perryplatt → u/agibsonccc]

For ND4j, what would the performance differences be compared to numpy with ffm and openBlas?

[u/agibsonccc → u/perryplatt]

Can you clarify what you're looking for? ND4j like numpy when it uses cpu has a few components: its own kernels for cpu, blas for matmul and other blas methods. It might not be clear but blas does *NOT* do most things. For blas when both use openblas you should see comparable performance.

Beyond that, nd4j for larger problems historically has been more performant for larger tasks because we use multi threaded for loops for most kernels that aren't blas. Numpy has always been about being the baseline.

That's why libraries like pytorch and co largely have superceded it for bigger compute problems in the python ecosystem. Most people don't need that level of performance and numpy just works with no configuration.

[u/International_Break2]

Could you use a openBlas or mkl jextract to try to perform the calculations if they are available?

[u/CutGroundbreaking305 → u/International_Break2]

actually i forgot u can use ffm api to bind native c/cpp code but i was like PURE JAVA!! and didnt think that for a sec .

though idea is good if it creates multiple files like .dll .so to make it run on any os/hardware then it defeats to purpose of not making bulky version .

[u/International_Break2 → u/CutGroundbreaking305]

The bindings could exist in their own jar and be optional. That way performance is available and there is always a fallback. For pure java, you would only need to make sure that the .so is already on the LD path.

[u/CutGroundbreaking305 → u/International_Break2]

for v0.1 i didnt think much when it comes to non java usage. Actually my main aim is to make what u said 2 jars one for just java other for java+openblas+lapack like multik in kotlin does. And i will say LD path idea is great as i dont need to bind natives in my lib by passing that to user's system thanks for that idea.

[u/International_Break2 → u/CutGroundbreaking305]

Hi, the only other idea that I have for you is, could there be a way that one could put in an arena into all of the math functions. That way Confined arenas could be used on these structures.

[u/CutGroundbreaking305 → u/International_Break2]

I didn't get it arenas in math functions? I mean we do sometimes need to create ofshared arenas in matmul but it's just that

And problem with confined arena is regarding threading we can't do multithreading with confined arenas

And rest of the methods don't need internal arenas as much

However if u mean giving add(..,arena) like this where we can give our arena I made better version where u can use a.add(b,resarray) where resarray is ur result ndarray with preferred arena type and a.addi(b) is basically a.add(b,a)

Can u give more clarity on this I didn't get it

[u/ankitkhandelwal6 → u/CutGroundbreaking305]

Why is bulk/size in MB a criteria? I understand the drive for a pure java solution, but not the "bulk reduction" criteria.

[u/CutGroundbreaking305 → u/ankitkhandelwal6]

actually idk

i mean yeah one i didnt want to make something like nd4j which is sometimes 300 mb and when u make an android app 300 mb is already just a library (got to know about nd4j api idk its good i guess) but my main point is pure java based numerical library after that anything else

[u/martinhaeusler]

It's a cool idea, but I'm not sure how "low level" you can go in Java while remaining portable across JVMs and CPU architectures. I think you'll sooner or later hit a point where you need to write a native function to achieve your goals. Numpy is also just a thin python wrapper around a C core library. That being said, people do crazy things on the JVM alone, just look at the top 10 of the 1 Billion Rows challenge.

[u/CutGroundbreaking305 → u/martinhaeusler]

same i am also not sure about how "low level" goes in java . I already hit few wall like operator overloading or cache missing but when i see my few methods being on par or even out perform industry standards like NumPy i feel that we can do things better. i am new to java so i dont know 1 Billion rows challenge but searched about it. I guess experimenting JVM limitations is one thing we can see which shocks us when some times it breaks general assumptions like C is always faster than Java.

[u/martinhaeusler → u/CutGroundbreaking305]

Java can in some cases outperform C because the just-in-time compiler has more information about the runtime behavior and the hardware at hand than any C compiler.

Some general tips:

• Use primitives (int, double, etc). Avoid boxing into wrappers (Integer, Double, etc) like the plague.

• Strongly prefer arrays over collections. Arrays are cumbersome to use but crazy optimized.

• Utilize techniques like SWAR (simd within a register) to batch process multiple numbers with smaller byte length

• You may want to consider looking into the Vector API. It's technically still in incubation status but it's been there for years now.

• Avoid strings. Not sure if this comes up in your library at all, but Double.parseDouble (and related operations) is slooow.

[u/CutGroundbreaking305 → u/martinhaeusler]

I mean entire performance of my code is on JIT 🙏

but i would like to say few things

1. i am not using primitives nor wrappers i am using some glue data types which are used to connect respective method and ffm needs ValueLayout types

2. yeah arrays are better than collections but i am using ffm to allocate off heap memory

3. I am using Vector api i mean its in the post

4. i dont use strings and Double.parseDoube() thanks for that though

[u/arkstack]

This is interesting territory - pure Java numerics on FFM + Vector API is exactly the kind of thing more people should be exploring, and shipping a v0.1 with actual tests and a JMH benchmark already in the repo is more than a lot of first libraries manage. A few observations.

The first thing that stands out is the type-specialization explosion: addFloat/addDouble/addInt * 4 ops * 2 (scalar/array) gives ~24 near-identical method bodies in ArithmaticOps, and the pattern repeats across
ReduceOps/MatMulOps/TrigOps/ExpOps. The natural instinct is "extract an interface and parametrise", but that path is closed in current Java - generics don't cover primitives, and the Vector API itself ships separate
FloatVector/DoubleVector/IntVector for the same reason. So the duplication isn't really a design choice; it's the language until Valhalla lands.

That said, I noticed templates/generate_*.py and the matching *.template.java files. You are generating this. The problem is the generated .java is checked in and the Python isn't wired into Maven, so the template-to-Java contract isn't enforced - somebody can edit ArithmaticOps.java directly and the templates silently drift. Move generation into a Maven exec step, or at least add a CI check that re-runs the scripts and diffs the output. Right now it's a quality gate that exists in principle but not in practice.

A few smaller things:

MemorySegment data, int[] shape, int[] strides are all public final on NDArray. The references are final, but MemorySegment writes through unimpeded and arrays are mutable - arr.shape[0] = 999 compiles and runs. For a lib whose invariants depend on shape/stride consistency, those want to be private with accessors.

MatmulBenchmark only measures your own matmul - the README's "faster than ND4J/NumPy on small/medium arrays" claim has no comparison JMH in the repo to back it. Worth either checking one in or softening the wording.

pom.xml sets source/target to 25 but the README says "Works on Java 22 or higher". Target 25 bytecode won't load on 22 - pick one.

Otherwise this is the right kind of thing to be working on - good luck with it.

[u/CutGroundbreaking305 → u/arkstack]

Thanks for the feedback. This was the reason why I made this as preview version. I was also not sure regrading public finals but i guess i will need to change few things. I will add benchmarks(i did benchmarking but i did some architecture so i needed to remove previous benchmarks). And yeah I will focus on 25 i will change things accordingly.

Thanks for the feed back.

[u/belayon40]

The blis library is a very fast matrix library. I’ve got an ffm wrapper for it already.

https://github.com/boulder-on/jblis

[u/CutGroundbreaking305 → u/belayon40]

Thanks I will definitely check this out (I knew some jdk dev made similar thing idk who)

but the thing is i wanted to go with pure java instead of some c/cpp bindings using vector api

idk how much i can achieve from that but when i will reach a full wall after which i cant do anything then i have no choice but to use something like what you made jblis

[u/VincentxH]

Very cool, beyond theorerical speculation; hope you push Java further than you expect.

[u/CutGroundbreaking305 → u/VincentxH]

thanks 😁

[u/Visual_Brain8809]

As I see it, I think I share your point of view and ultimate goal. I've been trying for quite some time to bring Java up to Python's level for AI applications, but unfortunately, I've encountered many performance issues that have prevented me from continuing my research. If you have any problems in a particular area of ​​your project, please share them, and we can see how we can help.

[u/CutGroundbreaking305 → u/Visual_Brain8809]

I am happy to take help

Make it on par with python libraries will take a lot of time and value classes will make it bit faster overall

Currently it is just adding bunch of methods but try to optimise them is a important part in this especially matrix multiplication which still fails to reach numpy level

[u/quafadas]

Have you considered luhenry‘s fork of netlib for the matmul part?

That falls back to a SIMD matrix multiplication if it can’t JNI to native. I think it also allows for strided representations of matrices which is critical to avoid deep copy / memory bound operations creeping into user code…

[u/CutGroundbreaking305 → u/quafadas]

I didnt know luhenry's fork of netlib but i will definitely check that out

i do use fallback's but since my main method is SIMD matmul i dont have fallbacks for that but generally i avoid deep copy but using non modulo method to do strided fallback for my basic math methods

[u/Mauer_Bluemchen]

You should be careful with Vector API, it still has spurious, non-deterministic "add-bugs" in jdk.incubator.vector.FloatVector on AMD/Intel CPUs.

[u/CutGroundbreaking305 → u/Mauer_Bluemchen]

Oh 😮 I use intel CPU I didn't get any issue maybe I will see this again

Anyways thanks for info

[u/Mauer_Bluemchen → u/CutGroundbreaking305]

It's clearly reproducible, there is a simple test class. Need to finally create the bug report, but currently no time because busy with other projects.

The error is not huge, but add results of identical input data can be different sometimes. Run into this when implementing vector multiplication for my NN library with test cases...

[u/nadrojriajr]

My java inference engine deliverance has matmul and other methods that are implemented using java panama as well as fork join pool to split the work:

https://github.com/edwardcapriolo/deliverance/blob/0bf805a6dbb1fa7555daa9beb0cd00bcc21ccb03/tensor/src/main/java/io/teknek/deliverance/tensor/operations/PanamaTensorOperations.java#L20

The benefit is an implementation that is optimized for many architectures including arm

[u/nadrojriajr]

Check out the gemmers here github.com/edwardcapriolo/deliverance they use panama

[u/MARSHALL_8976]

Nice Project but it has :

1. Unsafe Transposition of Non-Contiguous Matrices : When performing matrix multiplication, the code correctly checks if tensor a is contiguous (a.isContiguous() ? a.data : a.contiguous(arena).data), but it fails to apply this same check to tensor b. It passes b.data directly to fastTranspose2D_Float.

2. Logical Memory Corruption on reshape()

3. Buffer Over-Reads in copy() and equals()

Its a good project.

[u/Own_Opportunity4282 → u/[deleted]]

Can you help me build some apps tech guy :)

[u/MARSHALL_8976 → u/[deleted]]

You seem like a pro tech guy . Do you garbage collectors, Memory allocations , arrays ?