Log mismatched values during tensors_have check functions

Summary:
D48165168 added `bool tensors_have_same_.*(...)` functions as alternative to the `ET_CHECK_SAME_.*(...)` macros.

However, one key difference between the two is that the original macro will provide information about which values caused the check to fail, whereas the new functions don't.

This diff brings the new functions to parity of behaviour with the original macros by logging which values caused the check to return `false`.

Reviewed By: manuelcandales

Differential Revision: D48247982

fbshipit-source-id: 2f4f0ec59428d26949e44f1b04ea2ecd8d824b29

Author: SS-JIA

kernels/portable/cpu/scalar_utils.h

@@ -65,6 +65,20 @@ inline ScalarType get_scalar_dtype(Scalar scalar) {
   ET_CHECK_MSG(false, "Scalar must be Boolean, Integral or Floating.");
 }
 
+inline bool scalars_have_same_dtype(Scalar a, Scalar b) {
+  ScalarType a_dtype = get_scalar_dtype(a);
+  ScalarType b_dtype = get_scalar_dtype(b);
+  if (a_dtype == b_dtype) {
+    return true;
+  }
+  ET_LOG(
+      Error,
+      "Expected scalars to have the same dtype, but found %s and %s",
+      toString(a_dtype),
+      toString(b_dtype));
+  return false;
+}
+
 /**
  * Implement type promotion between a tensor's ScalarType with a Scalar.
  * If the Scalar contains a value in the same category of the tensor's

runtime/core/exec_aten/util/tensor_util.h

@@ -411,40 +411,199 @@ using ScalarType = exec_aten::ScalarType;
 // Utility functions for checking tensor attributes
 //
 
+inline bool tensor_can_cast_to(
+    exec_aten::Tensor a,
+    exec_aten::ScalarType dtype) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      torch::executor::canCast(a.scalar_type(), dtype),
+      "Tensor of dtype %s cannot cast to dtype %s",
+      torch::executor::toString(a.scalar_type()),
+      torch::executor::toString(dtype));
+
+  return true;
+}
+
+inline bool tensor_is_bool_type(exec_aten::Tensor t) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      t.scalar_type() == exec_aten::ScalarType::Bool,
+      "Expected to find bool type, but tensor has type %s",
+      torch::executor::toString(t.scalar_type()));
+
+  return true;
+}
+
+inline bool tensor_is_integral_type(
+    exec_aten::Tensor t,
+    bool includeBool = false) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      torch::executor::isIntegralType(t.scalar_type(), includeBool),
+      "Expected to find a integral type, but tensor has type %s",
+      torch::executor::toString(t.scalar_type()));
+
+  return true;
+}
+
+inline bool tensor_is_floating_type(exec_aten::Tensor t) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      torch::executor::isFloatingType(t.scalar_type()),
+      "Expected to find a floating type, but tensor has type %s",
+      torch::executor::toString(t.scalar_type()));
+
+  return true;
+}
+
+inline bool tensor_is_complex_type(exec_aten::Tensor t) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      torch::executor::isComplexType(t.scalar_type()),
+      "Expected to find a complex type, but tensor has type %s",
+      torch::executor::toString(t.scalar_type()));
+
+  return true;
+}
+
+inline bool tensor_is_bits_type(exec_aten::Tensor t) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      torch::executor::isBitsType(t.scalar_type()),
+      "Expected to find a bits type, but tensor has type %s",
+      torch::executor::toString(t.scalar_type()));
+
+  return true;
+}
+
 inline bool tensors_have_same_dtype(exec_aten::Tensor a, exec_aten::Tensor b) {
-  return a.scalar_type() == b.scalar_type();
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      a.scalar_type() == b.scalar_type(),
+      ET_TENSOR_CHECK_PREFIX__ ": dtype={%s, %s}",
+      torch::executor::toString(a.scalar_type()),
+      torch::executor::toString(b.scalar_type()));
+  return true;
 }
 
 inline bool tensors_have_same_dtype(
     exec_aten::Tensor a,
     exec_aten::Tensor b,
     exec_aten::Tensor c) {
-  return a.scalar_type() == b.scalar_type() &&
-      b.scalar_type() == c.scalar_type();
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      a.scalar_type() == b.scalar_type() && b.scalar_type() == c.scalar_type(),
+      ET_TENSOR_CHECK_PREFIX__ ": dtype={%s, %s, %s}",
+      torch::executor::toString(a.scalar_type()),
+      torch::executor::toString(b.scalar_type()),
+      torch::executor::toString(c.scalar_type()));
+  return true;
+}
+
+inline bool tensor_is_rank(exec_aten::Tensor t, size_t rank) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      t.dim() == rank,
+      "Expected tensor.dim() to be %zu, but got %zu",
+      static_cast<size_t>(rank),
+      static_cast<size_t>(t.dim()));
+
+  return true;
+}
+
+inline bool tensor_has_dim(exec_aten::Tensor t, int64_t d) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      d > 0 ? d < t.dim() : t.dim() + d >= 0,
+      "%zu-dim tensor does not have dim at index %zu",
+      static_cast<size_t>(t.dim()),
+      static_cast<size_t>(d));
+
+  return true;
+}
+
+inline bool tensors_have_same_size_at_dims(
+    exec_aten::Tensor a,
+    size_t dim_a,
+    exec_aten::Tensor b,
+    size_t dim_b) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      dim_a < a.dim(),
+      "Cannot retrieve dim %zu from tensor with dim %zu",
+      static_cast<size_t>(dim_a),
+      static_cast<size_t>(a.dim()));
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      dim_b < b.dim(),
+      "Cannot retrieve dim %zu from tensor with dim %zu",
+      static_cast<size_t>(dim_b),
+      static_cast<size_t>(b.dim()));
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      a.size(dim_a) == b.size(dim_b),
+      ET_TENSOR_CHECK_PREFIX__
+      ": a.size(%zu) = %zu does not match b.size(%zu) = %zu",
+      static_cast<size_t>(dim_a),
+      static_cast<size_t>(a.size(dim_a)),
+      static_cast<size_t>(dim_b),
+      static_cast<size_t>(b.size(dim_b)));
+
+  return true;
 }
 
 inline bool tensors_have_same_shape(exec_aten::Tensor a, exec_aten::Tensor b) {
-  if (a.numel() != b.numel()) {
-    return false;
-  }
-  if (a.numel() == 1) {
+  if (a.numel() == 1 && b.numel() == 1) {
     // PyTorch operators treat all scalar tensors as the same shape even if
     // they have different dims.
     return true;
   }
-  // Does a length comparison (ensuring dims are equal) and element-by-element
-  // comparison (ensuring sizes are equal).
-  if (a.sizes() != b.sizes()) {
+  if (!(a.sizes() == b.sizes() && a.numel() == b.numel())) {
+    ET_LOG(
+        Error,
+        ET_TENSOR_CHECK_PREFIX__ ": numel=(%zu,  %zu), dim=(%zu, %zu)",
+        static_cast<size_t>(a.numel()),
+        static_cast<size_t>(b.numel()),
+        static_cast<size_t>(a.dim()),
+        static_cast<size_t>(b.dim()));
+    for (size_t d = 0; d < ET_MIN2(a.dim(), b.dim()); ++d) {
+      ET_LOG(
+          Error,
+          "    size(%zu): (%zu, %zu)",
+          static_cast<size_t>(d),
+          static_cast<size_t>(a.size(d)),
+          static_cast<size_t>(b.size(d)));
+    }
+
     return false;
   }
+
   return true;
 }
 
 inline bool tensors_have_same_shape(
     exec_aten::Tensor a,
     exec_aten::Tensor b,
     exec_aten::Tensor c) {
-  return tensors_have_same_shape(a, b) && tensors_have_same_shape(b, c);
+  if (a.numel() == 1 && b.numel() == 1 && c.numel() == 1) {
+    // PyTorch operators treat all scalar tensors as the same shape even if
+    // they have different dims.
+    return true;
+  }
+  bool cond1 = (a.sizes() == b.sizes()) && (a.numel() == b.numel());
+  bool cond2 = (b.sizes() == c.sizes()) && (b.numel() == c.numel());
+
+  if (!(cond1 && cond2)) {
+    ET_LOG(
+        Error,
+        ET_TENSOR_CHECK_PREFIX__ ": numel=(%zu, %zu, %zu), dim=(%zu, %zu, %zu)",
+        static_cast<size_t>(a.numel()),
+        static_cast<size_t>(b.numel()),
+        static_cast<size_t>(c.numel()),
+        static_cast<size_t>(a.dim()),
+        static_cast<size_t>(b.dim()),
+        static_cast<size_t>(c.dim()));
+    for (size_t d = 0; d < ET_MIN3(a.dim(), b.dim(), c.dim()); ++d) {
+      ET_LOG(
+          Error,
+          "    size(%zu): (%zu, %zu, %zu)",
+          static_cast<size_t>(d),
+          static_cast<size_t>(a.size(d)),
+          static_cast<size_t>(b.size(d)),
+          static_cast<size_t>(c.size(d)));
+    }
+
+    return false;
+  }
+
+  return true;
 }
 
 inline bool tensors_have_same_shape_and_dtype(
@@ -463,14 +622,50 @@ inline bool tensors_have_same_shape_and_dtype(
 inline bool tensors_have_same_strides(
     exec_aten::Tensor a,
     exec_aten::Tensor b) {
-  return a.strides() == b.strides();
+  if (a.strides() != b.strides()) {
+    ET_LOG(
+        Error,
+        ET_TENSOR_CHECK_PREFIX__ ": dim=(%zu, %zu)",
+        static_cast<size_t>(a.dim()),
+        static_cast<size_t>(b.dim()));
+    for (size_t d = 0; d < ET_MIN2(a.dim(), b.dim()); ++d) {
+      ET_LOG(
+          Error,
+          "    stride(%zu): (%zu, %zu)",
+          static_cast<size_t>(d),
+          static_cast<size_t>(a.strides()[d]),
+          static_cast<size_t>(b.strides()[d]));
+    }
+
+    return false;
+  }
+  return true;
 }
 
 inline bool tensors_have_same_strides(
     exec_aten::Tensor a,
     exec_aten::Tensor b,
     exec_aten::Tensor c) {
-  return a.strides() == b.strides() && b.strides() == c.strides();
+  if (!(a.strides() == b.strides() && b.strides() == c.strides())) {
+    ET_LOG(
+        Error,
+        ET_TENSOR_CHECK_PREFIX__ ": dim=(%zu, %zu, %zu)",
+        static_cast<size_t>(a.dim()),
+        static_cast<size_t>(b.dim()),
+        static_cast<size_t>(c.dim()));
+    for (size_t d = 0; d < ET_MIN3(a.dim(), b.dim(), c.dim()); ++d) {
+      ET_LOG(
+          Error,
+          "    stride(%zu): (%zu, %zu, %zu)",
+          static_cast<size_t>(d),
+          static_cast<size_t>(a.strides()[d]),
+          static_cast<size_t>(b.strides()[d]),
+          static_cast<size_t>(c.strides()[d]));
+    }
+
+    return false;
+  }
+  return true;
 }
 
 inline bool tensor_is_contiguous(exec_aten::Tensor t) {
@@ -480,13 +675,21 @@ inline bool tensor_is_contiguous(exec_aten::Tensor t) {
   if (strides.size() == 0) {
     return true;
   }
-  if (strides[strides.size() - 1] != 1) {
-    return false;
-  }
-  for (auto i = strides.size() - 1; i > 0; --i) {
-    if (strides[i - 1] != strides[i] * sizes[i]) {
-      return false;
-    }
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      strides[strides.size() - 1] == 1,
+      "Tensor is not contiguous; the stride of the last dimension must be 1, "
+      "but got %zu",
+      static_cast<size_t>(strides[strides.size() - 1]));
+  for (int i = strides.size() - 1; i > 0; --i) {
+    ET_LOG_MSG_AND_RETURN_IF_FALSE(
+        strides[i - 1] == strides[i] * sizes[i],
+        "Tensor is not contiguous; the stride of dim %zu should be equal to "
+        "strides[%zu] * sizes[%zu] = %zu, but found %zu",
+        static_cast<size_t>(i - 1),
+        static_cast<size_t>(i),
+        static_cast<size_t>(i),
+        static_cast<size_t>(strides[i] * sizes[i]),
+        static_cast<size_t>(strides[i - 1]));
   }
   return true;
 }

runtime/core/exec_aten/util/test/tensor_util_test.cpp

@@ -9,6 +9,7 @@
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/platform/runtime.h>
 #include <executorch/test/utils/DeathTest.h>
 #include <cmath>
 #include <limits>
@@ -29,6 +30,12 @@ class TensorUtilTest : public ::testing::Test {
   TensorFactory<ScalarType::Float> tf_float_;
   TensorFactory<ScalarType::Double> tf_double_;
   TensorFactory<ScalarType::Bool> tf_bool_;
+
+  void SetUp() override {
+    // As some of these tests cause ET_LOG to be called, the PAL must be
+    // initialized first by calling runtime_init();
+    torch::executor::runtime_init();
+  }
 };
 
 TEST_F(TensorUtilTest, IdentityChecks) {
@@ -414,6 +421,30 @@ TEST_F(TensorUtilTest, BoolTensorNotScalarFails) {
 // Tests for utility functions that check tensor attributes
 //
 
+TEST_F(TensorUtilTest, TensorIsRankTest) {
+  using namespace torch::executor;
+  Tensor a = tf_float_.ones({2, 3, 5});
+
+  EXPECT_TRUE(tensor_is_rank(a, 3));
+  EXPECT_FALSE(tensor_is_rank(a, 0));
+  EXPECT_FALSE(tensor_is_rank(a, 5));
+}
+
+TEST_F(TensorUtilTest, TensorHasDimTest) {
+  using namespace torch::executor;
+  Tensor a = tf_float_.ones({2, 3, 5});
+
+  EXPECT_TRUE(tensor_has_dim(a, 2));
+  EXPECT_TRUE(tensor_has_dim(a, 1));
+  EXPECT_TRUE(tensor_has_dim(a, 0));
+  EXPECT_TRUE(tensor_has_dim(a, -1));
+  EXPECT_TRUE(tensor_has_dim(a, -2));
+  EXPECT_TRUE(tensor_has_dim(a, -3));
+
+  EXPECT_FALSE(tensor_has_dim(a, -4));
+  EXPECT_FALSE(tensor_has_dim(a, 3));
+}
+
 TEST_F(TensorUtilTest, TensorsHaveSameDtypeTest) {
   using namespace torch::executor;
   Tensor a = tf_float_.ones({2, 3});
@@ -427,6 +458,18 @@ TEST_F(TensorUtilTest, TensorsHaveSameDtypeTest) {
   EXPECT_FALSE(tensors_have_same_dtype(a, b, d));
 }
 
+TEST_F(TensorUtilTest, TensorsHaveSameSizeAtDimTest) {
+  using namespace torch::executor;
+  Tensor a = tf_float_.ones({2, 3, 4, 5});
+  Tensor b = tf_float_.ones({5, 4, 3, 2});
+
+  EXPECT_TRUE(tensors_have_same_size_at_dims(a, 0, b, 3));
+  EXPECT_TRUE(tensors_have_same_size_at_dims(a, 1, b, 2));
+  EXPECT_FALSE(tensors_have_same_size_at_dims(a, 1, b, 0));
+  EXPECT_FALSE(tensors_have_same_size_at_dims(a, 4, b, 0));
+  EXPECT_FALSE(tensors_have_same_size_at_dims(a, 2, b, 3));
+}
+
 TEST_F(TensorUtilTest, TensorsHaveSameShapeTest) {
   using namespace torch::executor;
   Tensor a = tf_float_.ones({2, 3});