From 54ddd1e30d99f4a38258f928b3ec573db14c4d33 Mon Sep 17 00:00:00 2001 From: Mengwei Liu Date: Mon, 23 Jun 2025 15:07:54 -0700 Subject: [PATCH 1/2] [llm] Add arange() tensor maker API As titled. `arange()` taking a `sizes` argument to be able to take custom tensor sizes. Differential Revision: [D77184741](https://our.internmc.facebook.com/intern/diff/D77184741/) [ghstack-poisoned] --- extension/tensor/tensor_ptr_maker.cpp | 106 ++++++++ extension/tensor/tensor_ptr_maker.h | 42 +++ .../tensor/test/tensor_ptr_maker_test.cpp | 242 ++++++++++++++++++ 3 files changed, 390 insertions(+) diff --git a/extension/tensor/tensor_ptr_maker.cpp b/extension/tensor/tensor_ptr_maker.cpp index 8e7c908bf43..c6e74d46a83 100644 --- a/extension/tensor/tensor_ptr_maker.cpp +++ b/extension/tensor/tensor_ptr_maker.cpp @@ -186,5 +186,111 @@ TensorPtr randint_strided( std::uniform_int_distribution(low, high - 1)); } +TensorPtr arange( + executorch::aten::Scalar start, + executorch::aten::Scalar end, + executorch::aten::Scalar step, + std::vector sizes, + executorch::aten::ScalarType type, + executorch::aten::TensorShapeDynamism dynamism) { + // Calculate the number of elements in the range + double start_val, end_val, step_val; + + if (start.isFloatingPoint()) { + start_val = start.to(); + } else if (start.isIntegral(/*includeBool=*/false)) { + start_val = static_cast(start.to()); + } else { + ET_CHECK_MSG(false, "start must be a number"); + } + + if (end.isFloatingPoint()) { + end_val = end.to(); + } else if (end.isIntegral(/*includeBool=*/false)) { + end_val = static_cast(end.to()); + } else { + ET_CHECK_MSG(false, "end must be a number"); + } + + if (step.isFloatingPoint()) { + step_val = step.to(); + } else if (step.isIntegral(/*includeBool=*/false)) { + step_val = static_cast(step.to()); + } else { + ET_CHECK_MSG(false, "step must be a number"); + } + + ET_CHECK_MSG(step_val != 0, "step cannot be zero"); + + // Calculate the number of elements + int64_t numel = + static_cast(std::ceil((end_val - start_val) / step_val)); + numel = std::max(int64_t(0), numel); // Ensure non-negative + + // Validate sizes compatibility with numel + if (!sizes.empty()) { + int64_t negative_one_count = 0; + int64_t negative_one_index = -1; + int64_t product = 1; + + // Count -1s and calculate product of positive dimensions + for (size_t i = 0; i < sizes.size(); ++i) { + if (sizes[i] == -1) { + negative_one_count++; + negative_one_index = static_cast(i); + } else if (sizes[i] <= 0) { + ET_CHECK_MSG(false, "sizes must contain positive integers or -1"); + } else { + product *= sizes[i]; + } + } + + // Check that there's at most one -1 + ET_CHECK_MSG(negative_one_count <= 1, "sizes can contain at most one -1"); + + if (negative_one_count == 1) { + // Infer the -1 dimension + ET_CHECK_MSG( + numel % product == 0, + "numel (%" PRId64 + ") is not divisible by the product of known dimensions (%" PRId64 ")", + numel, + product); + int64_t inferred_size = numel / product; + ET_CHECK_MSG( + inferred_size > 0, + "inferred dimension size must be positive, got %" PRId64, + inferred_size); + // Update the sizes vector with the inferred dimension + sizes[negative_one_index] = inferred_size; + } else { + // No -1, check exact match + ET_CHECK_MSG( + product == numel, + "product of sizes (%" PRId64 ") does not match numel (%" PRId64 ")", + product, + numel); + } + } + + // Create tensor with the provided sizes or default to 1D + std::vector tensor_sizes = sizes.empty() + ? std::vector{static_cast< + executorch::aten::SizesType>(numel)} + : sizes; + + auto tensor = empty(tensor_sizes, type, dynamism); + + // Fill the tensor with values from start to end with step + ET_SWITCH_REALHBBF16_TYPES(type, nullptr, "arange", CTYPE, [&] { + CTYPE* data = tensor->mutable_data_ptr(); + for (int64_t i = 0; i < numel; ++i) { + data[i] = static_cast(start_val + i * step_val); + } + }); + + return tensor; +} + } // namespace extension } // namespace executorch diff --git a/extension/tensor/tensor_ptr_maker.h b/extension/tensor/tensor_ptr_maker.h index eb3745d34e2..efc23ca6688 100644 --- a/extension/tensor/tensor_ptr_maker.h +++ b/extension/tensor/tensor_ptr_maker.h @@ -683,5 +683,47 @@ inline TensorPtr randint( return randint_strided(low, high, std::move(sizes), {}, type, dynamism); } +/** + * Creates a tensor with values from `start` to `end` (exclusive) with step size + * `step`. This API will error out if `sizes` is not compatible with the number + * of elements of the output tensor. If `sizes` is empty, the result tensor will + * be 1D. + * + * @param start The starting value of the sequence. + * @param end The ending value of the sequence (exclusive). + * @param step The step size between values in the sequence. + * @param sizes A vector specifying the size of each dimension. Only 1 + * occurrence of -1 is allowed, the sizes need to match the number of elements + * in the output tensor. + * @param type The scalar type of the tensor elements. + * @param dynamism Specifies whether the tensor's shape is static or dynamic. + * @return A TensorPtr instance managing the newly created Tensor. + */ +TensorPtr arange( + executorch::aten::Scalar start, + executorch::aten::Scalar end, + executorch::aten::Scalar step = 1, + std::vector sizes = {-1}, + executorch::aten::ScalarType type = executorch::aten::ScalarType::Float, + executorch::aten::TensorShapeDynamism dynamism = + executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND); + +/** + * Creates a 1D tensor (sizes=[max]) with values from 0 to `end` (exclusive) + * with step size 1. + * + * @param end The ending value of the sequence (exclusive). + * @param type The scalar type of the tensor elements. + * @param dynamism Specifies whether the tensor's shape is static or dynamic. + * @return A TensorPtr instance managing the newly created Tensor. + */ +inline TensorPtr arange( + executorch::aten::Scalar end, + executorch::aten::ScalarType type = executorch::aten::ScalarType::Float, + executorch::aten::TensorShapeDynamism dynamism = + executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) { + return arange(0, end, 1, {-1}, type, dynamism); +} + } // namespace extension } // namespace executorch diff --git a/extension/tensor/test/tensor_ptr_maker_test.cpp b/extension/tensor/test/tensor_ptr_maker_test.cpp index e17d18229df..1685447e253 100644 --- a/extension/tensor/test/tensor_ptr_maker_test.cpp +++ b/extension/tensor/test/tensor_ptr_maker_test.cpp @@ -482,3 +482,245 @@ TEST_F(TensorPtrMakerTest, CreateRandnTensorWithIntType) { EXPECT_EQ(val, 0); } } + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithDefaultStartAndStep) { + auto tensor = arange(5); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 5); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_EQ(val, static_cast(i)); + } +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithStartEndStep) { + auto tensor = arange(2, 10, 2); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 4); // (10-2)/2 = 4 elements + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_EQ(val, static_cast(2 + i * 2)); + } +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithNegativeStep) { + auto tensor = arange(5, 0, -1); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 5); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_EQ(val, static_cast(5 - i)); + } +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithIntType) { + auto tensor = arange(0, 5, 1, {-1}, executorch::aten::ScalarType::Int); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 5); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Int); + + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_EQ(val, i); + } +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithLongType) { + auto tensor = arange(0, 5, 1, {-1}, executorch::aten::ScalarType::Long); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 5); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Long); + + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_EQ(val, static_cast(i)); + } +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithDoubleType) { + auto tensor = + arange(0.5, 5.5, 0.5, {-1}, executorch::aten::ScalarType::Double); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 10); // (5.5-0.5)/0.5 = 10 elements + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Double); + + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_DOUBLE_EQ(val, 0.5 + i * 0.5); + } +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithEmptyRange) { + // End < start with positive step should error out + EXPECT_DEATH( + arange(5, 0, 1), + "inferred dimension size must be positive, got 0"); +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithValidSizes) { + // Test arange with compatible sizes + auto tensor = arange(0, 12, 1, {3, 4}); + + EXPECT_EQ(tensor->dim(), 2); + EXPECT_EQ(tensor->size(0), 3); + EXPECT_EQ(tensor->size(1), 4); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + // Check values are correct + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_EQ(val, static_cast(i)); + } +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithInferredDimension) { + // Test arange with -1 dimension that should be inferred + auto tensor = arange(0, 12, 1, {3, -1}); + + EXPECT_EQ(tensor->dim(), 2); + EXPECT_EQ(tensor->size(0), 3); + EXPECT_EQ(tensor->size(1), 4); // Should be inferred as 12/3 = 4 + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + // Check values are correct + for (auto i = 0; i < tensor->numel(); ++i) { + auto val = tensor->const_data_ptr()[i]; + EXPECT_EQ(val, static_cast(i)); + } +} + +TEST_F( + TensorPtrMakerTest, + CreateArangeTensorWithInferredDimensionDifferentPosition) { + // Test arange with -1 dimension in first position + auto tensor = arange(0, 20, 1, {-1, 5}); + + EXPECT_EQ(tensor->dim(), 2); + EXPECT_EQ(tensor->size(0), 4); // Should be inferred as 20/5 = 4 + EXPECT_EQ(tensor->size(1), 5); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); +} + +TEST_F( + TensorPtrMakerTest, + CreateArangeTensorWithMultipleDimensionsAndInference) { + // Test arange with 3D tensor and -1 dimension + auto tensor = arange(0, 24, 1, {2, -1, 3}); + + EXPECT_EQ(tensor->dim(), 3); + EXPECT_EQ(tensor->size(0), 2); + EXPECT_EQ(tensor->size(1), 4); // Should be inferred as 24/(2*3) = 4 + EXPECT_EQ(tensor->size(2), 3); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); +} + +TEST_F(TensorPtrMakerTest, ArangeTensorSizesValidationMultipleNegativeOnes) { + // Test that multiple -1s cause an error + EXPECT_DEATH( + arange(0, 12, 1, {-1, -1, 3}), "sizes can contain at most one -1"); +} + +TEST_F(TensorPtrMakerTest, ArangeTensorSizesValidationZeroDimension) { + // Test that zero dimensions cause an error + EXPECT_DEATH( + arange(0, 12, 1, {3, 0, 4}), + "sizes must contain positive integers or -1"); +} + +TEST_F(TensorPtrMakerTest, ArangeTensorSizesValidationNegativeDimension) { + // Test that negative dimensions (other than -1) cause an error + EXPECT_DEATH( + arange(0, 12, 1, {3, -2, 4}), + "sizes must contain positive integers or -1"); +} + +TEST_F(TensorPtrMakerTest, ArangeTensorSizesValidationIncompatibleProduct) { + // Test that incompatible product causes an error + EXPECT_DEATH( + arange(0, 12, 1, {3, 5}), // 3*5 = 15, but numel = 12 + "product of sizes \\(15\\) does not match numel \\(12\\)"); +} + +TEST_F(TensorPtrMakerTest, ArangeTensorSizesValidationNonDivisibleInference) { + // Test that non-divisible inference causes an error + EXPECT_DEATH( + arange(0, 13, 1, {3, -1, 2}), // 13 is not divisible by 3*2 = 6 + "numel \\(13\\) is not divisible by the product of known dimensions \\(6\\)"); +} + +TEST_F(TensorPtrMakerTest, ArangeTensorSizesValidationZeroInferredDimension) { + // Test that zero inferred dimension causes an error + EXPECT_DEATH( + arange( + 0, + 0, + 1, + {5, -1}), // numel = 0, so inferred dimension would be 0/5 = 0 + "inferred dimension size must be positive, got 0"); +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithFloatStepAndSizes) { + // Test arange with float step and sizes + auto tensor = + arange(0.0, 2.0, 0.5, {2, 2}, executorch::aten::ScalarType::Float); + + EXPECT_EQ(tensor->dim(), 2); + EXPECT_EQ(tensor->size(0), 2); + EXPECT_EQ(tensor->size(1), 2); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + // Check values: [0.0, 0.5, 1.0, 1.5] + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[0], 0.0f); + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[1], 0.5f); + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[2], 1.0f); + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[3], 1.5f); +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithNegativeStepAndSizes) { + // Test arange with negative step and sizes + auto tensor = arange(10, 6, -1, {2, 2}); + + EXPECT_EQ(tensor->dim(), 2); + EXPECT_EQ(tensor->size(0), 2); + EXPECT_EQ(tensor->size(1), 2); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + // Check values: [10, 9, 8, 7] + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[0], 10.0f); + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[1], 9.0f); + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[2], 8.0f); + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[3], 7.0f); +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithSingleElementAndSizes) { + // Test arange with single element and sizes + auto tensor = arange(5, 6, 1, {1, 1}); + + EXPECT_EQ(tensor->dim(), 2); + EXPECT_EQ(tensor->size(0), 1); + EXPECT_EQ(tensor->size(1), 1); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + + EXPECT_FLOAT_EQ(tensor->const_data_ptr()[0], 5.0f); +} + +TEST_F(TensorPtrMakerTest, CreateArangeTensorWithEmptyRangeAndSizes) { + // Test arange with empty range and sizes + auto tensor = arange(5, 0, 1, {}); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 0); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); +} From daf808ec2404a51dee84e0bb369e0ac74e339b95 Mon Sep 17 00:00:00 2001 From: Mengwei Liu Date: Mon, 23 Jun 2025 15:25:30 -0700 Subject: [PATCH 2/2] Update on "[llm] Add arange() tensor maker API" As titled. `arange()` taking a `sizes` argument to be able to take custom tensor sizes. Differential Revision: [D77184741](https://our.internmc.facebook.com/intern/diff/D77184741/) [ghstack-poisoned] --- extension/tensor/test/tensor_ptr_maker_test.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/extension/tensor/test/tensor_ptr_maker_test.cpp b/extension/tensor/test/tensor_ptr_maker_test.cpp index 1685447e253..07139477e09 100644 --- a/extension/tensor/test/tensor_ptr_maker_test.cpp +++ b/extension/tensor/test/tensor_ptr_maker_test.cpp @@ -565,8 +565,7 @@ TEST_F(TensorPtrMakerTest, CreateArangeTensorWithDoubleType) { TEST_F(TensorPtrMakerTest, CreateArangeTensorWithEmptyRange) { // End < start with positive step should error out EXPECT_DEATH( - arange(5, 0, 1), - "inferred dimension size must be positive, got 0"); + arange(5, 0, 1), "inferred dimension size must be positive, got 0"); } TEST_F(TensorPtrMakerTest, CreateArangeTensorWithValidSizes) {