Make make_tensor in broadcast utilities public and rename free_broadcast_tensor (#2785)

Summary:

This diff does a couple of things:
- Makes `make_tensor` a public function so that we can create temporary intermediate tensors in operators that need to do so. (Such as NMS that is implemented above in this stack)
- Renames `free_broadcast_tensor` to a more generic name `free_tensor`

Differential Revision: D55577026

Author: tarun292

kernels/portable/cpu/util/broadcast_util.cpp

@@ -19,16 +19,14 @@ namespace executor {
 using Tensor = exec_aten::Tensor;
 using ScalarType = exec_aten::ScalarType;
 
-void free_broadcast_tensor(const Tensor& broadcast_tensor) {
-  free((void*)broadcast_tensor.const_data_ptr());
-  free((void*)broadcast_tensor.sizes().data());
-  free((void*)broadcast_tensor.dim_order().data());
-  free((void*)broadcast_tensor.strides().data());
-  free(broadcast_tensor.unsafeGetTensorImpl());
+void free_tensor(const Tensor& tensor) {
+  free((void*)tensor.const_data_ptr());
+  free((void*)tensor.sizes().data());
+  free((void*)tensor.dim_order().data());
+  free((void*)tensor.strides().data());
+  free(tensor.unsafeGetTensorImpl());
 }
 
-namespace {
-
 Tensor make_tensor(
     const ArrayRef<Tensor::SizesType>& sizes,
     const ArrayRef<Tensor::DimOrderType>& dim_order,
@@ -74,8 +72,6 @@ Tensor make_tensor(
   return Tensor{tensor_impl};
 }
 
-} // namespace
-
 bool tensor_is_broadcastable_to(
     const exec_aten::ArrayRef<Tensor::SizesType> broadcast_from_shape,
     const exec_aten::ArrayRef<Tensor::SizesType> broadcast_to_shape) {

kernels/portable/cpu/util/broadcast_util.h

@@ -62,6 +62,21 @@ bool tensors_are_broadcastable_between(
  */
 bool tensors_are_broadcastable_between(const Tensor& a, const Tensor& b);
 
+/**
+ * Create a new tensor with the given sizes, dim_order, and strides.
+ *
+ * @param[in] sizes The sizes of the tensor.
+ * @param[in] dim_order The dim order of the tensor.
+ * @param[in] strides The strides of the tensor.
+ * @param[in] dtype The data type of the tensor.
+ * @returns A new tensor with the given sizes, dim_order, and strides.
+ */
+Tensor make_tensor(
+    const ArrayRef<Tensor::SizesType>& sizes,
+    const ArrayRef<Tensor::DimOrderType>& dim_order,
+    const ArrayRef<Tensor::StridesType>& strides,
+    const ScalarType& dtype);
+
 /**
  * DEPRECATED: Use `delinearize_index()` and `linearize_access_indexes()` for
  * index remapping to avoid memory allocation.
@@ -75,7 +90,7 @@ bool tensors_are_broadcastable_between(const Tensor& a, const Tensor& b);
  * @param[in] broadcast_to The tensor to which we want to broadcast to.
  * @returns A new tensor with the same shape as broadcast_to and the data
  * repeated as appropriate. This tensor contains dynamically allocated memory
- * and must be freed using free_broadcast_tensor.
+ * and must be freed using free_tensor.
  */
 __ET_DEPRECATED exec_aten::Tensor broadcast_tensor(
     const exec_aten::Tensor& broadcast_from,
@@ -202,8 +217,7 @@ __ET_NODISCARD inline Error resize_to_broadcast_target_size(
  * broadcast_tensor.
  * @returns void
  */
-__ET_DEPRECATED void free_broadcast_tensor(
-    const exec_aten::Tensor& broadcast_tensor);
+__ET_DEPRECATED void free_tensor(const exec_aten::Tensor& broadcast_tensor);
 
 /**
  * Delinearize a flattened index to per-dimension indexes.

kernels/portable/cpu/util/test/broadcast_test.cpp

@@ -32,11 +32,11 @@ TEST(BroadcastUtilTest, BroadcastTensor) {
 
   Tensor d = torch::executor::broadcast_tensor(a, c);
   EXPECT_TENSOR_DATA_EQ(d, tf.make({2, 2}, {2, 2, 2, 2}));
-  torch::executor::free_broadcast_tensor(d);
+  torch::executor::free_tensor(d);
 
   d = torch::executor::broadcast_tensor(b, c);
   EXPECT_TENSOR_DATA_EQ(d, tf.make({2, 2}, {2, 2, 2, 2}));
-  torch::executor::free_broadcast_tensor(d);
+  torch::executor::free_tensor(d);
 }
 
 TEST(BroadcastUtilTest, BroadcastableBetween) {
@@ -63,12 +63,12 @@ TEST(BroadcastUtilTest, BroadcastableToFrom) {
   ASSERT_TRUE(tensor_is_broadcastable_to(a, c));
   Tensor d = torch::executor::broadcast_tensor(a, c);
   EXPECT_TENSOR_DATA_EQ(d, tf.make({2, 2}, {2, 2, 2, 2}));
-  torch::executor::free_broadcast_tensor(d);
+  torch::executor::free_tensor(d);
 
   ASSERT_TRUE(tensor_is_broadcastable_to(b, c));
   d = torch::executor::broadcast_tensor(b, c);
   EXPECT_TENSOR_DATA_EQ(d, tf.make({2, 2}, {2, 2, 2, 2}));
-  torch::executor::free_broadcast_tensor(d);
+  torch::executor::free_tensor(d);
 }
 
 TEST(BroadcastUtilTest, NotBroadcastableTo) {