Cast the vector from deduced type to desired type if needed. (#5409)

Summary:
Pull Request resolved: #5409

.

Reviewed By: kirklandsign

Differential Revision: D62807302

fbshipit-source-id: ce71b88c7588367def22e3baf5b835e21b42c8bf
(cherry picked from commit e8a557c)

Author: shoumikhin

extension/tensor/tensor_impl_ptr.h

@@ -73,29 +73,81 @@ TensorImplPtr make_tensor_impl_ptr(
  * Creates a TensorImplPtr that manages a newly created TensorImpl with the
  * specified properties.
  *
- * This template overload is specialized for cases where the tensor data is
- * provided as a vector. The scalar type is automatically deduced from the
- * vector's data type. The deleter ensures that the data vector is properly
- * managed and its lifetime is tied to the TensorImpl.
+ * @param sizes A vector specifying the size of each dimension.
+ * @param data A pointer to the data buffer.
+ * @param type The scalar type of the tensor elements.
+ * @param dynamism Specifies the mutability of the tensor's shape.
+ * @param deleter A custom deleter function for managing the lifetime of the
+ * data buffer. If provided, this deleter is called when the managed TensorImpl
+ * is destroyed.
+ * @return A TensorImplPtr managing the newly created TensorImpl.
+ */
+inline TensorImplPtr make_tensor_impl_ptr(
+    std::vector<exec_aten::SizesType> sizes,
+    void* data,
+    exec_aten::ScalarType type = exec_aten::ScalarType::Float,
+    exec_aten::TensorShapeDynamism dynamism =
+        exec_aten::TensorShapeDynamism::DYNAMIC_BOUND,
+    std::function<void(void*)> deleter = nullptr) {
+  return make_tensor_impl_ptr(
+      std::move(sizes), data, {}, {}, type, dynamism, std::move(deleter));
+}
+
+/**
+ * Creates a TensorImplPtr that manages a newly created TensorImpl with the
+ * specified properties.
+ *
+ * This template overload is specialized for cases where tensor data is provided
+ * as a vector. If the specified `type` differs from the deduced type of the
+ * vector's elements, and casting is allowed, the data will be cast to the
+ * specified `type`. This allows for flexible creation of tensors with data
+ * vectors of one type and a different scalar type.
  *
  * @tparam T The C++ type of the tensor elements, deduced from the vector.
  * @param sizes A vector specifying the size of each dimension.
  * @param data A vector containing the tensor's data.
  * @param dim_order A vector specifying the order of dimensions.
  * @param strides A vector specifying the strides of each dimension.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
  * @param dynamism Specifies the mutability of the tensor's shape.
  * @return A TensorImplPtr that manages the newly created TensorImpl.
  */
-template <typename T = float>
+template <
+    typename T = float,
+    exec_aten::ScalarType deduced_type = runtime::CppTypeToScalarType<T>::value>
 TensorImplPtr make_tensor_impl_ptr(
     std::vector<exec_aten::SizesType> sizes,
     std::vector<T> data,
     std::vector<exec_aten::DimOrderType> dim_order = {},
     std::vector<exec_aten::StridesType> strides = {},
     exec_aten::TensorShapeDynamism dynamism =
         exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
-  constexpr exec_aten::ScalarType scalar_type =
-      runtime::CppTypeToScalarType<T>::value;
+  if (type != deduced_type) {
+    ET_CHECK_MSG(
+        runtime::canCast(deduced_type, type),
+        "Cannot cast deduced type to specified type.");
+    std::vector<uint8_t> casted_data(data.size() * runtime::elementSize(type));
+    ET_SWITCH_REALHBBF16_TYPES(
+        type, nullptr, "make_tensor_impl_ptr", CTYPE, [&] {
+          std::transform(
+              data.begin(),
+              data.end(),
+              reinterpret_cast<CTYPE*>(casted_data.data()),
+              [](const T& val) { return static_cast<CTYPE>(val); });
+        });
+    const auto raw_data_ptr = casted_data.data();
+    auto data_ptr =
+        std::make_shared<std::vector<uint8_t>>(std::move(casted_data));
+    return make_tensor_impl_ptr(
+        std::move(sizes),
+        raw_data_ptr,
+        std::move(dim_order),
+        std::move(strides),
+        type,
+        dynamism,
+        [data_ptr = std::move(data_ptr)](void*) {});
+  }
   const auto raw_data_ptr = data.data();
   auto data_ptr = std::make_shared<std::vector<T>>(std::move(data));
   return make_tensor_impl_ptr(
@@ -112,13 +164,16 @@ TensorImplPtr make_tensor_impl_ptr(
  * Creates a TensorImplPtr that manages a newly created TensorImpl with the
  * specified properties.
  *
- * This template overload is specialized for cases where the tensor data is
- * provided as a vector. The scalar type is automatically deduced from the
- * vector's data type. The deleter ensures that the data vector is properly
- * managed and its lifetime is tied to the TensorImpl.
+ * This template overload is specialized for cases where tensor data is provided
+ * as a vector. If the specified `type` differs from the deduced type of the
+ * vector's elements, and casting is allowed, the data will be cast to the
+ * specified `type`. This allows for flexible creation of tensors with data
+ * vectors of one type and a different scalar type.
  *
  * @tparam T The C++ type of the tensor elements, deduced from the vector.
  * @param data A vector containing the tensor's data.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
  * @param dynamism Specifies the mutability of the tensor's shape.
  * @return A TensorImplPtr that manages the newly created TensorImpl.
  */
@@ -127,21 +182,121 @@ TensorImplPtr make_tensor_impl_ptr(
     std::vector<T> data,
     exec_aten::TensorShapeDynamism dynamism =
         exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
-  constexpr exec_aten::ScalarType scalar_type =
-      runtime::CppTypeToScalarType<T>::value;
   std::vector<exec_aten::SizesType> sizes{exec_aten::SizesType(data.size())};
-  const auto raw_data_ptr = data.data();
-  auto data_ptr = std::make_shared<std::vector<T>>(std::move(data));
+  return make_tensor_impl_ptr(
+      std::move(sizes), std::move(data), {0}, {1}, type, dynamism);
+}
+
+/**
+ * Creates a TensorImplPtr that manages a newly created TensorImpl with the
+ * specified properties.
+ *
+ * This template overload is specialized for cases where tensor data is provided
+ * as an initializer list. If the specified `type` differs from the deduced type
+ * of the initializer list's elements, and casting is allowed, the data will be
+ * cast to the specified `type`. This allows for flexible creation of tensors
+ * with data initializer list of one type and a different scalar type.
+ *
+ * @tparam T The C++ type of the tensor elements, deduced from the initializer
+ * list.
+ * @param sizes A vector specifying the size of each dimension.
+ * @param list An initializer list containing the tensor's data.
+ * @param dim_order A vector specifying the order of dimensions.
+ * @param strides A vector specifying the strides of each dimension.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
+ * @param dynamism Specifies the mutability of the tensor's shape.
+ * @return A TensorImplPtr that manages the newly created TensorImpl.
+ */
+template <
+    typename T = float,
+    exec_aten::ScalarType deduced_type = runtime::CppTypeToScalarType<T>::value>
+inline TensorImplPtr make_tensor_impl_ptr(
+    std::vector<exec_aten::SizesType> sizes,
+    std::initializer_list<T> list,
+    std::vector<exec_aten::DimOrderType> dim_order = {},
+    std::vector<exec_aten::StridesType> strides = {},
+    exec_aten::ScalarType type = deduced_type,
+    exec_aten::TensorShapeDynamism dynamism =
+        exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
   return make_tensor_impl_ptr(
       scalar_type,
       std::move(sizes),
-      raw_data_ptr,
-      {0},
-      {1},
-      dynamism,
-      [data_ptr = std::move(data_ptr)](void*) {});
+      std::vector<T>(std::move(list)),
+      std::move(dim_order),
+      std::move(strides),
+      type,
+      dynamism);
 }
 
+/**
+ * Creates a TensorImplPtr that manages a newly created TensorImpl with the
+ * specified properties.
+ *
+ * This template overload is specialized for cases where tensor data is provided
+ * as an initializer list. If the specified `type` differs from the deduced type
+ * of the initializer list's elements, and casting is allowed, the data will be
+ * cast to the specified `type`. This allows for flexible creation of tensors
+ * with data initializer list of one type and a different scalar type.
+ *
+ * @tparam T The C++ type of the tensor elements, deduced from the initializer
+ * list.
+ * @param list An initializer list containing the tensor's data.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
+ * @param dynamism Specifies the mutability of the tensor's shape.
+ * @return A TensorImplPtr that manages the newly created TensorImpl.
+ */
+template <
+    typename T = float,
+    exec_aten::ScalarType deduced_type = runtime::CppTypeToScalarType<T>::value>
+inline TensorImplPtr make_tensor_impl_ptr(
+    std::initializer_list<T> list,
+    exec_aten::ScalarType type = deduced_type,
+    exec_aten::TensorShapeDynamism dynamism =
+        exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
+  std::vector<exec_aten::SizesType> sizes{exec_aten::SizesType(list.size())};
+  return make_tensor_impl_ptr(
+      std::move(sizes), std::move(list), {0}, {1}, type, dynamism);
+}
+
+/**
+ * Creates a TensorImplPtr to manage a Tensor with a single scalar value.
+ *
+ * @tparam T The C++ type of the scalar value.
+ * @param value The scalar value used for the Tensor.
+ * @return A TensorImplPtr managing the newly created TensorImpl.
+ */
+template <typename T>
+inline TensorImplPtr make_tensor_impl_ptr(T value) {
+  return make_tensor_impl_ptr({}, std::vector<T>{value});
+}
+
+/**
+ * Creates a TensorImplPtr that manages a newly created TensorImpl with the
+ * specified properties.
+ *
+ * This overload accepts a raw memory buffer stored in a std::vector<uint8_t>
+ * and a scalar type to interpret the data. The vector is managed, and its
+ * lifetime is tied to the TensorImpl.
+ *
+ * @param sizes A vector specifying the size of each dimension.
+ * @param data A vector containing the raw memory buffer for the tensor's data.
+ * @param dim_order A vector specifying the order of dimensions.
+ * @param strides A vector specifying the strides of each dimension.
+ * @param type The scalar type of the tensor elements.
+ * @param dynamism Specifies the mutability of the tensor's shape.
+ * @return A TensorImplPtr managing the newly created TensorImpl.
+ */
+TensorImplPtr make_tensor_impl_ptr(
+    std::vector<exec_aten::SizesType> sizes,
+    std::vector<uint8_t> data,
+    std::vector<exec_aten::DimOrderType> dim_order,
+    std::vector<exec_aten::StridesType> strides,
+    exec_aten::ScalarType type = exec_aten::ScalarType::Float,
+    exec_aten::TensorShapeDynamism dynamism =
+        exec_aten::TensorShapeDynamism::DYNAMIC_BOUND);
+
 /**
  * Creates a TensorImplPtr that manages a newly created TensorImpl with the
  * specified properties.

extension/tensor/tensor_ptr.h

@@ -171,13 +171,18 @@ inline TensorPtr make_tensor_ptr(
  *
  * This template overload is specialized for cases where the tensor data is
  * provided as a vector. The scalar type is automatically deduced from the
- * vector's data type.
+ * vector's data type. If the specified `type` differs from the deduced type of
+ * the vector's elements, and casting is allowed, the data will be cast to the
+ * specified `type`. This allows for flexible creation of tensors with data
+ * vectors of one type and a different scalar type.
  *
  * @tparam T The C++ type of the tensor elements, deduced from the vector.
  * @param sizes A vector specifying the size of each dimension.
  * @param data A vector containing the tensor's data.
  * @param dim_order A vector specifying the order of dimensions.
  * @param strides A vector specifying the strides of each dimension.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
  * @param dynamism Specifies the mutability of the tensor's shape.
  * @return A TensorPtr that manages the newly created TensorImpl.
  */
@@ -202,10 +207,15 @@ TensorPtr make_tensor_ptr(
  *
  * This template overload is specialized for cases where the tensor data is
  * provided as a vector. The scalar type is automatically deduced from the
- * vector's data type.
+ * vector's data type. If the specified `type` differs from the deduced type of
+ * the vector's elements, and casting is allowed, the data will be cast to the
+ * specified `type`. This allows for flexible creation of tensors with data
+ * vectors of one type and a different scalar type.
  *
  * @tparam T The C++ type of the tensor elements, deduced from the vector.
  * @param data A vector containing the tensor's data.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
  * @param dynamism Specifies the mutability of the tensor's shape.
  * @return A TensorPtr that manages the newly created TensorImpl.
  */
@@ -214,19 +224,67 @@ TensorPtr make_tensor_ptr(
     std::vector<T> data,
     exec_aten::TensorShapeDynamism dynamism =
         exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
-  return make_tensor_ptr(make_tensor_impl_ptr(std::move(data), dynamism));
+  return make_tensor_ptr(make_tensor_impl_ptr(std::move(data), type, dynamism));
+}
+
+/**
+ * Creates a TensorPtr that manages a Tensor with the specified properties.
+ *
+ * This template overload is specialized for cases where the tensor data is
+ * provided as an initializer list. The scalar type is automatically deduced
+ * from the initializer list's data type. If the specified `type` differs from
+ * the deduced type of the initializer list's elements, and casting is allowed,
+ * the data will be cast to the specified `type`. This allows for flexible
+ * creation of tensors with data vectors of one type and a different scalar
+ * type.
+ *
+ * @tparam T The C++ type of the tensor elements, deduced from the initializer
+ * list.
+ * @param sizes A vector specifying the size of each dimension.
+ * @param list An initializer list containing the tensor's data.
+ * @param dim_order A vector specifying the order of dimensions.
+ * @param strides A vector specifying the strides of each dimension.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
+ * @param dynamism Specifies the mutability of the tensor's shape.
+ * @return A TensorPtr that manages the newly created TensorImpl.
+ */
+template <
+    typename T = float,
+    exec_aten::ScalarType deduced_type = runtime::CppTypeToScalarType<T>::value>
+inline TensorPtr make_tensor_ptr(
+    std::vector<exec_aten::SizesType> sizes,
+    std::initializer_list<T> list,
+    std::vector<exec_aten::DimOrderType> dim_order = {},
+    std::vector<exec_aten::StridesType> strides = {},
+    exec_aten::ScalarType type = deduced_type,
+    exec_aten::TensorShapeDynamism dynamism =
+        exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) {
+  return make_tensor_ptr(make_tensor_impl_ptr(
+      std::move(sizes),
+      std::move(list),
+      std::move(dim_order),
+      std::move(strides),
+      type,
+      dynamism));
 }
 
 /**
  * Creates a TensorPtr that manages a Tensor with the specified properties.
  *
  * This template overload allows creating a Tensor from an initializer list
  * of data. The scalar type is automatically deduced from the type of the
- * initializer list's elements.
+ * initializer list's elements. If the specified `type` differs from
+ * the deduced type of the initializer list's elements, and casting is allowed,
+ * the data will be cast to the specified `type`. This allows for flexible
+ * creation of tensors with data vectors of one type and a different scalar
+ * type.
  *
  * @tparam T The C++ type of the tensor elements, deduced from the initializer
  * list.
- * @param data An initializer list containing the tensor's data.
+ * @param list An initializer list containing the tensor's data.
+ * @param type The scalar type of the tensor elements. If it differs from the
+ * deduced type, the data will be cast to this type if allowed.
  * @param dynamism Specifies the mutability of the tensor's shape.
  * @return A TensorPtr that manages the newly created TensorImpl.
  */