Revert "Revert "Add support for multiply on direct SymInt"" (pytorch#…

…79681) This relands Add support for multiply on direct SymInt pytorch#79493 Pull Request resolved: pytorch#79681 Approved by: https://github.com/Chillee
miladm · Jun 27, 2022 · 677123a · 677123a
1 parent 8ad8a82
commit 677123a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/c10/core/SymInt.cpp b/c10/core/SymInt.cpp
@@ -4,7 +4,7 @@
 
 namespace c10 {
 
-std::shared_ptr<SymbolicIntNode> SymInt::toSymbolicIntNode() {
+std::shared_ptr<SymbolicIntNode> SymInt::toSymbolicIntNode() const {
   auto& st = getSymIntTable();
   TORCH_CHECK(is_symbolic());
   return st.getNode(static_cast<uint64_t>(data_) & ~MASK);
@@ -25,6 +25,27 @@ SymInt SymInt::operator+(SymInt sci) const {
   return SymInt(data_ + sci.data_);
 }
 
+SymInt SymInt::operator*(SymInt sci) const {
+  if (!is_symbolic() && !sci.is_symbolic()) {
+    return SymInt(data_ * sci.data_);
+  }
+  // TODO: This is way to much boilerplate
+  std::shared_ptr<SymbolicIntNode> a =
+      is_symbolic() ? toSymbolicIntNode() : nullptr;
+  std::shared_ptr<SymbolicIntNode> b =
+      sci.is_symbolic() ? sci.toSymbolicIntNode() : nullptr;
+
+  SymbolicIntNode* common = a ? a.get() : b.get();
+  // TODO: technically we need to check that the classes match
+  if (!a) {
+    a = common->wrap(data_);
+  }
+  if (!b) {
+    b = common->wrap(sci.data_);
+  }
+  return SymInt::toSymInt(a->add(b));
+}
+
 bool SymInt::operator<(SymInt sci) const {
   TORCH_CHECK(
       !this->is_symbolic() && !sci.is_symbolic(),

diff --git a/c10/core/SymInt.h b/c10/core/SymInt.h
@@ -48,6 +48,7 @@ class C10_API SymInt {
   }
 
   SymInt operator+(SymInt sci) const;
+  SymInt operator*(SymInt sci) const;
   bool operator<(SymInt sci) const;
   void operator*=(SymInt sci);
 
@@ -56,7 +57,7 @@ class C10_API SymInt {
   bool operator==(int64_t sci) const;
   bool operator!=(int64_t sci) const;
 
-  std::shared_ptr<SymbolicIntNode> toSymbolicIntNode();
+  std::shared_ptr<SymbolicIntNode> toSymbolicIntNode() const;
   static c10::SymInt toSymInt(std::shared_ptr<SymbolicIntNode> sin);
 
   int64_t as_int_unchecked() const {