Optimize tensor.slice() (#1381)

* Optimize tensor.slice()

The performance of executing `tensor.slice()` is super poor, especially for
the 'logits' tensor with large dimensions.

```
const logits = outputs.logits.slice(null, -1, null);`
```

This is because currently implementation of the `slice` method manually iterates
through each element and calculate indices which is a big time consuming if
the tensor shape is large.

For cases like `slice(null, -1, null)`, where the slicing operation is
contiguous along certain dimensions, which can be optimized by bulk copy
by using `TypeArray.subarray()` and `TypeArray.set()`.

* nit

* Add a few more tensor slice unit tests

---------

Co-authored-by: Joshua Lochner <[email protected]>

Author: Honry

src/utils/tensor.js

@@ -443,15 +443,46 @@ export class Tensor {
         // Precompute strides
         const stride = this.stride();
 
-        for (let i = 0; i < newBufferSize; ++i) {
-            let originalIndex = 0;
-            for (let j = newDims.length - 1, num = i; j >= 0; --j) {
-                const size = newDims[j];
-                originalIndex += ((num % size) + newOffsets[j][0]) * stride[j];
-                num = Math.floor(num / size);
+        // Detect if the slice is contiguous
+        let isContiguous = true;
+        for (let i = 1; i < newDims.length; ++i) {
+            if (newOffsets[i][0] !== 0 || newOffsets[i][1] !== this.dims[i]) {
+                isContiguous = false;
+                break;
             }
-            data[i] = this_data[originalIndex];
         }
+
+        if (isContiguous) {
+            // Perform bulk copy for contiguous slices to improve performance
+            const start = newOffsets[0][0] * stride[0];
+            const end = newOffsets[0][1] * stride[0];
+
+            if (ArrayBuffer.isView(this_data)) {
+                // If this.data is a TypedArray, use subarray
+                // @ts-ignore
+                data.set(this_data.subarray(start, end));
+            } else if (Array.isArray(this_data)) {
+                // If this.data is a plain array, use slice
+                const slicedData = this_data.slice(start, end);
+                for (let i = 0; i < slicedData.length; ++i) {
+                    data[i] = slicedData[i];
+                }
+            } else {
+                throw new Error("Unsupported data type for slicing");
+            }
+        } else {
+            // Fallback to manual copying for non-contiguous slices
+            for (let i = 0; i < newBufferSize; ++i) {
+                let originalIndex = 0;
+                for (let j = newDims.length - 1, num = i; j >= 0; --j) {
+                    const size = newDims[j];
+                    originalIndex += ((num % size) + newOffsets[j][0]) * stride[j];
+                    num = Math.floor(num / size);
+                }
+                data[i] = this_data[originalIndex];
+            }
+        }
+
         return new Tensor(this.type, data, newTensorDims);
     }
 

tests/utils/tensor.test.js

@@ -59,15 +59,13 @@ describe("Tensor operations", () => {
       const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
       const t2 = t1.slice(1);
       const target = new Tensor("float32", [3, 4], [2]);
-
       compare(t2, target);
     });
 
     it("should return a range of rows", () => {
       const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
       const t2 = t1.slice([1, 3]);
       const target = new Tensor("float32", [3, 4, 5, 6], [2, 2]);
-
       compare(t2, target);
     });
 
@@ -78,9 +76,67 @@ describe("Tensor operations", () => {
         [4, 7],
       );
       const t2 = t1.slice([1, -1], [1, -1]);
-
       const target = new Tensor("float32", [9, 10, 11, 12, 13, 16, 17, 18, 19, 20], [2, 5]);
+      compare(t2, target);
+    });
+
+    it("should return the whole tensor when all indices are null/unset", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice();
+      compare(t2, t1);
+    });
+
+    it("should return the whole dimension when index is null", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice(null);
+      compare(t2, t1);
+    });
+
+    it("should slice from index to end when [start, null] is used", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice([1, null]);
+      const target = new Tensor("float32", [3, 4, 5, 6], [2, 2]);
+      compare(t2, target);
+    });
+
+    it("should slice from beginning to index when [null, end] is used", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice([null, 2]);
+      const target = new Tensor("float32", [1, 2, 3, 4], [2, 2]);
+      compare(t2, target);
+    });
+
+    it("should handle [null, null] as full slice", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice([null, null]);
+      compare(t2, t1);
+    });
+
+    it("should select a single element when a number is used in slice", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice(2, 1);
+      const target = new Tensor("float32", [6], []);
+      compare(t2, target);
+    });
 
+    it("should select a single row when a number is used in slice", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice(0);
+      const target = new Tensor("float32", [1, 2], [2]);
+      compare(t2, target);
+    });
+
+    it("should select a single column when a number is used in slice", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice(null, 1);
+      const target = new Tensor("float32", [2, 4, 6], [3]);
+      compare(t2, target);
+    });
+
+    it("should handle negative indices in slice", () => {
+      const t1 = new Tensor("float32", [1, 2, 3, 4, 5, 6], [3, 2]);
+      const t2 = t1.slice(-1);
+      const target = new Tensor("float32", [5, 6], [2]);
       compare(t2, target);
     });
   });