Add method with strides

Author: mayya-sharipova

x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/plugin/gpu/CuVSProviderDelegate.java

@@ -38,6 +38,17 @@ public CuVSMatrix.Builder<CuVSHostMatrix> newHostMatrixBuilder(long l, long l1,
         return delegate.newHostMatrixBuilder(l, l1, dataType);
     }
 
+    @Override
+    public CuVSMatrix.Builder<CuVSHostMatrix> newHostMatrixBuilder(
+        long size,
+        long columns,
+        int rowStride,
+        int columnStride,
+        CuVSMatrix.DataType dataType
+    ) {
+        return delegate.newHostMatrixBuilder(size, columns, rowStride, columnStride, dataType);
+    }
+
     @Override
     public CuVSMatrix.Builder<CuVSDeviceMatrix> newDeviceMatrixBuilder(
         CuVSResources cuVSResources,
@@ -65,6 +76,11 @@ public MethodHandle newNativeMatrixBuilder() {
         return delegate.newNativeMatrixBuilder();
     }
 
+    @Override
+    public MethodHandle newNativeMatrixBuilderWithStrides() {
+        return delegate.newNativeMatrixBuilderWithStrides();
+    }
+
     @Override
     public CuVSMatrix newMatrixFromArray(float[][] floats) {
         return delegate.newMatrixFromArray(floats);

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/DatasetUtils.java

@@ -22,6 +22,15 @@ static DatasetUtils getInstance() {
     /** Returns a Dataset over the vectors of type {@code dataType} in the input. */
     CuVSMatrix fromInput(MemorySegmentAccessInput input, int numVectors, int dims, CuVSMatrix.DataType dataType) throws IOException;
 
+    CuVSMatrix fromInput(
+        MemorySegmentAccessInput input,
+        int numVectors,
+        int dims,
+        int rowStride,
+        int columnStride,
+        CuVSMatrix.DataType dataType
+    ) throws IOException;
+
     /** Returns a Dataset over an input slice */
     CuVSMatrix fromSlice(MemorySegmentAccessInput input, long pos, long len, int numVectors, int dims, CuVSMatrix.DataType dataType)
         throws IOException;

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/DatasetUtilsImpl.java

@@ -21,6 +21,7 @@ public class DatasetUtilsImpl implements DatasetUtils {
     private static final DatasetUtils INSTANCE = new DatasetUtilsImpl();
 
     private static final MethodHandle createDataset$mh = CuVSProvider.provider().newNativeMatrixBuilder();
+    private static final MethodHandle createDatasetWithStrides$mh = CuVSProvider.provider().newNativeMatrixBuilderWithStrides();
 
     static DatasetUtils getInstance() {
         return INSTANCE;
@@ -40,6 +41,27 @@ static CuVSMatrix fromMemorySegment(MemorySegment memorySegment, int size, int d
         }
     }
 
+    static CuVSMatrix fromMemorySegment(
+        MemorySegment memorySegment,
+        int size,
+        int dimensions,
+        int rowStride,
+        int columnStride,
+        CuVSMatrix.DataType dataType
+    ) {
+        try {
+            return (CuVSMatrix) createDatasetWithStrides$mh.invokeExact(memorySegment, size, dimensions, rowStride, columnStride, dataType);
+        } catch (Throwable e) {
+            if (e instanceof Error err) {
+                throw err;
+            } else if (e instanceof RuntimeException re) {
+                throw re;
+            } else {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
     private DatasetUtilsImpl() {}
 
     @Override
@@ -50,6 +72,21 @@ public CuVSMatrix fromInput(MemorySegmentAccessInput input, int numVectors, int
         return createCuVSMatrix(input, 0L, input.length(), numVectors, dims, dataType);
     }
 
+    @Override
+    public CuVSMatrix fromInput(
+        MemorySegmentAccessInput input,
+        int numVectors,
+        int dims,
+        int rowStride,
+        int columnStride,
+        CuVSMatrix.DataType dataType
+    ) throws IOException {
+        if (numVectors < 0 || dims < 0) {
+            throwIllegalArgumentException(numVectors, dims);
+        }
+        return createCuVSMatrix(input, 0L, input.length(), numVectors, dims, rowStride, columnStride, dataType);
+    }
+
     @Override
     public CuVSMatrix fromSlice(MemorySegmentAccessInput input, long pos, long len, int numVectors, int dims, CuVSMatrix.DataType dataType)
         throws IOException {
@@ -76,6 +113,24 @@ private static CuVSMatrix createCuVSMatrix(
         return fromMemorySegment(ms, numVectors, dims, dataType);
     }
 
+    private static CuVSMatrix createCuVSMatrix(
+        MemorySegmentAccessInput input,
+        long pos,
+        long len,
+        int numVectors,
+        int dims,
+        int rowStride,
+        int columnStride,
+        CuVSMatrix.DataType dataType
+    ) throws IOException {
+        MemorySegment ms = input.segmentSliceOrNull(pos, len);
+        assert ms != null;
+        if (((long) numVectors * rowStride) > ms.byteSize()) {
+            throwIllegalArgumentException(ms, numVectors, dims);
+        }
+        return fromMemorySegment(ms, numVectors, dims, rowStride, columnStride, dataType);
+    }
+
     static void throwIllegalArgumentException(MemorySegment ms, int numVectors, int dims) {
         var s = "segment of size [" + ms.byteSize() + "] too small for expected " + numVectors + " float vectors of " + dims + " dims";
         throw new IllegalArgumentException(s);