PaddlePaddle
diff --git a/‎paddle/cinn/adt/CMakeLists.txt
Copy file name to clipboard
+14-2Lines changed: 14 additions & 2 deletions b/‎paddle/cinn/adt/CMakeLists.txt
Copy file name to clipboard
+14-2Lines changed: 14 additions & 2 deletions
diff --git a/‎paddle/cinn/adt/adapter_dynamic_tensor.h
Copy file name to clipboard
+46Lines changed: 46 additions & 0 deletions b/‎paddle/cinn/adt/adapter_dynamic_tensor.h
Copy file name to clipboard
+46Lines changed: 46 additions & 0 deletions
diff --git a/‎paddle/cinn/adt/adt.h
Copy file name to clipboard
+18-13Lines changed: 18 additions & 13 deletions b/‎paddle/cinn/adt/adt.h
Copy file name to clipboard
+18-13Lines changed: 18 additions & 13 deletions
diff --git a/‎paddle/cinn/adt/anchor_sd_equation_context.cc
Copy file name to clipboard
+27-71Lines changed: 27 additions & 71 deletions b/‎paddle/cinn/adt/anchor_sd_equation_context.cc
Copy file name to clipboard
+27-71Lines changed: 27 additions & 71 deletions
diff --git a/‎paddle/cinn/adt/anchor_sd_equation_context.h
Copy file name to clipboard
+1-17Lines changed: 1 addition & 17 deletions b/‎paddle/cinn/adt/anchor_sd_equation_context.h
Copy file name to clipboard
+1-17Lines changed: 1 addition & 17 deletions
@@ -13,6 +13,7 @@ if(NOT CINN_ONLY)
     equation_value.cc
     generate_map_expr.cc
     get_sub_reshape_dim_ranges.cc
+    graph_symbolic_dim_infer_ctx.cc
     igroup.cc
     index_expr_infer_context.cc
     kgroup.cc
@@ -23,6 +24,8 @@ if(NOT CINN_ONLY)
     schedule_descriptor.cc
     schedule_dim.cc
     schedule_mesh.cc
+    dim_expr.cc
+    dim_expr_simplifier.cc
     simplify_value.cc
     write_broadcast_disabled_bidirection_equation_generator.cc)
 
@@ -31,8 +34,17 @@ if(NOT CINN_ONLY)
 
   cinn_cc_test(tree_test SRCS tree_test.cc DEPS gtest glog)
 
-  cinn_cc_test(inline_translator_test SRCS inline_translator_test.cc DEPS
-               cinncore)
+  cinn_cc_test(inline_translator_test SRCS inline_translator_test.cc DEPS gtest
+               glog)
+
+  cinn_cc_test(
+    dim_expr_test
+    SRCS
+    dim_expr_test.cc
+    DEPS
+    gtest
+    glog
+    cinncore)
 
   message(STATUS "ADT srcs: ${cinnapi_src}")
 
 
@@ -0,0 +1,46 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include "glog/logging.h"
+
+#include "paddle/cinn/adt/adt.h"
+#include "paddle/cinn/adt/symbolic_dim.h"
+#include "paddle/cinn/hlir/framework/pir/group.h"
+
+namespace cinn::adt::adapter {
+
+struct DynamicTensor final {
+  ::pir::Value node_data;
+  const hlir::framework::pir::Group* group;
+
+  bool operator==(const DynamicTensor& other) const {
+    return this->node_data == other.node_data;
+  }
+
+  std::size_t GetRank() const {
+    return cinn::hlir::framework::pir::CompatibleInfo::ValueShape(node_data)
+        .size();
+  }
+
+  const std::vector<std::optional<DimExpr>>& GetShape() const {
+    return group->graph_symbolic_dim_infer_ctx()->GetTensorDimExprs(node_data);
+  }
+};
+
+inline std::size_t GetHashValueImpl(const DynamicTensor& tensor) {
+  return std::hash<::pir::Value>()(tensor.node_data);
+}
+
+}  // namespace cinn::adt::adapter
@@ -54,6 +54,16 @@ class Union {
     return variant_ >> match;
   }
 
