Add constraint related func into SymDimMgr. (#56727)

liuruyan · web-flow · commit 14ede2b94575 · 2023-09-06T14:11:43.000+08:00
* add symbolicDimProduct &amp; symbolicDimMgr without method shape_constraint related.

* split ddim in phi, add a target ddim, used by pd_type.

* add pd_type.cc to ir_shape CMakeLists.

* add dimOp, tieProductEqualOp. access constraint_func in SymbolTable.

* put DenseTensorType into builtin_type.

* add constraint related Mgr func.

* move to out assert.
diff --git a/paddle/ir/dialect/shape/ir/shape_op.cc b/paddle/ir/dialect/shape/ir/shape_op.cc
@@ -112,6 +112,7 @@ bool SymbolicDim::merge(SymbolicDim other) {
   if (!isDynamic() && !other.isDynamic() && getValue() != other.getValue())
     return false;
   if (isDynamic() && !other.isDynamic()) updateValue(other.getValue());
+  if (!isDynamic() && other.isDynamic()) other.updateValue(getValue());
 
   bool knownNonNegativeFlag =
       getKnownNonNegative() || other.getKnownNonNegative();
diff --git a/paddle/ir/dialect/shape/utils/shape_utils.cc b/paddle/ir/dialect/shape/utils/shape_utils.cc
@@ -46,6 +46,154 @@ const std::string SymbolTable::insert(ir::Operation* symbol) {
   return name;
 }
 
+bool SymbolicDimMgr::load() {
+  for (auto op_it = m_.block()->begin(); op_it != m_.block()->end(); op_it++) {
+    symbolTable_.insert(*op_it);
+    SymbolicDim op = (*op_it)->dyn_cast<SymbolicDim>();
+    if (!op) continue;
+    symbolDimUnionSet_[op] = op;
+    symbolNameSet_.insert(op.getSymName());
+  }
+  return loadShapeConstraintGraph();
+}
+
+bool SymbolicDimMgr::loadShapeConstraintGraph() {
+  // TODO(liujinnan): add more constraint function. currently, only support
+  // tie_product_equal.
+  auto constraint_vec =
+      symbolTable_.lookup<ir::dialect::TieProductEqualOp>("tie_product_equal");
+
+  if (!constraint_vec.size()) return true;
+
+  auto build_sym_product = [&](std::vector<ir::Value> range,
+                               SymbolicDimProduct& product) {
+    for (Value v : range) {
+      auto definingOp = v.GetDefiningOp();
+      if (auto constOp = definingOp->dyn_cast<ir::ConstantOp>()) {
+        product.factor *= constOp.value().dyn_cast<ir::Int32Attribute>().data();
+        continue;
+      } else if (auto dimOp = definingOp->dyn_cast<ir::dialect::DimOp>()) {
+        auto sym = symbolTable_.lookup<SymbolicDim>(dimOp.getName());
+        if (!sym) return false;
+        product.symbols.push_back(sym);
+        continue;
+      }
+      return false;
+    }
+    return true;
+  };
+  for (auto op : constraint_vec) {
+    SymbolicDimProduct lhs, rhs;
+    if (!build_sym_product(op.getLhs(), lhs) ||
+        !build_sym_product(op.getRhs(), rhs) ||
+        !mapSymbolicDimProductEqual(lhs, rhs))
+      return false;
+  }
+  return true;
+}
+
+int64_t gcd(int64_t m, int64_t n) {
+  if (!m) return n;
+  if (!n) return m;
+  return (m < n) ? gcd(m, n % m) : gcd(m % n, n);
+}
+
+bool SymbolicDimMgr::mapSymbolicDimProductEqual(const SymbolicDimProduct& lhs,
+                                                const SymbolicDimProduct& rhs) {
+  SymbolicDimProduct newLhs, newRhs;
+  std::tie(newLhs, newRhs) = simplifySymbolicDimProductPair(lhs, rhs);
+
+  // early return for identity case.
+  if (newLhs == newRhs) return true;
+
+  if (newLhs.factor == newRhs.factor && newLhs.symbols.size() == 1 &&
+      newRhs.symbols.size() == 1) {
+    return mapSymbolicDimEqual(newLhs.symbols[0], newRhs.symbols[0]);
+  } else if (newLhs.symbols.size() == 0 && newRhs.symbols.size() == 1 &&
+             newRhs.factor == 1) {
+    return mapSymbolicDimEqual(newConstantSymbolicDim(newLhs.factor),
+                               newRhs.symbols[0]);
+  } else if (newRhs.symbols.size() == 0 && newLhs.symbols.size() == 1 &&
+             newLhs.factor == 1) {
+    return mapSymbolicDimEqual(newConstantSymbolicDim(newRhs.factor),
+                               newLhs.symbols[0]);
+  }
+
+  productEqualityMap_[newLhs][newRhs] = productEqualityMap_[newRhs][newLhs] =
+      true;
+
+  productEqualityMapUpdated_ = false;
+  return true;
+}
+
+std::pair<SymbolicDimProduct, SymbolicDimProduct>
+SymbolicDimMgr::simplifySymbolicDimProductPair(const SymbolicDimProduct& x,
+                                               const SymbolicDimProduct& y) {
+  auto lhs = simplifySymbolicDimProduct(x);
+  auto rhs = simplifySymbolicDimProduct(y);
+
+  SymbolicDimProduct newLhs, newRhs;
+  int64_t gcdFactor = gcd(std::abs(lhs.factor), std::abs(rhs.factor));
+  if (!gcdFactor) return std::make_pair(std::move(newLhs), std::move(newRhs));
+  if (std::abs(lhs.factor) < std::abs(rhs.factor)) {
+    if (lhs.factor < 0) gcdFactor = -gcdFactor;
+  } else {
+    if (rhs.factor < 0) gcdFactor = -gcdFactor;
+  }
+
+  newLhs.factor = lhs.factor / gcdFactor;
+  newRhs.factor = rhs.factor / gcdFactor;
+
+  std::unordered_map<SymbolicDim, int, SymDimHasher> lhsSymbolMap;
+  std::unordered_map<SymbolicDim, int, SymDimHasher> rhsSymbolMap;
+  for (SymbolicDim op : lhs.symbols) ++lhsSymbolMap[op];
+  for (SymbolicDim op : rhs.symbols) ++rhsSymbolMap[op];
+
+  for (SymbolicDim op : lhs.symbols) {
+    auto it = rhsSymbolMap.find(op);
+    if (it != rhsSymbolMap.end() && op.getKnownNonSizeZero()) {
+      if (--it->second == 0) rhsSymbolMap.erase(it);
+      continue;
+    }
+    newLhs.symbols.push_back(op);
+  }
+
+  for (SymbolicDim op : rhs.symbols) {
+    auto it = lhsSymbolMap.find(op);
+    if (it != lhsSymbolMap.end() && op.getKnownNonSizeZero()) {
+      if (--it->second == 0) lhsSymbolMap.erase(it);
+      continue;
+    }
+    newRhs.symbols.push_back(op);
+  }
+
+  if (!newLhs.factor) newLhs.symbols.clear();
+  if (!newRhs.factor) newRhs.symbols.clear();
+
+  return std::make_pair(std::move(newLhs), std::move(newRhs));
+}
+
+SymbolicDimProduct SymbolicDimMgr::simplifySymbolicDimProduct(
+    const SymbolicDimProduct& x) {
+  std::vector<SymbolicDim> copied;
+  copied.reserve(x.symbols.size());
+  for (SymbolicDim op : x.symbols) copied.push_back(getRootSymbolicDim(op));
+
+  sort(copied.begin(), copied.end(), [&](SymbolicDim lhs, SymbolicDim rhs) {
+    return compareSymbolicDimNames(lhs.getSymName(), rhs.getSymName());
+  });
+  SymbolicDimProduct newX;
+  newX.factor = x.factor;
+  for (SymbolicDim op : copied) {
+    if (!op.isDynamic()) {
+      newX.factor *= op.getValue();
+    } else {
+      newX.symbols.push_back(op);
+    }
+  }
+  return newX;
+}
+
 const std::string SymbolicDimMgr::getNextName() {
   std::string name;
   do {
@@ -123,4 +271,154 @@ bool SymbolicDimMgr::mapSymbolicDimEqual(SymbolicDim lhs, SymbolicDim rhs) {
   return true;
 }
 
+SymbolicDimProduct* SymbolicDimMgr::symbolicDimProductDivide(
+    const SymbolicDimProduct& lhs, const SymbolicDimProduct& rhs) {
+  SymbolicDimProduct newLhs, newRhs;
+  std::tie(newLhs, newRhs) = simplifySymbolicDimProductPair(lhs, rhs);
+
+  if (newLhs.factor == 0 || newRhs.factor == 0) return nullptr;
+  if (newLhs.factor % newRhs.factor != 0) return nullptr;
+  if (newLhs.symbols.size() < newRhs.symbols.size()) return nullptr;
+
+  SymbolicDimProduct* result = new SymbolicDimProduct();
+  result->factor = newLhs.factor / newRhs.factor;
+
+  std::unordered_map<SymbolicDim, int, SymDimHasher> symProcMap;
+  for (SymbolicDim sym : newRhs.symbols) ++symProcMap[sym];
+
+  for (SymbolicDim sym : newLhs.symbols) {
+    auto it = symProcMap.find(sym);
+    if (it == symProcMap.end()) {
+      result->symbols.push_back(sym);
+      continue;
+    }
+    if (--it->second == 0) {
+      symProcMap.erase(it);
+      continue;
+    }
+  }
+
+  if (!symProcMap.empty()) return nullptr;
+  return result;
+}
+
+bool SymbolicDimMgr::isMultipleOfKnownSymbolicDimProductEqualPair(
+    const SymbolicDimProduct& lhs, const SymbolicDimProduct& rhs) {
+  for (auto& pairOutter : productEqualityMap_) {
+    const SymbolicDimProduct& x = pairOutter.first;
+    auto factorX = symbolicDimProductDivide(lhs, x);
+    if (!factorX) continue;
+    for (auto& pairInner : pairOutter.second) {
+      if (!pairInner.second) continue;
+      const SymbolicDimProduct& y = pairInner.first;
+      auto factorY = symbolicDimProductDivide(rhs, y);
+      if (!factorY || (*factorX) != (*factorY)) continue;
+      return true;
+    }
+  }
+
+  return false;
+}
+
+bool SymbolicDimMgr::updateProductEqualityMap() {
+  // early return if nothing is updated.
+  if (productEqualityMapUpdated_) return true;
+
+  SymbolicDimProductMap newMap;
+  std::unordered_set<SymbolicDimProduct, SymProductHasher> productSet;
+  for (auto& pairOutter : productEqualityMap_) {
+    const SymbolicDimProduct& x = pairOutter.first;
+    for (auto& pairInner : pairOutter.second) {
+      if (!pairInner.second) continue;
+      const SymbolicDimProduct& y = pairInner.first;
+      SymbolicDimProduct newX, newY;
+      std::tie(newX, newY) = simplifySymbolicDimProductPair(x, y);
+      if (newX == newY) continue;
+      newMap[newX][newY] = newMap[newY][newX] = true;
+      productSet.insert(newX);
+      productSet.insert(newY);
+    }
+  }
+  // hash function of SymbolicDimProduct is expensive, thus we map it to integer
+  // domain first.
+  std::unordered_map<const SymbolicDimProduct*, size_t> symProd2Idx;
+  std::vector<const SymbolicDimProduct*> idx2SymProd(productSet.size());
+  std::vector<size_t> idx2root(productSet.size());
+  for (auto& x : productSet) {
+    size_t idx = symProd2Idx.size();
+    symProd2Idx[&x] = idx;
+    idx2SymProd[idx] = &x;
+    idx2root[idx] = idx;
+  }
+
+  auto getRootIdx = [&](size_t root) {
+    std::vector<size_t> path;
+    while (idx2root[root] != root) {
+      path.push_back(root);
+      root = idx2root[root];
+    }
+    for (size_t idx : path) idx2root[idx] = root;
+    return root;
+  };
+
+  for (size_t x = 0; x < symProd2Idx.size(); ++x) {
+    auto& xProd = *idx2SymProd[x];
+    auto& rowMap = newMap[xProd];
+    size_t xRoot = getRootIdx(x);
+    for (size_t y = x; y < symProd2Idx.size(); ++y) {
+      auto& yProd = *idx2SymProd[y];
+      if (!rowMap[yProd]) continue;
+      idx2root[getRootIdx(y)] = xRoot;
+    }
+  }
+
+  for (size_t x = 0; x < symProd2Idx.size(); ++x)
+    for (size_t y = x; y < symProd2Idx.size(); ++y) {
+      if (getRootIdx(x) != getRootIdx(y)) continue;
+      auto& xSymProd = *idx2SymProd[x];
+      auto& ySymProd = *idx2SymProd[y];
+
+      newMap[xSymProd][ySymProd] = newMap[ySymProd][xSymProd] = true;
+    }
+
+  productEqualityMap_ = std::move(newMap);
+
+  for (auto& x : productSet)
+    for (auto& y : productSet) {
+      if (!productEqualityMap_[x][y]) continue;
+      productEqualityMap_[x][y] = productEqualityMap_[y][x] = false;
+      if (!isMultipleOfKnownSymbolicDimProductEqualPair(x, y)) {
+        productEqualityMap_[x][y] = productEqualityMap_[y][x] = true;
+      }
+    }
+
+  std::unordered_set<SymbolicDimProduct, SymProductHasher> toRemove;
+  for (auto& x : productSet) {
+    if (std::all_of(productSet.begin(),
+                    productSet.end(),
+                    [&](const SymbolicDimProduct& y) {
+                      return !productEqualityMap_[x][y];
+                    })) {
+      toRemove.insert(x);
+    }
+  }
+
+  for (auto& x : toRemove) {
+    productEqualityMap_.erase(x);
+  }
+
+  productEqualityMapUpdated_ = true;
+  return true;
+}
+
+bool SymbolicDimMgr::isSymbolicDimProductEqual(const SymbolicDimProduct& lhs,
+                                               const SymbolicDimProduct& rhs) {
+  SymbolicDimProduct newLhs, newRhs;
+  std::tie(newLhs, newRhs) = simplifySymbolicDimProductPair(lhs, rhs);
+
+  // early return for identity case.
+  if (newLhs == newRhs) return true;
+  IR_ENFORCE(updateProductEqualityMap(), "Update product equality map failed.");
+  return isMultipleOfKnownSymbolicDimProductEqualPair(newLhs, newRhs);
+}
 }  // namespace ir
diff --git a/paddle/ir/dialect/shape/utils/shape_utils.h b/paddle/ir/dialect/shape/utils/shape_utils.h
@@ -18,6 +18,7 @@
 #include <type_traits>
 #include <unordered_map>
 #include <unordered_set>
+#include "paddle/ir/core/builtin_attribute.h"
 #include "paddle/ir/core/builtin_op.h"
 #include "paddle/ir/core/utils.h"
 #include "paddle/ir/dialect/shape/ir/shape_op.h"
@@ -45,7 +46,6 @@ class SymbolTable {
  public:
   explicit SymbolTable(ir::Operation* symbolTableOp)
       : symbolTableOp_(symbolTableOp) {}
-
   template <typename T>
   typename std::enable_if<std::is_same<T, SymbolicDim>::value,
                           SymbolicDim>::type
@@ -97,16 +97,36 @@ struct SymProductHasher {
 class SymbolicDimMgr {
  public:
   explicit SymbolicDimMgr(ir::ModuleOp m);
+  bool load();
   SymbolicDim newSymbolicDim(const std::string& name = {});
   SymbolicDim newConstantSymbolicDim(int64_t val);
   std::vector<SymbolicDim> createSymbolicDimsForRankedValue(Value value);
   SymbolicDim getRootSymbolicDim(SymbolicDim symbol);
   bool isSymbolicDimEqual(SymbolicDim lhs, SymbolicDim rhs);
   SymbolTable& symbolTable() { return symbolTable_; }
   bool mapSymbolicDimEqual(SymbolicDim lhs, SymbolicDim rhs);
+  SymbolicDimProduct simplifySymbolicDimProduct(const SymbolicDimProduct& x);
+  std::pair<SymbolicDimProduct, SymbolicDimProduct>
+  simplifySymbolicDimProductPair(const SymbolicDimProduct& x,
+                                 const SymbolicDimProduct& y);
+  SymbolicDimProduct* symbolicDimProductDivide(const SymbolicDimProduct& x,
+                                               const SymbolicDimProduct& y);
+
+  bool save();  // TODO(liujinnan): load constraint func
+
+  bool isSymbolicDimProductEqual(const SymbolicDimProduct& lhs,
+                                 const SymbolicDimProduct& rhs);
+  bool mapSymbolicDimProductEqual(const SymbolicDimProduct& lhs,
+                                  const SymbolicDimProduct& rhs);
 
  private:
   const std::string getNextName();
+  bool updateProductEqualityMap();
+  bool isMultipleOfKnownSymbolicDimProductEqualPair(
+      const SymbolicDimProduct& lhs, const SymbolicDimProduct& rhs);
+  bool saveShapeConstraintGraph();  // TODO(liujinnan): load & save
+                                    // shape_constraint_func
+  bool loadShapeConstraintGraph();
 
  private:
   ir::ModuleOp m_;
@@ -127,6 +147,6 @@ class SymbolicDimMgr {
       std::unordered_map<SymbolicDimProduct, bool, SymProductHasher>,
       SymProductHasher>;
   SymbolicDimProductMap productEqualityMap_;
+  bool productEqualityMapUpdated_ = true;
 };
-
 }  // namespace ir
diff --git a/test/cpp/ir/shape_dialect/symbolic_op_test.cc b/test/cpp/ir/shape_dialect/symbolic_op_test.cc
