linearize_access_indexes only needs dim of broadcast_to (#270)

Summary:
Pull Request resolved: #270

The utility function `linearize_access_indexes` currently takes the `broadcast_to` tensor as the second argument. However, only the dimension of the tensor is needed.
Updating the function, to take only the dimension, so that it can be used in scenarios where a broadcast_to tensor is not available, and only a broadcast shape/ndim is known.

Reviewed By: kirklandsign

Differential Revision: D49159987

fbshipit-source-id: c6532b2a48b823234accdc3e69edc0e21147658a

Author: manuelcandales

kernels/portable/cpu/util/broadcast_util.cpp

@@ -263,11 +263,11 @@ void delinearize_index(
 
 size_t linearize_access_indexes(
     ArrayRef<size_t> indexes_broadcast_to,
-    const Tensor& broadcast_to,
+    ssize_t broadcast_to_ndim,
     const Tensor& broadcast_from) {
-  size_t num_skip_dims = broadcast_to.dim() - broadcast_from.dim();
+  size_t num_skip_dims = broadcast_to_ndim - broadcast_from.dim();
   ArrayRef<size_t> indexes_broadcast_from = indexes_broadcast_to.slice(
-      num_skip_dims, broadcast_to.dim() - num_skip_dims);
+      num_skip_dims, broadcast_to_ndim - num_skip_dims);
 
   ET_CHECK(indexes_broadcast_from.size() == broadcast_from.dim());
 

kernels/portable/cpu/util/broadcast_util.h

@@ -200,13 +200,13 @@ void delinearize_index(
  * broadcast_to tensor and itself.
  *
  * @param[in] indexes The tensor access indexes of broadcast_to tensor
- * @param[in] broadcast_to The tensor with the broadcasted shape.
+ * @param[in] broadcast_to_ndim The number of dims of the broadcasted shape.
  * @param[in] broadcast_from The tensor to be broadcasted.
  * @returns The flattend index for broadcast_from tensor.
  */
 size_t linearize_access_indexes(
     ArrayRef<size_t> indexes_broadcast_to,
-    const Tensor& broadcast_to,
+    ssize_t broadcast_to_ndim,
     const Tensor& broadcast_from);
 
 //
@@ -241,10 +241,10 @@ inline void apply_binary_elementwise_fn(
       delinearize_index(i, out, out_indexes, kTensorDimensionLimit);
 
       if (a_is_broadcasted) {
-        a_linear_index = linearize_access_indexes(out_indexes, out, a);
+        a_linear_index = linearize_access_indexes(out_indexes, out.dim(), a);
       }
       if (b_is_broadcasted) {
-        b_linear_index = linearize_access_indexes(out_indexes, out, b);
+        b_linear_index = linearize_access_indexes(out_indexes, out.dim(), b);
       }
     }
 
@@ -290,13 +290,13 @@ inline void apply_ternary_elementwise_fn(
       delinearize_index(i, out, out_indexes, kTensorDimensionLimit);
 
       if (a_is_broadcasted) {
-        a_linear_index = linearize_access_indexes(out_indexes, out, a);
+        a_linear_index = linearize_access_indexes(out_indexes, out.dim(), a);
       }
       if (b_is_broadcasted) {
-        b_linear_index = linearize_access_indexes(out_indexes, out, b);
+        b_linear_index = linearize_access_indexes(out_indexes, out.dim(), b);
       }
       if (c_is_broadcasted) {
-        c_linear_index = linearize_access_indexes(out_indexes, out, c);
+        c_linear_index = linearize_access_indexes(out_indexes, out.dim(), c);
       }
     }
 

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

@@ -169,7 +169,7 @@ TEST(BroadcastUtilTest, LinearizeIndex) {
     size_t test_indexes[] = {0, 0, 0, i};
     ArrayRef<size_t> broadcast_to_indexes(test_indexes);
     size_t linear_index = linearize_access_indexes(
-        broadcast_to_indexes, broadcast_to, broadcast_from);
+        broadcast_to_indexes, broadcast_to.dim(), broadcast_from);
     EXPECT_EQ(linear_index, 0);
   }
 
@@ -179,7 +179,7 @@ TEST(BroadcastUtilTest, LinearizeIndex) {
     size_t test_indexes[] = {0, i, 2, 3};
     ArrayRef<size_t> broadcast_to_indexes(test_indexes);
     size_t linear_index = linearize_access_indexes(
-        broadcast_to_indexes, broadcast_to, broadcast_from);
+        broadcast_to_indexes, broadcast_to.dim(), broadcast_from);
     EXPECT_EQ(linear_index, 2);
   }
 }