From a96fa5a8795e0b53dcd274b5ac82ea86dbab0c6d Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Wed, 30 Jul 2025 14:53:19 +0800 Subject: [PATCH 1/3] 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()`. --- src/utils/tensor.js | 45 ++++++++++++++++++++++++++++++++++++++------- 1 file changed, 38 insertions(+), 7 deletions(-) diff --git a/src/utils/tensor.js b/src/utils/tensor.js index 357aba9f7..d42e65352 100644 --- a/src/utils/tensor.js +++ b/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); } From 0fcc97f54b61dfda1fa53a1791b085e8e225721c Mon Sep 17 00:00:00 2001 From: Joshua Lochner <26504141+xenova@users.noreply.github.com> Date: Wed, 30 Jul 2025 14:59:57 -0400 Subject: [PATCH 2/3] nit --- src/utils/tensor.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/utils/tensor.js b/src/utils/tensor.js index d42e65352..ea822b6c6 100644 --- a/src/utils/tensor.js +++ b/src/utils/tensor.js @@ -464,7 +464,7 @@ export class Tensor { } 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++) { + for (let i = 0; i < slicedData.length; ++i) { data[i] = slicedData[i]; } } else { From cb3190b73513f6fcdfac9bb564ce717f38703a9a Mon Sep 17 00:00:00 2001 From: Joshua Lochner <26504141+xenova@users.noreply.github.com> Date: Wed, 30 Jul 2025 15:00:57 -0400 Subject: [PATCH 3/3] Add a few more tensor slice unit tests --- tests/utils/tensor.test.js | 62 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 59 insertions(+), 3 deletions(-) diff --git a/tests/utils/tensor.test.js b/tests/utils/tensor.test.js index dacd46789..008f46abb 100644 --- a/tests/utils/tensor.test.js +++ b/tests/utils/tensor.test.js @@ -59,7 +59,6 @@ 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); }); @@ -67,7 +66,6 @@ describe("Tensor operations", () => { 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); }); });