Hybrid API: Mix JIT Kernels with Native GPU Libraries

public final class MyLib {
    public static final String LIBRARY_NAME = "vendor/mylib";

    public static LibraryTaskDescriptor myOp(int n, FloatArray in, FloatArray out) {
        Access[] access = { Access.READ_ONLY, Access.READ_ONLY, Access.WRITE_ONLY };
        return new LibraryTaskDescriptor()
            .withLibrary(LIBRARY_NAME)
            .withFunction("myOp")
            .withParameters(new Object[] { n, in, out })
            .withAccess(access);
    }
}

public final class MyProvider implements TornadoLibraryProvider {
    public String libraryName() { return MyLib.LIBRARY_NAME; }

    public boolean canHandle(TornadoXPUDevice device) {
        return device instanceof TornadoNativeStreamSupport;
    }

    public LibraryContext createContext(TornadoXPUDevice device, long planId) {
        long stream = ((TornadoNativeStreamSupport) device).getNativeStream(planId);
        return new MyContext(MyNativeLib.createHandle(), stream);
    }

    public void prepare(LibraryTaskDescriptor d, LibraryContext ctx) {
        // Idempotent: create/cache per-shape plans + workspace BEFORE capture
    }

    public void dispatch(String fn, LibraryInvocation call) {
        long dIn  = call.getDevicePointer(1);   // reference args → raw device pointers
        long dOut = call.getDevicePointer(2);
        int  n    = (int) call.getArg(0);        // scalars → boxed values
        MyNativeLib.myOp(((MyContext) call.getContext()).handle, n, dIn, dOut);
    }

    public void destroyContext(LibraryContext ctx) { /* free handle, plans, workspace */ }
}

open module tornado.mylib {
    requires transitive tornado.api;
    requires tornado.runtime;
    exports vendor.mylib;
    provides uk.ac.manchester.tornado.runtime.library.spi.TornadoLibraryProvider
        with vendor.mylib.provider.MyProvider;
}