+  template <typename __T>
+  const __T& Get() const {
+    return std::get<__T>(variant_);
+  }
+
+  template <typename __T>
+  bool Has() const {
+    return std::holds_alternative<__T>(variant_);
+  }
+
   const std::variant<Ts...>& variant() const { return variant_; }
 
  private:
@@ -74,6 +84,7 @@ class Tuple {
             std::make_shared<std::tuple<Ts...>>(std::forward<Args>(args)...)) {}
 
   const std::tuple<Ts...>& tuple() const { return *tuple_; }
+  std::tuple<Ts...>* mut_tuple() { return &*tuple_; }
 
   template <std::size_t I>
   const auto& Get() const {
@@ -180,18 +191,18 @@ class List final {
                                bool> = true>                                   \
     class_name(Arg&& arg) : variant_(std::forward<Arg>(arg)) {}                \
                                                                                \
-    template <typename T>                                                      \
-    const T& Get() const {                                                     \
-      return std::get<T>(variant_);                                            \
+    template <typename __T>                                                    \
+    const __T& Get() const {                                                   \
+      return std::get<__T>(variant_);                                          \
     }                                                                          \
                                                                                \
-    template <typename T>                                                      \
+    template <typename __T>                                                    \
     bool Has() const {                                                         \
-      return std::holds_alternative<T>(variant_);                              \
+      return std::holds_alternative<__T>(variant_);                            \
     }                                                                          \
                                                                                \
-    template <typename T>                                                      \
-    auto Visit(const T& visitor) const {                                       \
+    template <typename __T>                                                    \
+    auto Visit(const __T& visitor) const {                                     \
       return std::visit(visitor, variant_);                                    \
     }                                                                          \
                                                                                \
@@ -236,12 +247,6 @@ bool UnionEqual(const UnionT& lhs, const UnionT& rhs) {
     using Tuple<T0, T1>::Tuple;        \
   }
 
-DEFINE_ADT_UNARY(Neg);
-DEFINE_ADT_BINARY(Add);
-DEFINE_ADT_BINARY(Mul);
-DEFINE_ADT_BINARY(Div);
-DEFINE_ADT_BINARY(Mod);
-
 #define OVERLOAD_OPERATOR_EQ_NE(type, function)              \
   inline bool operator==(const type& lhs, const type& rhs) { \
     return function(lhs, rhs);                               \
 
@@ -16,21 +16,17 @@
 
 namespace cinn::adt::config {
 
-void GenerateScheduleMeshEquations(
-    const ScheduleMesh& sched_mesh,
-    const List<Iterator>& tmp_anchor_iterators,
-    const List<Iterator>& sd_iterators,
-    Equations* equations,
-    std::unordered_map<Dim, const Constant>* dim2constant);
+void GenerateScheduleMeshEquations(const ScheduleMesh& sched_mesh,
+                                   const List<Iterator>& tmp_anchor_iterators,
+                                   const List<Iterator>& sd_iterators,
+                                   Equations* equations);
 
 namespace {
 
-void GenerateScheduleMeshEquationsImpl(
-    const List<ScheduleDim>& sched_dims,
-    const List<Iterator>& input_iterators,
-    const List<Iterator>& output_iterators,
-    Equations* equations,
-    std::unordered_map<Dim, const Constant>* dim2constant) {
+void GenerateScheduleMeshEquationsImpl(const List<ScheduleDim>& sched_dims,
+                                       const List<Iterator>& input_iterators,
+                                       const List<Iterator>& output_iterators,
+                                       Equations* equations) {
   CHECK_EQ(input_iterators->size(), output_iterators->size());
   for (std::size_t i = 0; i < output_iterators->size(); ++i) {
     Equal(input_iterators->at(i), output_iterators->at(i), equations);
@@ -41,45 +37,26 @@ void GenerateScheduleMeshEquationsImpl(
     const ScheduleMeshReshape<ScheduleMesh>& sched_reshape,
     const List<Iterator>& input_iterators,
     const List<Iterator>& output_iterators,
-    Equations* equations,
-    std::unordered_map<Dim, const Constant>* dim2constant) {
+    Equations* equations) {
   const auto& [middle_sched_mesh, shape] = sched_reshape.tuple();
   List<Iterator> middle_iterators =
       MakeIterators(GetOutputRank(middle_sched_mesh));
-  List<Dim> middle_dims = MakeDims(GetOutputRank(middle_sched_mesh));
+  List<DimExpr> middle_dims = GetOutputDimValues(middle_sched_mesh);
   CHECK_EQ(shape.value()->size(), output_iterators->size());
-  List<Dim> output_dims = MakeDims(output_iterators->size());
-  {
-    List<Constant> middle_dim_values = GetOutputDimValues(middle_sched_mesh);
-    for (std::size_t i = 0; i < middle_dim_values->size(); ++i) {
-      CHECK(dim2constant->emplace(middle_dims->at(i), middle_dim_values->at(i))
-                .second);
-    }
-
-    List<Constant> output_dim_values =
-        GetOutputDimValues(ScheduleMesh{sched_reshape});
-    for (std::size_t i = 0; i < output_dims->size(); ++i) {
-      CHECK(dim2constant->emplace(output_dims->at(i), output_dim_values->at(i))
-                .second);
-    }
-  }
+  List<DimExpr> output_dims = GetOutputDimValues(ScheduleMesh{sched_reshape});
   const auto& middle_index = MakeDot(middle_iterators, middle_dims, equations);
   const auto& output_index = MakeDot(output_iterators, output_dims, equations);
   Equal(middle_index, output_index, equations);
 
-  GenerateScheduleMeshEquations(middle_sched_mesh,
-                                input_iterators,
-                                middle_iterators,
-                                equations,
-                                dim2constant);
+  GenerateScheduleMeshEquations(
+      middle_sched_mesh, input_iterators, middle_iterators, equations);
 }
 
 void GenerateScheduleMeshEquationsImpl(
     const ScheduleMeshTranspose<ScheduleMesh>& sched_transpose,
     const List<Iterator>& input_iterators,
     const List<Iterator>& output_iterators,
-    Equations* equations,
-    std::unordered_map<Dim, const Constant>* dim2constant) {
+    Equations* equations) {
   const auto& [sched_mesh, perm] = sched_transpose.tuple();
   CHECK_EQ(GetOutputRank(sched_mesh), output_iterators->size());
   List<Iterator> middle_iterators = MakeIterators(output_iterators->size());
@@ -89,73 +66,52 @@ void GenerateScheduleMeshEquationsImpl(
           equations);
   }
   GenerateScheduleMeshEquations(
-      sched_mesh, input_iterators, middle_iterators, equations, dim2constant);
+      sched_mesh, input_iterators, middle_iterators, equations);
 }
 
 void GenerateScheduleMeshEquationsImpl(
     const ScheduleMeshPadding<ScheduleMesh>& sched_padding,
     const List<Iterator>& input_iterators,
     const List<Iterator>& output_iterators,
-    Equations* equations,
-    std::unordered_map<Dim, const Constant>* dim2constant) {
+    Equations* equations) {
   const auto& [sched_mesh, _] = sched_padding.tuple();
   CHECK_EQ(GetOutputRank(sched_mesh), output_iterators->size());
   List<Iterator> middle_iterators = MakeIterators(output_iterators->size());
   for (std::size_t i = 0; i < output_iterators->size(); ++i) {
     Equal(middle_iterators->at(i), output_iterators->at(i), equations);
   }
   GenerateScheduleMeshEquations(
-      sched_mesh, input_iterators, middle_iterators, equations, dim2constant);
+      sched_mesh, input_iterators, middle_iterators, equations);
 }
 
 }  // namespace
 
-void GenerateScheduleMeshEquations(
-    const ScheduleMesh& sched_mesh,
-    const List<Iterator>& tmp_anchor_iterators,
-    const List<Iterator>& sd_iterators,
-    Equations* equations,
-    std::unordered_map<Dim, const Constant>* dim2constant) {
+void GenerateScheduleMeshEquations(const ScheduleMesh& sched_mesh,
+                                   const List<Iterator>& tmp_anchor_iterators,
+                                   const List<Iterator>& sd_iterators,
+                                   Equations* equations) {
   return std::visit(
       [&](const auto& impl) {
         return GenerateScheduleMeshEquationsImpl(
-            impl, tmp_anchor_iterators, sd_iterators, equations, dim2constant);
+            impl, tmp_anchor_iterators, sd_iterators, equations);
       },
       sched_mesh.variant());
 }
 
-void AnchorSdEquationContext::InitDim2Constant(const ScheduleMesh& sched_mesh) {
-  const auto& AddDimValue = [&](const List<Dim>& dims,
-                                const List<Constant>& dim_values) {
-    CHECK_EQ(dims->size(), dim_values->size());
-    for (std::size_t i = 0; i < dims->size(); ++i) {
-      CHECK(dim2constant_.emplace(dims->at(i), dim_values->at(i)).second);
-    }
-  };
-
-  const auto& anchor_dim_values =
-      GetOutputDimValues(GetInputScheduleMesh(sched_mesh));
-  AddDimValue(anchor_dims_, anchor_dim_values);
-
-  const auto& sd_dim_values = GetOutputDimValues(sched_mesh);
-  AddDimValue(sd_dims_, sd_dim_values);
-}
-
 void AnchorSdEquationContext::GenerateSdEquation(const ScheduleMesh& sched_mesh,
                                                  const Index& anchor_index) {
   const auto& tmp_anchor_iterators = MakeIterators(GetInputRank(sched_mesh));
 
   {
+    const auto& anchor_dim_values =
+        GetOutputDimValues(GetInputScheduleMesh(sched_mesh));
     const auto& tmp_anchor_index =
-        MakeDot(tmp_anchor_iterators, anchor_dims_, &equations_);
+        MakeDot(tmp_anchor_iterators, anchor_dim_values, &equations_);
     Equal(tmp_anchor_index, anchor_index, &equations_);
   }
 
-  GenerateScheduleMeshEquations(sched_mesh,
-                                tmp_anchor_iterators,
-                                sd_iterators_,
-                                &equations_,
-                                &dim2constant_);
+  GenerateScheduleMeshEquations(
+      sched_mesh, tmp_anchor_iterators, sd_iterators_, &equations_);
 }
 
 }  // namespace cinn::adt::config
@@ -31,36 +31,20 @@ class AnchorSdEquationContext final {
 
   AnchorSdEquationContext(const ScheduleMesh& sched_mesh,
                           const AnchorIndex& anchor_index)
-      : sd_dims_(MakeDims(GetOutputRank(sched_mesh))),
-        sd_iterators_(MakeIterators(GetOutputRank(sched_mesh))),
-        anchor_dims_(MakeDims(GetInputRank(sched_mesh))) {
-    InitDim2Constant(sched_mesh);
+      : sd_iterators_(MakeIterators(GetOutputRank(sched_mesh))) {
     GenerateSdEquation(sched_mesh, anchor_index);
   }
 
-  const List<Dim>& sd_dims() const { return sd_dims_; }
-
-  const List<Dim>& anchor_dims() const { return anchor_dims_; }
-
   const List<Iterator>& sd_iterators() const { return sd_iterators_; }
 
   const Equations& equations() const { return equations_; }
 
-  const std::unordered_map<Dim, const Constant>& dim2constant() const {
-    return dim2constant_;
-  }
-
  private:
-  void InitDim2Constant(const ScheduleMesh& sched_mesh);
-
   void GenerateSdEquation(const ScheduleMesh& sched_mesh,
                           const Index& tensor_index);
 
-  List<Dim> sd_dims_;
   List<Iterator> sd_iterators_;
-  List<Dim> anchor_dims_;
   Equations equations_;
-  std::unordered_map<Dim, const Constant> dim2constant_;
 };
 
 }  // namespace cinn::adt::config