Add TensorNames::checkUnique, operator<< (TensorName) (#29124)

zou3519 · facebook-github-bot · commit 3233a058fae3 · 2019-11-05T07:24:19.000-08:00
Summary: Pull Request resolved: #29124 TensorNames::checkUnique gives a nice error message if there are duplicate names. Adding operator<< on TensorName cleans up some code. A TensorName gets printed out as: "'H' (index 2 of ['N', 'C', 'H', 'W'])" for example. Test Plan: - New c++ tests. test with `build/bin/NamedTensor_test`. Differential Revision: D18311868 Pulled By: zou3519 fbshipit-source-id: 5be197dba227f0328b40d7f66e78fffefe4dbd00
diff --git a/aten/src/ATen/TensorNames.cpp b/aten/src/ATen/TensorNames.cpp
@@ -26,11 +26,8 @@ const TensorName& TensorName::unify(const TensorName& other, const char* op_name
     const auto it = std::find(other.origin_.begin(), other.origin_.end(), name_);
     TORCH_CHECK(it == other.origin_.end(),
         op_name, ":",
-        " Cannot match ", name_,
-        " (at index ", origin_idx_, " of ", origin_, ")",
-        " with ", other.name_,
-        " (at index ", other.origin_idx_, " of ", other.origin_, ")",
-        " because the latter names already has ", name_, ".",
+        " Cannot match ", *this, " with ", other,
+        " because the latter names already have ", name_, ".",
         " Are your tensors misaligned?");
     return *this;
   }
@@ -43,10 +40,8 @@ const TensorName& TensorName::unify(const TensorName& other, const char* op_name
   // unify(A, B)
   TORCH_CHECK(name_ == other.name_,
       op_name, ":",
-      " Expected ", name_,
-      " (at index ", origin_idx_, " of ", origin_, ")",
-      " to match ", other.name_,
-      " (at index ", other.origin_idx_, " of ", other.origin_, ")",
+      " Expected ", *this,
+      " to match ", other,
       " but they do not match.");
   return *this;
 }
@@ -89,6 +84,34 @@ void TensorNames::append(TensorName&& name) {
   names_.emplace_back(name);
 }
 
+void TensorNames::checkUnique(const char* op_name) const {
+  // O(N^2), but named tensors can have at most N = 64 dimensions, so this
+  // doesn't matter unless benchmarking tells us it does. The alternative is
+  // to create some sort of set data structure but the overhead of that
+  // might dominate for small sizes.
+  for (auto it = names_.begin(); it != names_.end(); ++it) {
+    const auto name = it->toDimname();
+    if (name.isWildcard()) continue;
+
+    auto dup = std::find_if(it + 1, names_.end(),
+        [&](const TensorName& other) { return other.toDimname() == name; });
+    TORCH_CHECK(dup == names_.end(),
+        op_name, ": ",
+        "Attempted to propagate dims ", *it, " and ", dup, " to the output, ",
+        "but that would create a tensor with duplicate names ", toDimnameVec(),
+        ". Please rename your inputs with Tensor.rename to prevent this.");
+  }
+}
+
+// Let's say the TensorName represents 'C' in ['N', 'C', 'H, 'W'].
+// It should print like:
+// 'C' (index 1 of ['N', 'C', 'H', 'W'])
+std::ostream& operator<<(std::ostream& out, const TensorName& tensorname) {
+  out << tensorname.name_ << " (index " << tensorname.origin_idx_ << " of ";
+  out << tensorname.origin_ << ")";
+  return out;
+}
+
 std::vector<Dimname> TensorNames::toDimnameVec() const {
   std::vector<Dimname> result;
   result.reserve(names_.size());
diff --git a/aten/src/ATen/TensorNames.h b/aten/src/ATen/TensorNames.h
@@ -34,13 +34,17 @@ struct CAFFE2_API TensorName {
       origin_idx_(origin_idx),
       name_(origin[maybe_wrap_dim(origin_idx, origin.size())]) {}
 
+  // op_name is only used for error reporting.
   const TensorName& unify(const TensorName& other, const char* op_name) const;
   Dimname toDimname() const;
 
  private:
   ArrayRef<Dimname> origin_;
   int origin_idx_;
   Dimname name_;
+  CAFFE2_API friend std::ostream& operator<<(
+      std::ostream& out,
+      const TensorName& tensorname);
 };
 
 using TensorNameVec = SmallVector<TensorName, 10>;
@@ -52,7 +56,9 @@ struct CAFFE2_API TensorNames {
   // `names`, NOT names[start:end], because the original tensor's names are `names`.
   explicit TensorNames(ArrayRef<Dimname> names, int64_t start, int64_t end);
 
+  // op_name is only used for error reporting.
   TensorNames unifyFromRight(const TensorNames& other, const char* op_name) const;
+  void checkUnique(const char* op_name) const;
 
   void append(TensorName&& name);
   std::vector<Dimname> toDimnameVec() const;
@@ -62,6 +68,7 @@ struct CAFFE2_API TensorNames {
 
   TensorNameVec names_;
 };
+
 #endif
 
 }} // namespace at::namedinference
diff --git a/aten/src/ATen/test/NamedTensor_test.cpp b/aten/src/ATen/test/NamedTensor_test.cpp
@@ -12,6 +12,8 @@ using at::Dimname;
 using at::DimnameList;
 using at::NamedTensorMeta;
 using at::Symbol;
+using at::namedinference::TensorName;
+using at::namedinference::TensorNames;
 using c10::guts::make_unique;
 
 TEST(NamedTensorTest, defaultMetadata) {
@@ -209,5 +211,42 @@ TEST(NamedTensorTest, NoNamesGuard) {
   ASSERT_TRUE(at::NamesMode::is_enabled());
 }
 
+static std::vector<Dimname> nchw() {
+  auto N = dimnameFromString("N");
+  auto C = dimnameFromString("C");
+  auto H = dimnameFromString("H");
+  auto W = dimnameFromString("W");
+  return { N, C, H, W };
+}
+
+TEST(NamedTensorTest, TensorNamePrint) {
+  auto names = nchw();
+  {
+    auto N = TensorName(names, 0);
+    ASSERT_EQ(
+        c10::str(N),
+        "'N' (index 0 of ['N', 'C', 'H', 'W'])");
+  }
+  {
+    auto H = TensorName(names, 2);
+    ASSERT_EQ(
+        c10::str(H),
+        "'H' (index 2 of ['N', 'C', 'H', 'W'])");
+  }
+}
+
+TEST(NamedTensorTest, TensorNamesCheckUnique) {
+  auto names = nchw();
+  {
+    // smoke test to check that this doesn't throw
+    TensorNames(names).checkUnique("op_name");
+  }
+  {
+    std::vector<Dimname> nchh = { names[0], names[1], names[2], names[2] };
+    auto tensornames = TensorNames(nchh);
+    ASSERT_THROW(tensornames.checkUnique("op_name"), c10::Error);
+  }
+}
+
 
 #endif
