MemorySegment plugin and standalone usage of MemorySegment

[andrii0lomakin]

Good day.

I found out that if I use Vector API and TornadoVM, the last one stops working if I use it after the first.
My conclusions are preliminary, but nonetheless, I would ask you to confirm them or deny them.

My flow of execution is the following.

1. I use Vector API that uses MemorySegment underneath.
2. I call TornadoVM part and plugin is not called. I suppose because of optimizations to avoid to resolve plugins every time. I am especially curious about the case when code is jitted (but I am pretty sure that in my case code was not jitted because I checked that interpreter in reality executes MemorySegment byte code during execution of TVM kernel though MemorySegment resolution for MemorySegment methods was not called, though code can be deoptimized of course at this moment).
3. Kernels are not executed correctly.

That raises several questions:

1. Is it possible to use MemorySegment for native memory access once it is extended by TorandoVM ?
2. Is it possible to ensure the correct work of TornadoVM with the Vector API part that uses MemorySegments ? I find this part important because the application can use libraries that work with Vector API.

[mikepapadim]

Hello @andrii0lomakin ,

Thank you for the information. It is possible to share a reproducer of the work-flow that triggers this incosistency?

I used extensively TornadoVM along with the VectorAPI in this LLM inference example without having any data incosistencies https://github.com/mikepapadim/llama2.tornadovm.java

Also, how do you call TornadoVM with the MemorySegment?
Passing a MemorySegment directly to a TaskGraph is not currently supported.
It needs to be wrapped into one of the TornadoNativeTypes. You can create a TornadoNative from a MemorySegment and then pass it to the TaskGraph.

Did you have a look here: https://github.com/beehive-lab/TornadoVM/blob/master/tornado-unittests/src/main/java/uk/ac/manchester/tornado/unittests/vector/api/TestVectorAPI.java

[andrii0lomakin]

It needs to be wrapped into one of the TornadoNativeTypes. You can create a TornadoNative from a MemorySegment and then pass it to the >TaskGraph.

That is exactly how I did. But Float(Int)Array are accessed through MemorySegment that is compiled to GPU code, is it correct ?

Thank you for the information. It is possible to share a reproducer of the work-flow that triggers this incosistency?

That is quite easy to do, I will create branch where I enable my unit tests related to VectorAPI.

[mikepapadim]

It needs to be wrapped into one of the TornadoNativeTypes. You can create a TornadoNative from a MemorySegment and then pass it to the >TaskGraph.

That is exactly how I did. But Float(Int)Array are accessed through MemorySegment that is compiled to GPU code, is it correct ?
If you can share the code snippet it would be great so I can have a look direclty into it.
One thing to double check is that the API for each type has two different methods to get the segment (getSegmentWithHeader and getSegment).

Thank you for the information. It is possible to share a reproducer of the work-flow that triggers this incosistency?

That is quite easy to do, I will create branch where I enable my unit tests related to VectorAPI.

[andrii0lomakin]

@mikepapadim could you confirm that the usage of plugins for MemorySegment does not prevent its usage in Vector API or directly
https://github.com/beehive-lab/TornadoVM/blob/master/tornado-drivers/opencl/src/main/java/uk/ac/manchester/tornado/drivers/opencl/graal/compiler/plugins/OCLGraphBuilderPlugins.java#L408 ?

[mikepapadim]

Good catch!
I just added a test for that.
We should throw an exception to prevent it's usage directly.

Actually, we are currently redesign this part of the plugins while preparing for JDK22 and Graal 24.0.1, so we are going to handle it approapriatly there.

[andrii0lomakin]

@mikepapadim Probably I expressed myself incorrectly.
I have debugged the current implementation, and from what I see, there is no plugin for FloatArray and IntArray, at least for memory access and memory access is handled directly by MemorySegment plugin(s).

Which why I asked if altering of byte code in those plugins can prevent to work with MemorySegment in other cases?

[andrii0lomakin]

If that is true (which I do not claim but ask), do not you find it "dangerous" to change byte code generation for it? Because JDK libraries can rely on the internals of implementation of those classes.

[mikepapadim]

If that is true (which I do not claim but ask), do not you find it "dangerous" to change byte code generation for it? Because JDK libraries can rely on the internals of implementation of those classes.

Actually, for the JDK22 port we have a proxy class for the MemorySegments and we register/handle the plugins for accesses only through this proxy class

[andrii0lomakin]

Actually, for the JDK22 port we have a [proxy class for the MemorySegments

@mikepapadim you read my mind :-))
My exact quotation to my colleague - "They can use wrapper over memory segment to avoid this problems but still support work in case of absence of devices to use"
Thank you.

Do you have any ETA for this change ?

[mikepapadim]

Actually, for the JDK22 port we have a [proxy class for the MemorySegments

@mikepapadim you read my mind :-)) My exact quotation to my colleague - "They can use wrapper over memory segment to avoid this problems but still support work in case of absence of devices to use" Thank you.

Do you have any ETA for this change ?

Currently, we made this change as part of the JDK22 port. We havent discuss yet if we are going to backport in the current version of Tornado. @jjfumero