Repository: alibaba/BladeDISC (language: python)

Issue:
[Design Doc] Shape constraint IR and related optimizations
# Background
Dynamic shape problem is ubiquitous when deploying AI workloads. Almost all our in-house models have dynamic shape problem in some degree. Dynamic shape problem can be divided into two categories.
- Input shape of the model is dynamic
   * different batch size in inference scenarios
   * different image size, sequence length, etc
- Shape is dynamic by design
  * For seq2seq model, the number of decoding step is dynamic
  * For CV detection model, the number of detected instants is dynamic
  * For recommendation model, shape is affected by the value of its inputs (e.g. `tf.UniqueOp`)
  
To solve the above problem,  BladeDISC, an e2e dynamic shape compiler, is proposed.  We already have seen many advantages of using BladeDISC over the previous static-shape-based compiler (e.g. XLA) in many pratical scenarios. However, there are still many rooms to improvement, especially in terms of the performance. In dynamic shape semantics, the concrete shape is out of reach, making many optimization strategies used in static shape semantics invalid. According to our previous experience in optimizing dynamic shape AI workloads, we found that shape constraint,  the relationship between different symbolic shape (or dimension size),  is very helpful in terms of performance optimization. 
For example:
- (partial) shape inference, propagating some known information to make best use of the existing information.
  * e.g. `mhlo.add(tensor<?x10xf32>, tensor<10x?xf32>) -> tensor<?x?xf32>` can be simplifier to `mhlo.add(tensor<10x10xf32>, tensor<10x10xf32>) -> tensor<10x10xf32>`.
- graph optimization. For example, removing redundant `mhlo.dynamic_broadcast_in_dim` op if we can infer that the output and input of the broadcast op have the same shape.
- layout optimization. For example, grouping shape-compatible ops, using same layout inside the group, saving redundant layout conversion between ops inside the groups.
- Fusion decision. Making use of the shape constraint information, we can do better fusion decision in dynamic shape semantics, e.g., only fusing shape-compatible ops into one kernel. 
- Codegen optimization. For example, removing many redundant index calculation if  we know some symbolic dimensions have the same size, which is crucial when the fusion pattern is relative large.
- and many others ...

In this post, we will first give a brief introduction to the definition of the shape constraint we used in the post and where we can get the shape constraint information. Then, we will show the limitations of the current strategy of making use of  the shape constraint information and why we need to use shape constraint IR. At last, we will describe the details of  shape constraint IR design and the related optimizations.

# What is shape constraint?
There are two kinds of shape constraint:
- Structure shape constraint,  the predicate (relationship) between different symbolic dimensions.
  * dimension size equality
    + the size of one dimension of a tensor is equal to the size of one dimension of another tensor.
    + the size of one dimension of a tensor is equal to the size of another dimension of the same tensor.
  * Number elements equality
    + the number of elements of a tensor is equal to the number of elements of another tensor.
  * group multiplication equality
    + e.g. reshape([a, b, c, d]) -> [a*b, c*d]
  * and many others ...
- Shape distribution constraint
  * dim size % 4 = 0;
  * Likely values of the symbolic dimension (shape);
  * Possible shape ranges of the symbolic dimension (shape);
  * and many others ...

# Where to get shape constraint information?
## Semantics of  the definition of operation
* The inputs and outputs of an element op of `mhlo` dialect should have same shape
* The input and output of a `mhlo.dynamic_reshape` op should have same number of elements
* The inputs and output of a `mhlo.concat` op should have same dim size for each non-concat dim
* and many others...
![image](https://user-images.githubusercontent.com/15364516/168256459-5af1e105-b0d5-4879-9bc3-52e55ad49c13.png)


## Injected by frontend converter
BladeDISC Frontend converter lowers a coarse grained op to a bunch of fine grained MHLO ops. Not only the data computation but also shape constraint  should be lowered during dialect conversions.
For example, `tf.SplitOp` will be lowered into a series of mho.RealDynamicSliceOp. According to the definition of the `tf.SplitOp`, we can know that all the outputs of the split op should have same shape and `the split dimension of the input % (number of outputs) == 0`.
![image](https://user-images.githubusercontent.com/15364516/168256068-272c6c45-fec9-450d-886e-1988a651cd68.png)

In the current implementation,  we only do data computation IR lowering, losing some shape constraint information, which in turn leads to losing some opportunities to do further optimizations.

## Symbolic Shape Analysis
There are two kinds of symbolic shape analysis ways.
- Propagation existing shape constraint to find more shape constraint
  * Input -> output, input -> input, output -> input
- Shape computation IR analysis
  * As shown in the below example, we can know more about the dynamic reshape op by analyzing the shape computation IR (to be concrete, the IR used to calculate the target shape of the `dynamic_reshape` op).
 
![image](https://user-images.githubusercontent.com/15364516/168256138-ac3d62fe-9db4-47a4-ae68-77e16483891c.png)

## Injected at JIT compilation time
The input shapes are actually known at JIT compilation time. We may choose to inject such information as likely shape values of the inputs of the graph and then propagate such information.

## Provided by end users 
User may provide more information about the distribution of the shape of inputs (e.g. shape range) or relationships of the shape of inputs (e.g. two inputs have same shape). And then, the compiler can make use of such information to do better optimizations.

# Why shape constraint IR
Although we have already made use of the shape constraint in some degree, there are still many limitations of the current implementation. In the current implementation, The shape constraint information itself is not stored using IR and we obtains the shape constraint information by doing global analysis on the data computation IR each time we need it.  The above design choice leads to following problems:
- shape constraint information loss during dialect conversion. During the whole optimization pipeline, the data computation IR is lowered many times (e.g. tf/torch dialect -> mhlo, mhlo -> lmhlo).  For each lowering, we only do data computation IR lowering currently, losing some shape constraint information. 
- the result of shape constraint analysis is not stable after IR mutation (e.g. after running canonicalization pass). Such inconsistent result may lead to wired behaviour when making use of shape constraint information across passes. 
- doing analysis on the data computation IR globally to get shape constraint information is more challenging or even impossible after lowering to buffer level (e.g. lmhlo dialect). It's very hard to do analysis by chasing a series of load/store ops, especially when there exists control flow ops.
- shape distribution constraint can not always be analyzed from the data computation IR. And we do not have a unify way to describe such information across different dialects.

To solve the above problem, we propose using shape constraint IR to describe shape constraint information explicitly.  We'll give more details about the design in following sections.

# Shape constraint IR design
## Basic idea
- using type instead of dedicated shape constraint op to encode  shape constraint information into data computation IR.
  * Type-based schema not introduce any other ops, which means existing pattern-based passes (e.g. `matmul+BiadAdd` -> `FusedMatmulBiasAdd` ) not do need to take shape constraint ops into consideration.
```
// type-based encoding schema
// each unique reference attribute (e.g. @S0, @S1) represents a unique symbolic dim.
%2 = mhlo.add(%0, %1) : (tensor<?x?xf32, [@S0, @S1]>, tensor<?x?xf32, [@S0, @S1]>) -> tensor<?x?xf32, [@S0, @S1]>

// op-based encoding schema
%d0 = tensor.dim %0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %0, %c1 : tensor<?x?xf32>
// explicitly binding each dimension size with the corresponding tensor 
%n0 = disc_shape.tie_shape(%0, %d0, %d1) : tensor<?x?xf32>
%n1 =  disc_shape.tie_shape(%1, %d0, %d1) : tensor<?x?xf32>
%2 = mhlo.add(%n1, %n2) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
%n2 =  disc_shape.tie_shape(%2, %d0, %d1) : tensor<?x?xf32>
```
- Separating shape constraint IR and data computation IR. 
   * Some shape constraint information (e.g. shape distribution constraint) is not always possible to be bound with data computation IR.
   * A unified way to describe shape constraint information, no matter what dialects the data computation IR used.  And thus, the shape constraint information will no longer be lost when lowering data computation IR during the whole optimization pipeline.

```
func @main() {
  ...
  %0 = any_dialect.any_operation(...) : tensor<?x?xf32, [@S0, @S1]>
  ...
}

disc_shape.SymbolDim @S0 {
  range list : [[...], [...], ...]
  likely_values : [...]
  ...
  symbolic_shape_graph: @shape_constraint_graph
}

disc_shape.SymbolDim @S1 {
  range list : [[...], [...], ...]
  likely_values : [...]
  ...
  symbolic_shape_graph: @shape_constraint_graph
}

// A separated function to store shape constraint predicates between different symbolic dimensions.
// Each symbolic dim is either bound to a `disc_shape.dim` op or `disc_shape.bind_dim`
func @shape_constraint_graph(...) {
  %0 = disc_shape.dim() {ref: @S0} : index
  %1 = disc_shape.dim() {ref: @S1} : index
  disc_shape.tie_predicate_divisible(d0, d1) // d0 % d1 == 0
  
  // other tie_* ops
  //   disc_shape.tie_predicate_eq(d0, d1)  // d0 == d1
  //   disc_shape.tie_predicate_lt(d0, d1)  // dim less than
  //   disc_shape.tie_predicate_mul_eq(d0, d1, d2, ...) // d0 = d1 * d2 * ...
  //   // d0 * d1 = s0 * s1
  //   disc_shape.tie_predicate_group_mul_eq([d0, d1, ..], [s0, s1, ...])
  //   // d0 = affine.apply(d1, d2, ...) {affine_attr = ...}
  //   disc_shape.tie_predicate_affine_eq(d0, d1, d2, ...) {affine_attr = ...}
}
```
- explicitly materialize all shape computation IR on tensor level instead of on buffer level. Currently, we only materialize all shape computation IR after bufferization (mhlo -> lmhlo). It means we have to do some shape computation IR analysis and optimization on buffer level, which could be much difficult than doing same thing on tensor level since we have to chasing each load/store ops and buffer alias to do analysis.

## Basic flow
Most logic is implemented inside the shape optimization pass which is shown as following diagram.
The shape optimization pass is tensor-level only pass, and is reenterable. The whole pass can be divided into two stages.
- stage one: Explicitly materialize shape computation IR on tensor level.
- stage two: shape constraint analysis and shape computation IR optimization on tensor level.
![image](https://user-images.githubusercontent.com/15364516/168513885-ea810cbe-10f9-4080-aa8b-46479c40e948.png)

### stage one: explicitly materialize shape computation IR on tensor level
It's main functions:
- (partial) shape inference, (input -> output only in the phase)
   *  no need to do shape inference manually anymore.
   *  taking advantage of existing `canonicalize pattern & cse pass`
   *  example: 
       +  `concat([tensor<?x10xf32>, tensor<?x10xf32>], axis=1) -> tensor<?x20xf32>`
- enable more opportunities to optimize shape computation IR
   *  again, taking advantage of existing `canonicalize pattern & cse pass`
   *  examples:
   ```
    // An unified solution to handle partial known shape information
    - mhlo.real_dynamic_slice: 
       - %out = %in[:,1:]
    - mhlo.dynamic_conv
       - if stride = 1, padding = kernel size, then output spatial size == input spatial size
    - dynamic pad
       - if only padding along some of the axes,
       - then output and input have some size along the other axes
    - ...
    // take `mhlo.real_dynamicm_slice:  %out = %in[:,1:]` as an example
    %d0 = %tensor.dim %in, %c0 : tensor<?x?xf32>
    %d1 = %tensor.dim %in, %c1 : tensor<?x?xf32>
    %begin = tensor.from_elements %c0, %c1 : tensor<2xindex>
    %end = tensor.from_elements %d0, %d1 : tensor<2xindex>
    %strides = tensor.from_elements %c1, %c1 : tensor<2xindex>
    %out = mhlo.real_dynamic_slice(%in, %begin, %end, %strides)
  
    // we only need to provide a `tensor.extract_elements + tensor.from_elements`
    // rewrite pattern, and then we can know that the output and input have same
    // dim size along the first axis by taking advantage of existing
    // `canonicalize pattern & cse pass`
   ```
It's basic flow:
- For a given data computation function (e.g. the main function), we insert a `disc_shape.tie_shape` for each value having `RankedTensorType` type. An example is shown as below.
```
// before inserting tie_shape op
func @main(%arg0 : tensor<?xf32>) -> tensor<?xf32> {
  %0 = "mhlo.some_ops"(%arg0) : tensor<?xf32>
  return %0 : tensor<?xf32>
}

// after inserting tie_shape op
func @main(%arg0 : tensor<?xf32>) -> tensor<?xf32> {
  %c0 = constant 0 : index
  // insert tie_shape op for %arg0
  %d0_arg0 = tensor.dim %arg0, %c0 : tensor<?xf32>
  %new_arg0 = disc_shape.tie_shape(%arg0, %d0_arg0)
  
  %result = "mhlo.some_ops"(%new_arg0) : tensor<?xf32>
  
  // insert tie_shape op for %result
  %d0_result = tensor.dim %result, %c0 : tensor<?xf32>
  %new_result = disc_shape.tie_shape(%result, %d0_result)
  return %new_result : tensor<?xf32>
}
```

- recursively  resolve tensor.dim/shape.shape_of op based on shape interface implementation. basic idea is shown as below.
```
// take `tensor.dim` as an example, same to shape.shape_of
  
for each tensor.dim op:
  operand = dimOp.getOperand(0)
  definingOp = operand.getDefiningOp()
  if (definingOp has shape interface implementation)
    materialize the shape interface implementation
        
// reference implementation pass:
// mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp
```
### stage two: shape constraint analysis and shape computation IR optimization on tensor level
- It’s an iterative process
- The whole process can be divided into five steps:
  * canonicalization stage
    + including normal canonicalization pattern & CSE
    + and a bunch of new rewrite patterns
       - Scalarize shape computation IR whenever possible
         + shape.shape_of -> a bunch of tensor.dim ops
         + shape.broadcast(shape_a, shape_b)
            - Dim-wise broadcast + tensor.from_elements
        + bcast/reshape simplifier
       - ...
  * load existing shape constraint IR
  * global shape analysis
    + analysis shape computation IR -> find more constraint
    + inject shape constraint implied by `mhlo` op definition
    + ...
  * global  shape optimization
    + using same SSA value for symbolic equal dims
    + refine types with partial known information globally
  * save the updated shape constraint into IR


Discussion / hints:
@wyzero It's excellent work! Do you want to submit a PR for this design doc? that is better to review and we can merge this design doc into the codebase later.


> @wyzero It's excellent work! Do you want to submit a PR for this design doc? that is better to review and we can merge this design doc into the codebase later.

Yeah, I'll send a PR later.

Fix patch:
diff --git a/tao_compiler/mlir/disc/BUILD b/tao_compiler/mlir/disc/BUILD
index 0ff0e731ab3..220492cfb1a 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -175,6 +175,7 @@ cc_library(
     deps = [
         ":cycle_detector",
         ":placement_utils",
+        "//tensorflow/core:lib",
         "@llvm-project//mlir:MemRefDialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Pass",
@@ -327,6 +328,7 @@ cc_library(
         ":codegen_utils",
         ":disc_lhlo_elemental_utils",
         ":disc_supported_list",
+        ":disc_util",
         ":fusion_utils",
         ":input_inline_fusion_pattern",
         ":pass_details",
@@ -527,6 +529,7 @@ cc_library(
     hdrs = ["transforms/passes.h"],
     deps = [
         ":codegen_utils",
+        ":disc_util",
         ":fusion_utils",
         ":pass_details",
         ":placement_utils",
@@ -676,6 +679,7 @@ cc_library(
     srcs = ["transforms/shape_utils.cc"],
     hdrs = ["transforms/shape_utils.h"],
     deps = [
+        ":disc_shape_optimization_utils",
         ":disc_util",
         ":mhlo_disc",
         ":pass_details",
@@ -721,6 +725,7 @@ cc_library(
     srcs = ["transforms/lhlo_fusion.cc"],
     deps = [
         ":disc_shape_optimization_utils",
+        ":disc_util",
         ":fusion_utils",
         ":pass_details",
         ":placement_utils",
@@ -1517,6 +1522,7 @@ cc_library(
     deps = [
         ":codegen_utils",
         ":disc_lhlo_elemental_utils",
+        ":disc_util",
         ":fusion_utils",
         ":pass_details",
         ":placement_utils",
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index 457a148668d..3c8eb5b381b 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -189,6 +189,8 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
 
   DefaultTimingManager tm;
   applyDefaultTimingManagerCLOptions(tm);
+  // Records elapsed time for each pass in the passpipe
+  tm.setEnabled(true);
   TimingScope timing = tm.getRootScope();
   PassManager pm(m.getContext());
   applyPassManagerCLOptions(pm);
@@ -228,10 +230,7 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(createCSEPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
 
-  bool enable_shape_constraint_ir = false;
-  tensorflow::ReadBoolFromEnvVar("DISC_ENABLE_SHAPE_CONSTRAINT_IR",
-                                 enable_shape_constraint_ir,
-                                 &enable_shape_constraint_ir);
+  bool enable_shape_constraint_ir = useShapeConstraintIR();
   if (!enable_shape_constraint_ir) {
     // propagate some known shape information.
     pm.addPass(disc_ral::createDiscShapeSimplifierPass());
@@ -553,8 +552,6 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
 
   if (failed(pm.run(m))) return failure();
-
-  TimingScope outputTiming = timing.nest("Output");
   return success();
 }
 
diff --git a/tao_compiler/mlir/disc/disc_util.cc b/tao_compiler/mlir/disc/disc_util.cc
index 859bcdc8837..5e7b88731cf 100644
--- a/tao_compiler/mlir/disc/disc_util.cc
+++ b/tao_compiler/mlir/disc/disc_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "tensorflow/compiler/mlir/disc/transforms/placement_utils.h"
+#include "tensorflow/core/util/env_var.h"
 
 namespace mlir {
 namespace disc_ral {
@@ -189,5 +190,16 @@ SmallVector<Value, 4> GetAllPossibleUsedValues(Operation* op) {
   return values;
 }
 
+bool useShapeConstraintIR() {
+  static bool enabled = []() {
+    bool enable_shape_constraint_ir = false;
+    tensorflow::ReadBoolFromEnvVar("DISC_ENABLE_SHAPE_CONSTRAINT_IR",
+                                   enable_shape_constraint_ir,
+                                   &enable_shape_constraint_ir);
+    return enable_shape_constraint_ir;
+  }();
+  return enabled;
+}
+
 }  // namespace disc_ral
 }  // namespace mlir
diff --git a/tao_compiler/mlir/disc/disc_util.h b/tao_compiler/mlir/disc/disc_util.h
index 9172902e754..d204fdf6945 100644
--- a/tao_compiler/mlir/disc/disc_util.h
+++ b/tao_compiler/mlir/disc/disc_util.h
@@ -90,6 +90,9 @@ llvm::Optional<int32_t> TryMergeNode(GraphCycles* graph_cycles, int32_t a,
 
 SmallVector<Value, 4> GetAllPossibleUsedValues(Operation* op);
 
+// Returns true if the shape constraint IR is enabled.
+bool useShapeConstraintIR();
+
 }  // namespace disc_ral
 }  // namespace mlir
 
diff --git a/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc b/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
index 8604ff24ec2..3628051a55f 100644
--- a/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
@@ -481,8 +481,7 @@ class DotBatchMergeConverter {
     // TODO: use shape-symbol rather than dim-value of m/n/k in this struct
     // after reconstructing ShapeAnalysis. We actually do not need to
     // distinguish m/n/k but only need to use the overall shape-symbol here in
-    // the future. We do not use shape-symbol currently because of the
-    // incomplete implementation of currently ShapeAnalysis data structure...
+    // the future.
     DimValue m_dim;
     DimValue n_dim;
     DimValue contracting_dim;
@@ -535,7 +534,8 @@ class DotBatchMergeConverter {
                          MergingShapeHash>;
 
  private:
-  bool buildMergingShapeMap(Block* block, ShapeAnalysis& analysis,
+  bool buildMergingShapeMap(Block* block,
+                            ShapeAnalysisDeprecated& analysisDeprecated,
                             MergingShapeEqualMap& equal_merge_shape_map);
   bool applyMerging(DotCluster& cluster);
   Value expandDim0(OpBuilder& builder, Location& loc, Value value);
@@ -545,9 +545,10 @@ class DotBatchMergeConverter {
 };
 
 bool DotBatchMergeConverter::run() {
-  ShapeAnalysis analysis(func_);
-  if (failed(analysis.run())) {
-    LLVM_DEBUG(llvm::dbgs() << "ShapeAnalysis failes for dot merge.\n");
+  ShapeAnalysisDeprecated analysisDeprecated(func_);
+  if (failed(analysisDeprecated.run())) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "ShapeAnalysisDeprecated failes for dot merge.\n");
     return false;
   }
 
@@ -557,7 +558,8 @@ bool DotBatchMergeConverter::run() {
   for (Block* block : blocks) {
     // A map to help to cluster dots with same shape and dim-numbers together.
     MergingShapeEqualMap equal_merge_shape_map;
-    if (!buildMergingShapeMap(block, analysis, equal_merge_shape_map)) {
+    if (!buildMergingShapeMap(block, analysisDeprecated,
+                              equal_merge_shape_map)) {
       continue;
     }
     // Find merging clusters.
@@ -574,7 +576,7 @@ bool DotBatchMergeConverter::run() {
 }
 
 bool DotBatchMergeConverter::buildMergingShapeMap(
-    Block* block, ShapeAnalysis& analysis,
+    Block* block, ShapeAnalysisDeprecated& analysisDeprecated,
     MergingShapeEqualMap& equal_merge_shape_map) {
   block->walk([&](mhlo::DotGeneralOp op) {
     MergingShape dot_shape;
@@ -590,14 +592,15 @@ bool DotBatchMergeConverter::buildMergingShapeMap(
     // Initialize `batching_dims`.
     SmallVector<DimValue>& batching_dims = dot_shape.batching_dims;
     for (auto dim : lhs_batch_dims) {
-      batching_dims.emplace_back(std::move(analysis.getDimValue(lhs, dim)));
+      batching_dims.emplace_back(
+          std::move(analysisDeprecated.getDimValue(lhs, dim)));
     }
     // Initialize `contracting_dim`.
     auto lhs_contracting_dims =
         dot_shape.dimension_numbers.getLhsContractingDimensions();
     assert(lhs_contracting_dims.size() == 1);
     dot_shape.contracting_dim =
-        analysis.getDimValue(lhs, lhs_contracting_dims[0]);
+        analysisDeprecated.getDimValue(lhs, lhs_contracting_dims[0]);
     // Initialize `m_dim`.
     int64_t lhs_rank = lhs.getType().cast<RankedTensorType>().getRank();
     assert(lhs_batch_dims.size() + 2 == lhs_rank);
@@ -606,7 +609,7 @@ bool DotBatchMergeConverter::buildMergingShapeMap(
     for (int64_t i = 0; i < lhs_rank; i++) {
       if ((lhs_batch_dims_set.find(i) == lhs_batch_dims_set.end()) &&
           (i != lhs_contracting_dims[0])) {
-        dot_shape.m_dim = analysis.getDimValue(lhs, i);
+        dot_shape.m_dim = analysisDeprecated.getDimValue(lhs, i);
         break;
       }
     }
@@ -621,7 +624,7 @@ bool DotBatchMergeConverter::buildMergingShapeMap(
     for (int64_t i = 0; i < rhs_rank; i++) {
       if ((rhs_batch_dims_set.find(i) == rhs_batch_dims_set.end()) &&
           (i != rhs_contracting_dims[0])) {
-        dot_shape.n_dim = analysis.getDimValue(rhs, i);
+        dot_shape.n_dim = analysisDeprecated.getDimValue(rhs, i);
         break;
       }
     }
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
index 32a5efd12c4..1a12a42defa 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
@@ -535,6 +535,12 @@ struct CopyLikeOpConvertor : public OpRewritePattern<OpTy> {
 
     Value newValue = CastMemRefTo(rewriter, loc, dispatchOp->getResult(0),
                                   targetType, dimSizes);
+    if (Operation* resultDefiningOp = result.getDefiningOp()) {
+      auto attrName = disc_shape::SymbolicDimOp::getSymbolicDimAttrName();
+      if (resultDefiningOp->hasAttr(attrName))
+        newValue.getDefiningOp()->setAttr(attrName,
+                                          resultDefiningOp->getAttr(attrName));
+    }
 
     // replace non-shape-consumer users of original result
     for (Operation* user : users) user->replaceUsesOfWith(result, newValue);
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc b/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
index ba25bd59497..669d43cd685 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 // This file implements the logic to do some shape optimizations on tensor
 // level.
+#include <chrono>
 #include <unordered_set>
 #include <utility>
 
@@ -40,6 +41,8 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/disc_shape_optimization_utils.h"
 
+#define DEBUG_TYPE "disc-shape-optimization"
+
 namespace mlir {
 namespace disc_ral {
 
@@ -1281,37 +1284,88 @@ LogicalResult optimizeShapeComputation(ModuleOp m, FuncOp main,
   bool changed;
   do {
     changed = false;
+    LLVM_DEBUG(std::chrono::steady_clock::time_point begin =
+                   std::chrono::steady_clock::now());
     if (failed(runCanonicalizer(m, runner))) {
       return failure();
     }
+    LLVM_DEBUG(std::chrono::steady_clock::time_point end =
+                   std::chrono::steady_clock::now());
+    LLVM_DEBUG(llvm::dbgs()
+               << "  runCanonicalizer takes: "
+               << std::chrono::duration_cast<std::chrono::microseconds>(end -
+                                                                        begin)
+                      .count()
+               << " us\n");
 
     LLVM_DEBUG(
         llvm::dbgs()
         << "Module after runCanonicalizer in optimize-shape-computation:\n"
         << m << "\n");
 
+    LLVM_DEBUG(begin = std::chrono::steady_clock::now());
     SymbolicDimMgr mgr(m);
+    LLVM_DEBUG(end = std::chrono::steady_clock::now());
+    LLVM_DEBUG(llvm::dbgs()
+               << "  Building SymbolicDimMgr takes: "
+               << std::chrono::duration_cast<std::chrono::microseconds>(end -
+                                                                        begin)
+                      .count()
+               << " us\n");
+
+    LLVM_DEBUG(begin = std::chrono::steady_clock::now());
     if (failed(mgr.load())) {
       return m.emitError() << "fail to load shape constraint IR\n";
     }
+    LLVM_DEBUG(end = std::chrono::steady_clock::now());
+    LLVM_DEBUG(llvm::dbgs()
+               << "  SymbolicDimMgr.load() takes: "
+               << std::chrono::duration_cast<std::chrono::microseconds>(end -
+                                                                        begin)
+                      .count()
+               << " us\n");
 
+    LLVM_DEBUG(begin = std::chrono::steady_clock::now());
     ShapeComputationIRAnalysis analysis(main, mgr);
     if (failed(analysis.run())) {
       return m.emitError() << "fail to analysis shape computation IR\n";
     }
+    LLVM_DEBUG(end = std::chrono::steady_clock::now());
+    LLVM_DEBUG(llvm::dbgs()
+               << "  Building ShapeComputationIRAnalysis takes: "
+               << std::chrono::duration_cast<std::chrono::microseconds>(end -
+                                                                        begin)
+                      .count()
+               << " us\n");
 
+    LLVM_DEBUG(begin = std::chrono::steady_clock::now());
     if (failed(applyShapeComputationOptimization(analysis, changed))) {
       return m.emitError() << "fail to optimize shape computation IR\n";
     }
+    LLVM_DEBUG(end = std::chrono::steady_clock::now());
+    LLVM_DEBUG(llvm::dbgs()
+               << "  applyShapeComputationOptimization takes: "
+               << std::chrono::duration_cast<std::chrono::microseconds>(end -
+                                                                        begin)
+                      .count()
+               << " us\n");
 
     LLVM_DEBUG(
         llvm::dbgs()
         << "Module after apply-shape-opt in optimize-shape-computation:\n"
         << m << "\n");
 
+    LLVM_DEBUG(begin = std::chrono::steady_clock::now());
     if (failed(mgr.save())) {
       return m.emitError() << "fail to save shape constraint IR\n";
     }
+    LLVM_DEBUG(end = std::chrono::steady_clock::now());
+    LLVM_DEBUG(llvm::dbgs()
+               << "  SymbolicDimMgr.save() takes: "
+               << std::chrono::duration_cast<std::chrono::microseconds>(end -
+                                                                        begin)
+                      .count()
+               << " us\n");
 
     LLVM_DEBUG(llvm::dbgs()
                << "Module after save-shape-ir in optimize-shape-computation:\n"
@@ -1443,15 +1497,26 @@ void DiscShapeOptimizationPass::runOnOperation() {
     return;
   }
 
+  LLVM_DEBUG(std::chrono::steady_clock::time_point begin =
+                 std::chrono::steady_clock::now());
   // Stage #1: Explictily materialize shape computation IR on tensor level
   if (failed(materializeShapeComputation(m, main))) {
     signalPassFailure();
     return;
   }
+  LLVM_DEBUG(std::chrono::steady_clock::time_point end =
+                 std::chrono::steady_clock::now());
+  LLVM_DEBUG(
+      llvm::dbgs() << "materializeShapeComputation takes: "
+                   << std::chrono::duration_cast<std::chrono::microseconds>(
+                          end - begin)
+                          .count()
+                   << " us\n");
   LLVM_DEBUG(llvm::dbgs() << "Module after materialize shape computation:\n"
                           << m << "\n");
 
   // Stage #2: Optimize shape computation IR on tensor level
+  LLVM_DEBUG(begin = std::chrono::steady_clock::now());
   PassPipelineRunner runner = [this](OpPassManager& dynamicPM, ModuleOp m) {
     return runPipeline(dynamicPM, m);
   };
@@ -1459,13 +1524,28 @@ void DiscShapeOptimizationPass::runOnOperation() {
     signalPassFailure();
     return;
   }
+  LLVM_DEBUG(end = std::chrono::steady_clock::now());
+  LLVM_DEBUG(
+      llvm::dbgs() << "optimizeShapeComputation takes: "
+                   << std::chrono::duration_cast<std::chrono::microseconds>(
+                          end - begin)
+                          .count()
+                   << " us\n");
   LLVM_DEBUG(llvm::dbgs() << "Module after shape optimizaiton:\n" << m << "\n");
 
   // Stage #3: clean up
+  LLVM_DEBUG(begin = std::chrono::steady_clock::now());
   if (failed(cleanUp(m, keep_tie_shape_))) {
     signalPassFailure();
     return;
   }
+  LLVM_DEBUG(end = std::chrono::steady_clock::now());
+  LLVM_DEBUG(
+      llvm::dbgs() << "cleanUp takes: "
+                   << std::chrono::duration_cast<std::chrono::microseconds>(
+                          end - begin)
+                          .count()
+                   << " us\n");
   LLVM_DEBUG(llvm::dbgs() << "Module after cleanup:\n" << m << "\n");
 }
 
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
index 55bc06dc530..533510097f8 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpDefinition.h"
@@ -31,6 +32,8 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/disc/IR/disc_shape_ops.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 
+#define DEBUG_TYPE "disc-shape-optimization-utils"
+
 namespace mlir {
 namespace disc_ral {
 
@@ -74,7 +77,7 @@ llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
   return os;
 }
 
-SymbolicDimMgr::SymbolicDimMgr(ModuleOp m) : m_(m) {}
+SymbolicDimMgr::SymbolicDimMgr(ModuleOp m) : m_(m), symbolTable_(m_) {}
 
 LogicalResult SymbolicDimMgr::load() {
   m_.walk([&](disc_shape::SymbolicDimOp op) {
@@ -92,29 +95,35 @@ std::string SymbolicDimMgr::getNextName() {
   return name;
 }
 
-SymbolicDimOp SymbolicDimMgr::newSymbolicDim() {
+SymbolicDimOp SymbolicDimMgr::newSymbolicDim(StringRef name) {
   OpBuilder b(m_);
-  auto symbol = b.create<SymbolicDimOp>(m_.getLoc(), getNextName());
+  auto symbol =
+      b.create<SymbolicDimOp>(m_.getLoc(), name.empty() ? getNextName() : name);
   symbolDimUnionSet_[symbol] = symbol;
-  m_.push_back(symbol);
+  symbolTable_.insert(symbol);
   return symbol;
 }
 
 SymbolicDimOp SymbolicDimMgr::newConstantSymbolicDim(int64_t val) {
   auto it = constantSymbolicDimMap_.find(val);
   if (it == constantSymbolicDimMap_.end()) {
-    it = constantSymbolicDimMap_.insert(std::make_pair(val, newSymbolicDim()))
+    auto name = (llvm::Twine("C") + llvm::Twine(val)).str();
+    it = constantSymbolicDimMap_
+             .insert(std::make_pair(val, newSymbolicDim(name)))
              .first;
     it->second.setDimSize(val);
-    it->second.setName((llvm::Twine("C") + llvm::Twine(val)).str());
   }
   return getRootSymbolicDim(it->second);
 }
 
 SymbolicDimOp SymbolicDimMgr::getRootSymbolicDim(SymbolicDimOp symbol) {
   SymbolicDimOp current = symbol;
-  while (symbolDimUnionSet_[current] != current)
+  SmallVector<SymbolicDimOp> path;
+  while (symbolDimUnionSet_[current] != current) {
+    path.push_back(current);
     current = symbolDimUnionSet_[current];
+  }
+  for (SymbolicDimOp sym : path) symbolDimUnionSet_[sym] = current;
   return current;
 }
 
@@ -137,6 +146,7 @@ LogicalResult SymbolicDimMgr::mapSymbolicDimEqual(SymbolicDimOp lhs,
       if (failed(rhsRoot.Merge(lhsRoot))) return failure();
       symbolDimUnionSet_[lhsRoot] = rhsRoot;
     }
+    productEqualityMapUpdated_ = false;
   }
   return success();
 }
@@ -151,6 +161,7 @@ int64_t gcd(int64_t m, int64_t n) {
 SymbolicDimProduct SymbolicDimMgr::simplifySymbolicDimProduct(
     const SymbolicDimProduct& x) {
   SmallVector<SymbolicDimOp> copied;
+  copied.reserve(x.symbols.size());
   // TODO(disc): expand op if it can be factorized into other symbol dims.
   for (SymbolicDimOp op : x.symbols) copied.push_back(getRootSymbolicDim(op));
 
@@ -317,10 +328,14 @@ LogicalResult SymbolicDimMgr::mapSymbolicDimProductEqual(
   productEqualityMap_[newLhs][newRhs] = productEqualityMap_[newRhs][newLhs] =
       true;
 
-  return updateProductEqualityMap();
+  productEqualityMapUpdated_ = false;
+  return success();
 }
 
 LogicalResult SymbolicDimMgr::updateProductEqualityMap() {
+  // early return if nothing is updated.
+  if (productEqualityMapUpdated_) return success();
+
   SymbolicDimProductMap newMap;
   DenseSet<SymbolicDimProduct> productSet;
   for (auto& pairOutter : productEqualityMap_) {
@@ -385,30 +400,50 @@ LogicalResult SymbolicDimMgr::updateProductEqualityMap() {
       LLVM_DEBUG(llvm::dbgs() << "Pair: x = " << x << "\ny = " << y << "\n");
     }
   }
+  productEqualityMapUpdated_ = true;
   return success();
 }
 
 LogicalResult SymbolicDimMgr::save() {
   // replace all uses of a symbolic dim op with its root symbolic dim op
-  if (failed(walkRankedTensorValue(m_, [&](Value value, RankedTensorType ty,
-                                           ArrayAttr attrs) {
-        SmallVector<Attribute> newAttrs;
-        for (Attribute attr : attrs) {
-          auto sym =
-              m_.lookupSymbol<SymbolicDimOp>(attr.cast<FlatSymbolRefAttr>());
-          assert(sym);
-          SymbolicDimOp root = getRootSymbolicDim(sym);
-          FlatSymbolRefAttr rootSymbol = FlatSymbolRefAttr::get(root);
-          newAttrs.push_back(rootSymbol);
-        }
-        auto symbolicShapeAttr = ArrayAttr::get(value.getContext(), newAttrs);
-        auto newTy = RankedTensorType::get(ty.getShape(), ty.getElementType(),
-                                           symbolicShapeAttr);
-        value.setType(newTy);
-        return success();
-      }))) {
+  using Name2SymbolFn = std::function<SymbolicDimOp(StringRef)>;
+  auto updateAttrs = [&](ArrayAttr attrs, Name2SymbolFn fn) {
+    SmallVector<Attribute> newAttrs;
+    for (Attribute attr : attrs) {
+      auto sym = fn(attr.cast<FlatSymbolRefAttr>().getValue());
+      assert(sym);
+      SymbolicDimOp root = getRootSymbolicDim(sym);
+      FlatSymbolRefAttr rootSymbol = FlatSymbolRefAttr::get(root);
+      newAttrs.push_back(rootSymbol);
+    }
+    return ArrayAttr::get(m_->getContext(), newAttrs);
+  };
+  // update attributes attached in values.
+  if (failed(walkRankedTensorValue(
+          m_, [&](Value value, RankedTensorType ty, ArrayAttr attrs) {
+            auto symbolicShapeAttr = updateAttrs(attrs, [&](StringRef name) {
+              return symbolTable_.lookup<SymbolicDimOp>(name);
+            });
+            auto newTy = RankedTensorType::get(
+                ty.getShape(), ty.getElementType(), symbolicShapeAttr);
+            value.setType(newTy);
+            return success();
+          }))) {
     return failure();
   }
+  // update attributes attached in operations.
+  m_.walk([&](Operation* op) {
+    auto attrs =
+        op->getAttrOfType<ArrayAttr>(SymbolicDimOp::getSymbolicDimAttrName());
+    if (!attrs) return;
+    auto symbolicShapeAttr = updateAttrs(attrs, [&](StringRef name) {
+      return symbolTable_.lookup<SymbolicDimOp>(name);
+    });
+    op->setAttr(SymbolicDimOp::getSymbolicDimAttrName(), symbolicShapeAttr);
+  });
+  // update shape production equality
+  if (failed(updateProductEqualityMap()))
+    return m_->emitError() << "fail to update prodcut euqal map\n";
 
   // Update function type
   if (failed(updateFunctionType(m_))) return failure();
@@ -416,26 +451,57 @@ LogicalResult SymbolicDimMgr::save() {
   // collect symbolic dim ops that are referred by other ops/types.
   DenseSet<SymbolicDimOp> usedSymbolicOps;
   SmallVector<std::string> usedSymbolNames;
-  if (failed(walkRankedTensorValue(m_, [&](Value value, RankedTensorType ty,
-                                           ArrayAttr attrs) {
-        SmallVector<Attribute> newAttrs;
-        for (Attribute attr : attrs) {
-          auto sym =
-              m_.lookupSymbol<SymbolicDimOp>(attr.cast<FlatSymbolRefAttr>());
-          assert(sym);
-          if (usedSymbolicOps.insert(sym).second)
-            usedSymbolNames.push_back(sym.getName().str());
-        }
-        return success();
-      }))) {
+  // collect uses in values.
+  auto collectUsedSymbols = [&](ArrayAttr attrs) {
+    for (Attribute attr : attrs) {
+      auto sym = symbolTable_.lookup<SymbolicDimOp>(
+          attr.cast<FlatSymbolRefAttr>().getValue());
+      assert(sym);
+      if (usedSymbolicOps.insert(sym).second)
+        usedSymbolNames.push_back(sym.getName().str());
+    }
+  };
+  if (failed(walkRankedTensorValue(
+          m_, [&](Value value, RankedTensorType ty, ArrayAttr attrs) {
+            collectUsedSymbols(attrs);
+            return success();
+          }))) {
     return failure();
   }
+  // collect uses in operations.
+  m_.walk([&](Operation* op) {
+    auto attrs =
+        op->getAttrOfType<ArrayAttr>(SymbolicDimOp::getSymbolicDimAttrName());
+    if (!attrs) return;
+    collectUsedSymbols(attrs);
+  });
 
   // remove symbolic dim ops that are known not used by any other ops/types.
   for (auto& p : symbolDimUnionSet_) {
     if (!usedSymbolicOps.count(p.first)) p.first->erase();
   }
 
+  // remove unused shape production equality
+  SmallVector<SymbolicDimProduct> candidates;
+  for (auto& outter : productEqualityMap_) {
+    if (llvm::any_of(outter.first.symbols, [&](SymbolicDimOp sym) {
+          return usedSymbolicOps.count(sym) == 0;
+        }))
+      candidates.push_back(outter.first);
+  }
+  for (auto& prod : candidates) productEqualityMap_.erase(prod);
+
+  for (auto& outter : productEqualityMap_) {
+    SmallVector<SymbolicDimProduct> candidates;
+    for (auto& inner : outter.second) {
+      if (llvm::any_of(inner.first.symbols, [&](SymbolicDimOp sym) {
+            return usedSymbolicOps.count(sym) == 0;
+          }))
+        candidates.push_back(outter.first);
+    }
+    for (auto& prod : candidates) outter.second.erase(prod);
+  }
+
   // canonicalize the name of symbolic dim ops
   llvm::sort(usedSymbolNames,
              [&](const std::string& lhs, const std::string& rhs) {
@@ -443,34 +509,41 @@ LogicalResult SymbolicDimMgr::save() {
              });
   int numNonConstDims = 0;
   std::unordered_map<std::string, std::string> nameMapping;
-  for (const auto& en : llvm::enumerate(usedSymbolNames)) {
-    if (en.value().size() > 0 && en.value()[0] == 'C') {
-      nameMapping[en.value()] = en.value();
+  for (const auto& name : usedSymbolNames) {
+    if (name.size() > 0 && name[0] == 'C') {
+      nameMapping[name] = name;
     } else {
-      nameMapping[en.value()] =
+      nameMapping[name] =
           (llvm::Twine("S") + llvm::Twine(numNonConstDims++)).str();
     }
   }
-  for (SymbolicDimOp op : usedSymbolicOps)
-    op.setName(nameMapping[op.getName().str()]);
+  std::unordered_map<std::string, SymbolicDimOp> name2Symbol;
+  for (SymbolicDimOp op : usedSymbolicOps) {
+    auto name = op.getName().str();
+    op.setName(nameMapping[name]);
+    name2Symbol[name] = op;
+  }
 
   // replace the name of a symbolic dim op to its new name.
-  if (failed(walkRankedTensorValue(m_, [&](Value value, RankedTensorType ty,
-                                           ArrayAttr attrs) {
-        SmallVector<Attribute> newAttrs;
-        for (Attribute attr : attrs) {
-          auto sym = m_.lookupSymbol<SymbolicDimOp>(
-              nameMapping[attr.cast<FlatSymbolRefAttr>().getValue().str()]);
-          assert(sym);
-          FlatSymbolRefAttr newAttr = FlatSymbolRefAttr::get(sym);
-          newAttrs.push_back(newAttr);
-        }
-        auto symbolicShapeAttr = ArrayAttr::get(value.getContext(), newAttrs);
-        auto newTy = RankedTensorType::get(ty.getShape(), ty.getElementType(),
-                                           symbolicShapeAttr);
-        value.setType(newTy);
-        return success();
-      }))) {
+  // update attributes attached to operations.
+  m_.walk([&](Operation* op) {
+    auto attrs =
+        op->getAttrOfType<ArrayAttr>(SymbolicDimOp::getSymbolicDimAttrName());
+    if (!attrs) return;
+    auto symbolicShapeAttr = updateAttrs(
+        attrs, [&](StringRef name) { return name2Symbol[name.str()]; });
+    op->setAttr(SymbolicDimOp::getSymbolicDimAttrName(), symbolicShapeAttr);
+  });
+  // update attributes attached to values.
+  if (failed(walkRankedTensorValue(
+          m_, [&](Value value, RankedTensorType ty, ArrayAttr attrs) {
+            auto symbolicShapeAttr = updateAttrs(
+                attrs, [&](StringRef name) { return name2Symbol[name.str()]; });
+            auto newTy = RankedTensorType::get(
+                ty.getShape(), ty.getElementType(), symbolicShapeAttr);
+            value.setType(newTy);
+            return success();
+          }))) {
     return failure();
   }
 
@@ -486,12 +559,9 @@ LogicalResult SymbolicDimMgr::save() {
 }
 
 LogicalResult SymbolicDimMgr::saveShapeConstraintGraph() {
-  if (failed(updateProductEqualityMap()))
-    return m_->emitError() << "fail to update prodcut euqal map\n";
-
   // first try to remove the old shape constraint graph
   StringRef funcName = getShapeConstraintGraphFunctionName();
-  if (auto func = m_.lookupSymbol<FuncOp>(funcName)) func->erase();
+  if (auto func = symbolTable_.lookup<FuncOp>(funcName)) func->erase();
 
   OpBuilder moduleBuilder(m_);
   auto funcTy = FunctionType::get(m_.getContext(), {}, {});
@@ -534,7 +604,7 @@ LogicalResult SymbolicDimMgr::saveShapeConstraintGraph() {
 
 LogicalResult SymbolicDimMgr::loadShapeConstraintGraph() {
   StringRef funcName = getShapeConstraintGraphFunctionName();
-  auto func = m_.lookupSymbol<FuncOp>(funcName);
+  auto func = symbolTable_.lookup<FuncOp>(funcName);
   if (!func) return success();
 
   auto build_sym_product = [&](ValueRange range,
@@ -545,7 +615,7 @@ LogicalResult SymbolicDimMgr::loadShapeConstraintGraph() {
         product.factor *= constOp->getAttrOfType<IntegerAttr>("value").getInt();
         continue;
       } else if (auto dimOp = dyn_cast_or_null<disc_shape::DimOp>(definingOp)) {
-        auto sym = m_.lookupSymbol<SymbolicDimOp>(dimOp.name());
+        auto sym = symbolTable_.lookup<SymbolicDimOp>(dimOp.name());
         if (!sym) return dimOp->emitError() << "fail to find symbolic dim op\n";
         product.symbols.push_back(sym);
         continue;
@@ -578,7 +648,8 @@ SymbolicDimMgr::getOrCreateSymbolicDimsForRankedValue(Value value) {
                        .getEncoding()
                        .dyn_cast_or_null<ArrayAttr>()) {
     for (Attribute attr : attrs) {
-      auto sym = m_.lookupSymbol<SymbolicDimOp>(attr.cast<FlatSymbolRefAttr>());
+      auto sym = symbolTable_.lookup<SymbolicDimOp>(
+          attr.cast<FlatSymbolRefAttr>().getValue());
       assert(sym);
       symbols.push_back(sym);
     }
@@ -594,6 +665,72 @@ SymbolicDimMgr::getOrCreateSymbolicDimsForRankedValue(Value value) {
   return symbols;
 }
 
+SymbolicDimOp SymbolicDimMgr::cloneSymbol(SymbolicDimOp symbol) {
+  // early return for static shape symbol
+  if (!symbol.isDynamic()) return symbol;
+  SymbolicDimOp newSymbol = newSymbolicDim();
+
+  // copied all attributes expect the name.
+  for (const NamedAttribute& attr : symbol->getAttrs()) {
+    StringRef name = attr.getName().strref();
+    if (name == "sym_name") continue;
+    newSymbol->setAttr(name, attr.getValue());
+  }
+  return newSymbol;
+}
+
+LogicalResult SymbolicDimMgr::cloneSymbolGroup(
+    const DenseSet<SymbolicDimOp>& symbols,
+    DenseMap<SymbolicDimOp, SymbolicDimOp>& mapping) {
+  DenseMap<SymbolicDimOp, SmallVector<SymbolicDimOp>> root2Symbol;
+  for (SymbolicDimOp symbol : symbols) {
+    SymbolicDimOp newSymbol = cloneSymbol(symbol);
+    mapping[symbol] = newSymbol;
+    SymbolicDimOp root = getRootSymbolicDim(symbol);
+    root2Symbol[root].push_back(newSymbol);
+  }
+
+  // copy dim equality predicate.
+  for (auto& it : root2Symbol) {
+    if (it.second.size() <= 1) continue;
+    for (SymbolicDimOp symbol : it.second)
+      if (failed(mapSymbolicDimEqual(symbol, it.second[0]))) return failure();
+  }
+
+  // copy dim product equality predicate
+  SmallVector<std::pair<SymbolicDimProduct, SymbolicDimProduct>> candidates;
+  for (auto& outter : productEqualityMap_) {
+    if (llvm::any_of(outter.first.symbols, [&](SymbolicDimOp sym) {
+          return symbols.count(sym) == 0;
+        }))
+      continue;
+    for (auto& inner : outter.second) {
+      if (!inner.second) continue;
+      if (llvm::any_of(inner.first.symbols, [&](SymbolicDimOp sym) {
+            return symbols.count(sym) == 0;
+          }))
+        continue;
+      auto copySymbolicDimProduct = [&](SymbolicDimProduct& prod) {
+        SymbolicDimProduct newProd;
+        newProd.factor = prod.factor;
+        for (SymbolicDimOp sym : prod.symbols) {
+          newProd.symbols.push_back(mapping[sym]);
+        }
+        return newProd;
+      };
+      candidates.emplace_back(copySymbolicDimProduct(outter.first),
+                              copySymbolicDimProduct(inner.first));
+    }
+  }
+
+  for (auto& it : candidates) {
+    productEqualityMap_[it.first][it.second] = true;
+    productEqualityMap_[it.second][it.first] = true;
+  }
+
+  return success();
+}
+
 llvm::Optional<SmallVector<FlatSymbolRefAttr>> getRankedValueSymbolicDimRefs(
     Value value) {
   auto ty = value.getType().dyn_cast<RankedTensorType>();
@@ -610,6 +747,43 @@ llvm::Optional<SmallVector<FlatSymbolRefAttr>> getRankedValueSymbolicDimRefs(
   return symbols;
 }
 
+llvm::Optional<SmallVector<FlatSymbolRefAttr>> getMemRefValueSymbolicDimRefs(
+    Value value) {
+  auto ty = value.getType().dyn_cast<MemRefType>();
+  Operation* op = value.getDefiningOp();
+  if (!ty || !op) return {};
+  auto attrs = op->getAttrOfType<ArrayAttr>(
+      disc_shape::SymbolicDimOp::getSymbolicDimAttrName());
+  if (!attrs || attrs.size() != ty.getRank()) return {};
+
+  SmallVector<FlatSymbolRefAttr> symbols;
+  symbols.reserve(attrs.size());
+  for (const auto& attr : attrs) {
+    auto symbol = attr.dyn_cast<FlatSymbolRefAttr>();
+    if (!symbol) return {};
+    symbols.push_back(symbol);
+  }
+  return symbols;
+}
+
+ArrayAttr makeSymbolicDimOpRefArrayAttr(
+    const SmallVector<SymbolicDimOp>& symbols) {
+  assert(!symbols.empty());
+  SmallVector<Attribute> attrs;
+  attrs.reserve(symbols.size());
+  for (auto& sym : symbols) {
+    attrs.push_back(FlatSymbolRefAttr::get(sym));
+  }
+  return ArrayAttr::get(symbols[0]->getContext(), attrs);
+}
+
+void attachSymbolicDimOpRefArrayAttrOnOperation(
+    Operation* op, const SmallVector<SymbolicDimOp>& symbols) {
+  assert(op != nullptr);
+  auto symbolicShapeAttr = makeSymbolicDimOpRefArrayAttr(symbols);
+  op->setAttr(SymbolicDimOp::getSymbolicDimAttrName(), symbolicShapeAttr);
+}
+
 // Updates the function types according to the types of entry block arguments
 // and the types of operands of the return op of the func op. This function
 // suppose that there is only one block inside the function region.
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.h b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.h
index f29e21aad64..60fd021e0ba 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.h
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.h
@@ -31,6 +31,21 @@ using disc_shape::SymbolicDimOp;
 // attached symbolic dim ref attributes.
 llvm::Optional<SmallVector<FlatSymbolRefAttr>> getRankedValueSymbolicDimRefs(
     Value value);
+// Return the symbolicDim ref attribute if there is an attached disc
+// shape-constraint specific attribute filed. Return nullptr if there isn't an
+// attached symbolic dim ref attributes.
+llvm::Optional<SmallVector<FlatSymbolRefAttr>> getMemRefValueSymbolicDimRefs(
+    Value value);
+
+// Returns a ArrayAttr composed of a list of ref attributes to corresponding
+// symbols;
+ArrayAttr makeSymbolicDimOpRefArrayAttr(
+    const SmallVector<SymbolicDimOp>& symbols);
+
+// Inserts or updates disc symbolic dim attr of the op. The  disc symbolic dim
+// is created from the provided symbols.
+void attachSymbolicDimOpRefArrayAttrOnOperation(
+    Operation* op, const SmallVector<SymbolicDimOp>& symbols);
 
 using Visitor = std::function<LogicalResult(Value value, RankedTensorType ty,
                                             ArrayAttr attrs)>;
@@ -157,7 +172,7 @@ class SymbolicDimMgr {
 
   // Returns a new symbolicDim instance. The returned symbolicDim is owned by
   // this mgr.
-  SymbolicDimOp newSymbolicDim();
+  SymbolicDimOp newSymbolicDim(StringRef name = {});
 
   // Returns a symbolicDim which have static dim size == `val`.
   SymbolicDimOp newConstantSymbolicDim(int64_t val);
@@ -212,6 +227,17 @@ class SymbolicDimMgr {
   // Returns the name of the shape constraint graph
   static StringRef getShapeConstraintGraphFunctionName();
 
+  // Returns a clone of the original symbol
+  SymbolicDimOp cloneSymbol(SymbolicDimOp symbol);
+
+  // Clones a group of symbols and the relationships among the symbols in the
+  // group. Returns ok if success, otherwise failure.
+  LogicalResult cloneSymbolGroup(
+      const DenseSet<SymbolicDimOp>& symbols,
+      DenseMap<SymbolicDimOp, SymbolicDimOp>& mapping);
+
+  SymbolTable& symbolTable() { return symbolTable_; }
+
  private:
   // Returns next unique name for a new SymbolicDim op.
   std::string getNextName();
@@ -235,6 +261,9 @@ class SymbolicDimMgr {
   // The module this SymbolicDimMgr runs on.
   ModuleOp m_;
 
+  // A symbol table corresponding to related module op.
+  SymbolTable symbolTable_;
+
   // A unique id to generate unique name.
   int64_t nextSymbolicOpIdx_ = 0;
 
@@ -258,6 +287,8 @@ class SymbolicDimMgr {
   using SymbolicDimProductMap =
       DenseMap<SymbolicDimProduct, DenseMap<SymbolicDimProduct, bool>>;
   SymbolicDimProductMap productEqualityMap_;
+  // Lazily update `productEqualityMap_` since it takes a lot of time.
+  bool productEqualityMapUpdated_ = true;
 };
 
 }  // namespace disc_ral
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc b/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc
index bbc2c8e7f52..f8944341aba 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc
@@ -379,9 +379,9 @@ struct ShapeSimplifierPass
 
   void runOnOperation() override;
 
-  LogicalResult applyShapeAnalysis(ShapeAnalysis&, bool&);
+  LogicalResult applyShapeAnalysis(ShapeAnalysisDeprecated&, bool&);
 
-  LogicalResult applySymbolicShapeOptimization(ShapeAnalysis&, bool&);
+  LogicalResult applySymbolicShapeOptimization(ShapeAnalysisDeprecated&, bool&);
 };
 
 void ShapeSimplifierPass::populateShapeRefinerPatterns(
@@ -423,7 +423,7 @@ void ShapeSimplifierPass::runOnOperation() {
     (void)applyPatternsAndFoldGreedily(main, std::move(patterns));
 
     // Stage #2: propagate shape information globally
-    ShapeAnalysis analysis(main);
+    ShapeAnalysisDeprecated analysis(main);
     if (failed(analysis.run())) {
       // error message should be generated inside the above function call.
       signalPassFailure();
@@ -454,8 +454,8 @@ void ShapeSimplifierPass::runOnOperation() {
   }
 }
 
-LogicalResult ShapeSimplifierPass::applyShapeAnalysis(ShapeAnalysis& analysis,
-                                                      bool& changed) {
+LogicalResult ShapeSimplifierPass::applyShapeAnalysis(
+    ShapeAnalysisDeprecated& analysis, bool& changed) {
   func::FuncOp func = dyn_cast_or_null<func::FuncOp>(analysis.getOperation());
   if (func == nullptr) {
     return failure();
@@ -500,7 +500,7 @@ LogicalResult ShapeSimplifierPass::applyShapeAnalysis(ShapeAnalysis& analysis,
 }
 
 LogicalResult ShapeSimplifierPass::applySymbolicShapeOptimization(
-    ShapeAnalysis& analysis, bool& changed) {
+    ShapeAnalysisDeprecated& analysis, bool& changed) {
   func::FuncOp func = dyn_cast_or_null<func::FuncOp>(analysis.getOperation());
   if (func == nullptr) {
     return failure();
diff --git a/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc b/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc
index ca4170320c3..f67b0815f21 100644
--- a/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc
@@ -19,6 +19,7 @@
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/MLIRContext.h"  // TF:llvm-project
 #include "mlir/Pass/Pass.h"       // TF:local_config_mlir
+#include "tensorflow/compiler/mlir/disc/disc_util.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/codegen_utils.h"
 #include "tensorflow/compiler/mlir/disc/transforms/fusion_utils.h"
@@ -29,6 +30,7 @@ namespace mlir {
 namespace disc_ral {
 namespace {
 
+using disc_shape::SymbolicDimOp;
 using lmhlo::DynamicBroadcastInDimOp;
 using lmhlo::FusionOp;
 
@@ -91,10 +93,72 @@ Value createViewLike(OpBuilder& b, Location loc, Value from, Value to) {
   return CastMemRefTo(b, loc, from, targetType, toShape);
 }
 
+struct ShapeConstraintIRCloneContext {
+  BlockAndValueMapping valueMapping;
+  DenseMap<Value, SmallVector<SymbolicDimOp>> value2Symbols;
+  DenseMap<SymbolicDimOp, SymbolicDimOp> symbolMapping;
+  SymbolicDimMgr* mgr = nullptr;
+};
+
+FusionOp cloneFusion(OpBuilder& b, FusionOp op,
+                     ShapeConstraintIRCloneContext* ctx = nullptr) {
+  if (ctx == nullptr) {
+    return dyn_cast<FusionOp>(b.clone(*op.getOperation()));
+  }
+
+  SymbolTable& table = ctx->mgr->symbolTable();
+  DenseSet<SymbolicDimOp> originalSymbols;
+  // 1, create a local view of original (dynamic shape) buffers.
+  for (Operation& op : op.region().front()) {
+    for (Value operand : op.getOperands()) {
+      auto ty = operand.getType().dyn_cast<MemRefType>();
+      // skip staitc shape buffers.
+      if (!ty || ty.hasStaticShape()) continue;
+      if (ctx->valueMapping.lookupOrNull(operand)) continue;
+      auto symbols = getMemRefValueSymbolicDimRefs(operand);
+      if (!symbols) continue;
+      ctx->valueMapping.map(operand,
+                            createViewLike(b, op.getLoc(), operand, operand));
+      auto& symbolOps = ctx->value2Symbols[operand];
+      for (const auto& sym : *symbols) {
+        auto symOp = table.lookup<SymbolicDimOp>(sym.getValue());
+        assert(symOp);
+        originalSymbols.insert(symOp);
+        symbolOps.push_back(symOp);
+      }
+    }
+  }
+  // 2, clone symbol group
+  auto status = ctx->mgr->cloneSymbolGroup(originalSymbols, ctx->symbolMapping);
+  assert(!failed(status));
+
+  // 3, attach new symbol attrs for the cloned operands.
+  for (const auto& it : ctx->valueMapping.getValueMap()) {
+    auto symIt = ctx->value2Symbols.find(it.first);
+    assert(symIt != ctx->value2Symbols.end());
+    Operation* definingOp = it.second.getDefiningOp();
+    assert(definingOp != nullptr);
+    SmallVector<Attribute> newAttrs;
+    SmallVector<SymbolicDimOp> newSymbols;
+    for (SymbolicDimOp sym : symIt->second) {
+      SymbolicDimOp newSymbol = ctx->symbolMapping[sym];
+      newSymbols.push_back(newSymbol);
+      newAttrs.push_back(FlatSymbolRefAttr::get(newSymbol));
+    }
+    auto symbolicShapeAttr = ArrayAttr::get(op->getContext(), newAttrs);
+    StringRef attrName = disc_shape::SymbolicDimOp::getSymbolicDimAttrName();
+    definingOp->setAttr(attrName, symbolicShapeAttr);
+    ctx->value2Symbols[it.second] = std::move(newSymbols);
+  }
+
+  return dyn_cast<FusionOp>(b.clone(*op.getOperation(), ctx->valueMapping));
+}
+
 FusionOp cloneWithBroadcastSimplifying(
     OpBuilder& b, FusionOp fusion_op,
-    SmallVectorImpl<Operation*>& broadcast_ops) {
-  FusionOp cloned = dyn_cast<FusionOp>(b.clone(*fusion_op.getOperation()));
+    SmallVectorImpl<Operation*>& broadcast_ops,
+    ShapeConstraintIRCloneContext* ctx = nullptr) {
+  FusionOp cloned = cloneFusion(b, fusion_op, ctx);
 
   // Collects all candidate broadcast ops inside the fusion op.
   cloned.walk([&](DynamicBroadcastInDimOp op) {
@@ -113,6 +177,17 @@ FusionOp cloneWithBroadcastSimplifying(
         continue;
       }
       user->replaceUsesOfWith(result, operand);
+      if (ctx != nullptr) {
+        auto lhsIt = ctx->value2Symbols.find(result);
+        auto rhsIt = ctx->value2Symbols.find(operand);
+        if (lhsIt == ctx->value2Symbols.end() ||
+            rhsIt == ctx->value2Symbols.end())
+          continue;
+
+        assert(lhsIt->second.size() == rhsIt->second.size());
+        for (const auto& z : llvm::zip(lhsIt->second, rhsIt->second))
+          ctx->mgr->mapSymbolicDimEqual(std::get<0>(z), std::get<1>(z));
+      }
     }
   }
   return cloned;
@@ -176,9 +251,11 @@ struct DiscSpecializeFusionWithSpeculationPass
 
     // Clone the fusion op and mark all candidate broadcast ops within the
     // fusion op.
+    ShapeConstraintIRCloneContext ctx;
+    ctx.mgr = symbolMgr.get();
     SmallVector<Operation*, 4> broadcast_ops;
-    FusionOp cloned =
-        cloneWithBroadcastSimplifying(b, fusion_op, broadcast_ops);
+    FusionOp cloned = cloneWithBroadcastSimplifying(
+        b, fusion_op, broadcast_ops, useShapeConstraintIR() ? &ctx : nullptr);
     addFusionTag(b, cloned, "no_ib");
 
     // Generate the predition.
@@ -205,32 +282,87 @@ struct DiscSpecializeFusionWithSpeculationPass
     fusion_op.getOperation()->moveBefore(else_block, else_block->begin());
 
     DenseMap<Value, Value> viewMap;
-    SmallVector<Operation*, 4> op_list;
-    for (Operation& op : cloned.region().front()) op_list.push_back(&op);
-    OpListShapeAnalysis op_list_shape_analysis(op_list);
-    for (Operation* op : op_list) {
-      for (Value operand : op->getOperands()) {
-        if (viewMap.find(operand) != viewMap.end()) continue;
-        Value leader =
-            op_list_shape_analysis.GetLeaderValueWithSameShape(operand);
-        if (!leader || leader == operand) continue;
-        // TODO(disc): handle mismatch type case
-        auto leaderTy = leader.getType().dyn_cast<MemRefType>();
-        auto operandTy = operand.getType().dyn_cast<MemRefType>();
-        if (!leaderTy || !operandTy ||
-            leaderTy.getRank() != operandTy.getRank())
-          continue;
-        bool sameShape = true;
-        for (auto&& en : llvm::zip(leaderTy.getShape(), operandTy.getShape())) {
-          if (std::get<0>(en) != std::get<1>(en)) {
-            sameShape = false;
-            break;
+    if (useShapeConstraintIR()) {
+      OpBuilder viewBuilder(cloned);
+      // build symbolicDimOp to its corresponding SSA value map
+      DenseMap<SymbolicDimOp, Value> symbolicDim2SSAValue;
+      for (Operation& op : cloned.region().front()) {
+        for (Value operand : op.getOperands()) {
+          auto it = ctx.value2Symbols.find(operand);
+          if (it == ctx.value2Symbols.end()) continue;
+          for (const auto& en : llvm::enumerate(it->second)) {
+            SymbolicDimOp rootSym = ctx.mgr->getRootSymbolicDim(en.value());
+            if (en.value() != rootSym) continue;
+            if (rootSym.isDynamic()) {
+              symbolicDim2SSAValue[rootSym] = viewBuilder.create<memref::DimOp>(
+                  op.getLoc(), operand, en.index());
+            } else {
+              symbolicDim2SSAValue[rootSym] =
+                  viewBuilder.create<arith::ConstantIndexOp>(
+                      op.getLoc(), rootSym.getDimSize());
+            }
           }
         }
-        if (!sameShape) continue;
-        OpBuilder viewBuilder(cloned);
-        viewMap[operand] =
-            createViewLike(viewBuilder, op->getLoc(), operand, leader);
+      }
+
+      for (Operation& op : cloned.region().front()) {
+        for (Value operand : op.getOperands()) {
+          if (viewMap.find(operand) != viewMap.end()) continue;
+          auto it = ctx.value2Symbols.find(operand);
+          if (it == ctx.value2Symbols.end()) continue;
+          SmallVector<int64_t> newShape;
+          SmallVector<Value> newShapeValues;
+          SmallVector<Attribute> newAttrs;
+          for (SymbolicDimOp sym : it->second) {
+            SymbolicDimOp rootSym = ctx.mgr->getRootSymbolicDim(sym);
+            newAttrs.push_back(FlatSymbolRefAttr::get(rootSym));
+            newShapeValues.push_back(symbolicDim2SSAValue[rootSym]);
+            newShape.push_back(rootSym.isDynamic() ? ShapedType::kDynamicSize
+                                                   : rootSym.getDimSize());
+          }
+          auto oldType = operand.getType().cast<MemRefType>();
+          auto newType =
+              MemRefType::get(newShape, oldType.getElementType(),
+                              oldType.getLayout(), oldType.getMemorySpace());
+          // TODO(disc): handle mismatch type case
+          if (newType != oldType) continue;
+          viewMap[operand] = CastMemRefTo(viewBuilder, op.getLoc(), operand,
+                                          newType, newShapeValues);
+          auto symbolicShapeAttr = ArrayAttr::get(op.getContext(), newAttrs);
+          StringRef attrName = SymbolicDimOp::getSymbolicDimAttrName();
+          viewMap[operand].getDefiningOp()->setAttr(attrName,
+                                                    symbolicShapeAttr);
+        }
+      }
+    } else {
+      SmallVector<Operation*, 4> op_list;
+      for (Operation& op : cloned.region().front()) op_list.push_back(&op);
+      OpListShapeAnalysis op_list_shape_analysis(op_list);
+      for (Operation* op : op_list) {
+        for (Value operand : op->getOperands()) {
+          if (viewMap.find(operand) != viewMap.end()) continue;
+          Value leader =
+              op_list_shape_analysis.GetLeaderValueWithSameShape(operand);
+          if (!leader || leader == operand) continue;
+          // TODO(disc): handle mismatch type case
+          auto leaderTy = leader.getType().dyn_cast<MemRefType>();
+          auto operandTy = operand.getType().dyn_cast<MemRefType>();
+          if (!leaderTy || !operandTy ||
+              leaderTy.getRank() != operandTy.getRank())
+            continue;
+          bool sameShape = true;
+          for (auto&& en :
+               llvm::zip(leaderTy.getShape(), operandTy.getShape())) {
+            if (std::get<0>(en) != std::get<1>(en)) {
+              sameShape = false;
+              break;
+            }
+          }
+          if (!sameShape) continue;
+          OpBuilder viewBuilder(cloned);
+          viewMap[operand] =
+              createViewLike(viewBuilder, op->getLoc(), operand, leader);
+        }
       }
     }
 
@@ -556,6 +688,22 @@ struct DiscSpecializeFusionWithSpeculationPass
   }
 
   void runOnOperation() override {
+    if (useShapeConstraintIR()) {
+      func::FuncOp func = getOperation();
+      // skip shape constraint graph
+      if (func.getName() ==
+          SymbolicDimMgr::getShapeConstraintGraphFunctionName())
+        return;
+
+      auto m = func->getParentOfType<ModuleOp>();
+      symbolMgr.reset(new SymbolicDimMgr(m));
+      if (failed(symbolMgr->load())) {
+        getOperation()->emitError() << "fail to load shape constraint IR\n";
+        signalPassFailure();
+        return;
+      }
+    }
+
     // Stage #1: broadcast simplifier with speculation.
     Speculator(
         &DiscSpecializeFusionWithSpeculationPass::DoBroadcastSpeculation);
@@ -573,7 +721,17 @@ struct DiscSpecializeFusionWithSpeculationPass
     // Stage #4: speculation of vectorization/tiling.
     Speculator(
         &DiscSpecializeFusionWithSpeculationPass::DoVectorizeOrTileSpeculation);
+
+    if (useShapeConstraintIR()) {
+      if (failed(symbolMgr->save())) {
+        getOperation()->emitError() << "fail to load shape constraint IR\n";
+        signalPassFailure();
+        return;
+      }
+    }
   }
+
+  std::shared_ptr<SymbolicDimMgr> symbolMgr;
 };
 
 }  // namespace
diff --git a/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc b/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc
index f6da9cee529..11269253588 100644
--- a/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_stitch_fusion.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LogicalResult.h"
+#include "tensorflow/compiler/mlir/disc/disc_util.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/codegen_utils.h"
 #include "tensorflow/compiler/mlir/disc/transforms/fusion_utils.h"
@@ -58,13 +59,10 @@ bool isOnGpu(Operation* op) {
          memory_space.cast<StringAttr>().getValue() == placement_utils::kGpu;
 }
 
-LogicalResult HandleCpuFusionOp(OpBuilder& b, Operation* fusion) {
+LogicalResult HandleCpuFusionOp(OpBuilder& b, Operation* fusion,
+                                ShapeAnalysis& shapeAnalysis) {
   auto fusionOp = cast<lmhlo::FusionOp>(fusion);
   assert(fusionOp);
-  // TODO: use FuncOp that contains `fusionOp` to construct shape-analysis,
-  // which will use global information for shape equality and decomposition
-  // analysis.
-  ShapeAnalysis shapeAnalysis(fusionOp);
   FusionPattern fusionPattern(fusionOp, &shapeAnalysis);
   if (!fusionPattern.isStitchFusion()) {
     // skip non-stitch fusion pattern.
@@ -79,7 +77,7 @@ LogicalResult HandleCpuFusionOp(OpBuilder& b, Operation* fusion) {
     return success();
   }
 
-  StitchCPUAnalysis stitchAnalysis(fusionPattern);
+  StitchCPUAnalysis stitchAnalysis(fusionPattern, shapeAnalysis);
   if (!stitchAnalysis.doCodeGeneration(b, fusionOp)) {
     LLVM_DEBUG(llvm::dbgs() << "stitchAnalysis failed to doCodeGeneration\n");
     return failure();
@@ -101,6 +99,13 @@ struct DiscStitchFusion : public DiscStitchFusionBase<DiscStitchFusion> {
         cpu_fusion_worklist.push_back(op);
     });
 
+    std::unique_ptr<ShapeAnalysis> shapeAnalysisPtr;
+    if (!gpu_fusion_worklist.empty() || !cpu_fusion_worklist.empty()) {
+      if (useShapeConstraintIR()) {
+        shapeAnalysisPtr.reset(new ShapeConstraintIRAnalysis(func));
+      }
+    }
+
     for (Operation* fusion : gpu_fusion_worklist) {
       // TODO(disc): handling stitch fusion on GPU.
       signalPassFailure();
@@ -108,8 +113,15 @@ struct DiscStitchFusion : public DiscStitchFusionBase<DiscStitchFusion> {
     }
 
     for (Operation* fusion : cpu_fusion_worklist) {
+      if (!useShapeConstraintIR()) {
+        // TODO: use FuncOp that contains `fusionOp` to construct
+        // shape-analysis, which will use global information for shape equality
+        // and decomposition analysis.
+        shapeAnalysisPtr.reset(new ShapeAnalysisDeprecated{fusion});
+      }
+
       // Error message should be emitted inside the function.
-      if (failed(HandleCpuFusionOp(b, fusion))) {
+      if (failed(HandleCpuFusionOp(b, fusion, *shapeAnalysisPtr))) {
         signalPassFailure();
         return;
       }
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
index 12743ed9c51..8d49e51910f 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -39,6 +39,7 @@ namespace mlir {
 namespace disc_ral {
 
 using namespace lmhlo;
+using disc_shape::SymbolicDimOp;
 using placement_utils::kDiscPlaceAssignment;
 using placement_utils::kDiscShapeCalcAttr;
 
@@ -1089,7 +1090,7 @@ bool BaseFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
         Value shape = getEffectiveShape(target, result);
         return checkSameShape(lhs, rhs, target)
                    ? shapeAnalysis.isShapeEqual(ref_shape, shape)
-                   : shapeAnalysis.HasSameNumElements(ref_shape, shape);
+                   : shapeAnalysis.isSameNumElements(ref_shape, shape);
       })) {
     return false;
   }
@@ -1335,7 +1336,7 @@ bool StitchCpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
     return false;
   }
 
-  StitchCPUAnalysis stitchAnalysis(target);
+  StitchCPUAnalysis stitchAnalysis(target, shapeAnalysis);
   if (!stitchAnalysis.fusibilityAnalysis()) {
     LLVM_DEBUG(llvm::dbgs() << "fusibilityAnalysis failed\n");
     return false;
@@ -2449,9 +2450,15 @@ bool StitchCPUAnalysis::emitInOutTiles(OpBuilder& b, Location loc,
                 fusionPattern_.getOperands().end());
   inOuts.insert(inOuts.end(), fusionPattern_.getResults().begin(),
                 fusionPattern_.getResults().end());
+  SymbolicDimMgr* mgr = nullptr;
+  if (auto shapeIRAnalysis =
+          dynamic_cast<ShapeConstraintIRAnalysis*>(&shapeAnalysis_)) {
+    mgr = &shapeIRAnalysis->symbolicDimMgr();
+  }
   for (Value in : inOuts) {
     auto& valueViewStore = viewStore[in];
     auto& tileInfo = tilePlan_[in];
+    auto inSymbols = getMemRefValueSymbolicDimRefs(in);
     for (int id : parallelPlan_[in]) {
       auto& parallelInfo = parallelInfoStore_[id];
       auto& indices = parallelInfo.symbolIndices;
@@ -2472,9 +2479,11 @@ bool StitchCPUAnalysis::emitInOutTiles(OpBuilder& b, Location loc,
       // logic can be simplified to:
       //   %inputView = input[i][j][0:n][0:m] : memref<1x1x?x?xf32>
 
+      bool staticShape = true;
       SmallVector<OpFoldResult> subViewSizes;
       SmallVector<OpFoldResult> subViewOffsets;
       SmallVector<OpFoldResult> subViewSteps;
+      SmallVector<SymbolicDimOp> subViewSymbols;
       int parallelIdx = 0;
       for (int d = 0; d < rank; ++d) {
         subViewSteps.push_back(getIntAttr(1));
@@ -2484,6 +2493,10 @@ bool StitchCPUAnalysis::emitInOutTiles(OpBuilder& b, Location loc,
           auto& parallelIndex = parallelIndexStore_[parallelIt->second];
           subViewSizes.push_back(getIntAttr(parallelIndex.step));
           subViewOffsets.push_back(parallelInfo.symbolIndices[parallelIdx]);
+          if (inSymbols && mgr) {
+            subViewSymbols.push_back(
+                mgr->newConstantSymbolicDim(parallelIndex.step));
+          }
           ++parallelIdx;
         } else if (tileIt != tileInfo.tileSizes.end()) {
           subViewOffsets.push_back(getIntAttr(0));
@@ -2491,19 +2504,37 @@ bool StitchCPUAnalysis::emitInOutTiles(OpBuilder& b, Location loc,
             if (inType.getShape()[d] == ShapedType::kDynamicSize) {
               Value dimSize = b.create<memref::DimOp>(loc, in, d);
               subViewSizes.push_back(dimSize);
+              if (inSymbols && mgr) {
+                subViewSymbols.push_back(
+                    mgr->symbolTable().lookup<SymbolicDimOp>(
+                        (*inSymbols)[d].getValue()));
+                staticShape = false;
+              }
             } else {
               subViewSizes.push_back(getIntAttr(inType.getShape()[d]));
+              if (inSymbols && mgr) {
+                subViewSymbols.push_back(
+                    mgr->newConstantSymbolicDim(inType.getShape()[d]));
+              }
             }
           } else {
             subViewSizes.push_back(getIntAttr(tileIt->second));
+            if (inSymbols && mgr) {
+              subViewSymbols.push_back(
+                  mgr->newConstantSymbolicDim(tileIt->second));
+            }
           }
         } else {
           assert(false && "unexpected situation\n");
           return false;
         }
       }
-      valueViewStore[indices] = b.create<memref::SubViewOp>(
-          loc, in, subViewOffsets, subViewSizes, subViewSteps);
+      auto view = b.create<memref::SubViewOp>(loc, in, subViewOffsets,
+                                              subViewSizes, subViewSteps);
+      if (inSymbols && mgr && !staticShape) {
+        attachSymbolicDimOpRefArrayAttrOnOperation(view, subViewSymbols);
+      }
+      valueViewStore[indices] = view.getResult();
     }
   }
   return true;
@@ -2527,20 +2558,43 @@ Value StitchCPUAnalysis::emitTileBuffer(OpBuilder& b, Location loc, Value val) {
     return totalSize < 64;
   }();
 
+  bool staticShape = true;
+  SymbolicDimMgr* mgr = nullptr;
+  if (auto shapeIRAnalysis =
+          dynamic_cast<ShapeConstraintIRAnalysis*>(&shapeAnalysis_)) {
+    mgr = &shapeIRAnalysis->symbolicDimMgr();
+  }
+  auto valSymbols = getMemRefValueSymbolicDimRefs(val);
   SmallVector<int64_t> newTyShape;
   SmallVector<Value> dynDims;
+  SmallVector<SymbolicDimOp> tileSymbols;
   for (int d = 0; d < ty.getRank(); ++d) {
     auto it = tileInfo.tileSizes.find(d);
     if (it == tileInfo.tileSizes.end()) {
       newTyShape.push_back(1);
+      if (mgr && valSymbols) {
+        tileSymbols.push_back(mgr->newConstantSymbolicDim(1));
+      }
       continue;
     }
     if (it->second != ShapedType::kDynamicSize) {
       newTyShape.push_back(it->second);
+      if (mgr && valSymbols) {
+        tileSymbols.push_back(mgr->newConstantSymbolicDim(it->second));
+      }
       continue;
     }
     if (ty.getShape()[d] == ShapedType::kDynamicSize) {
       dynDims.push_back(b.create<memref::DimOp>(loc, val, d));
+      if (mgr && valSymbols) {
+        tileSymbols.push_back(mgr->symbolTable().lookup<SymbolicDimOp>(
+            (*valSymbols)[d].getValue()));
+        staticShape = false;
+      }
+    } else {
+      if (mgr && valSymbols) {
+        tileSymbols.push_back(mgr->newConstantSymbolicDim(ty.getShape()[d]));
+      }
     }
     newTyShape.push_back(ty.getShape()[d]);
   }
@@ -2557,6 +2611,10 @@ Value StitchCPUAnalysis::emitTileBuffer(OpBuilder& b, Location loc, Value val) {
     } else {
       tileBuffer = b.create<memref::AllocOp>(loc, newTy, dynDims);
     }
+    if (mgr && valSymbols && !staticShape) {
+      attachSymbolicDimOpRefArrayAttrOnOperation(tileBuffer.getDefiningOp(),
+                                                 tileSymbols);
+    }
   }
 
   return tileBuffer;
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.h b/tao_compiler/mlir/disc/transforms/fusion_utils.h
index 5c79ad4e32f..c4eb293c268 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.h
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.h
@@ -496,8 +496,9 @@ struct ParallelInfo {
 
 class StitchCPUAnalysis {
  public:
-  explicit StitchCPUAnalysis(FusionPattern& fusionPattern)
-      : fusionPattern_(fusionPattern) {}
+  explicit StitchCPUAnalysis(FusionPattern& fusionPattern,
+                             ShapeAnalysis& shapeAnalysis)
+      : fusionPattern_(fusionPattern), shapeAnalysis_(shapeAnalysis) {}
 
   // Returns true if the fusion pattern is a valid stitch pattern.
   bool fusibilityAnalysis();
@@ -597,6 +598,7 @@ class StitchCPUAnalysis {
 
  private:
   FusionPattern& fusionPattern_;
+  ShapeAnalysis& shapeAnalysis_;
 
   // Used for roots dominant analysis
   Value dominantValue_;
@@ -648,10 +650,6 @@ class StitchGpuFusionStrategy : public FusionStrategy {
  private:
   virtual Value getEffectiveShape(FusionPattern& target, Value value);
 
-  bool tileInfoPropagateI2O(
-      ShapeAnalysis& shapeAnalysis, DenseMap<Value, TileInfo>& tile_plan,
-      Operation* op, int64_t input_index,
-      SmallVector<std::pair<Value, TileInfo>, 4>& out_info);
   bool tileCoverInfoPropagateO2I(
       ShapeAnalysis& shapeAnalysis, DenseMap<Value, TileInfo>& tile_plan,
       Operation* op, SmallVector<std::pair<Value, TileInfo>, 4>& in_info,
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc b/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc
index 0c770cd1768..1340547035b 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc
@@ -57,233 +57,6 @@ Value StitchGpuFusionStrategy::getEffectiveShape(FusionPattern& target,
   return isa<lmhlo::ReduceOp>(result_op) ? result_op->getOperand(0) : v;
 }
 
-bool StitchGpuFusionStrategy::tileInfoPropagateI2O(
-    ShapeAnalysis& shapeAnalysis, DenseMap<Value, TileInfo>& tile_plan,
-    Operation* op, int64_t input_index,
-    SmallVector<std::pair<Value, TileInfo>, 4>& out_info) {
-  if (isa<lmhlo::ConstOp>(op)) {
-    return true;
-  }
-  if (isElementWise(op) ||
-      isa<lmhlo::RealDynamicSliceOp, lmhlo::SliceOp, lmhlo::ConcatenateOp>(
-          op)) {
-    Value in_value = op->getOperand(input_index);
-    Value out_value = cast<lmhlo::LmhloOp>(op).getResultBuffer();
-    auto tile_info = tile_plan[in_value];
-    out_info.emplace_back(out_value, tile_info);
-  } else if (isa<lmhlo::BroadcastInDimOp, lmhlo::DynamicBroadcastInDimOp>(op)) {
-    if (input_index != 0) {
-      return true;
-    }
-    Value in_value = op->getOperand(0);
-    Value out_value = cast<lmhlo::LmhloOp>(op).getResultBuffer();
-    auto& in_tile = tile_plan[in_value];
-    TileInfo out_tile;
-    auto dimAttr = op->getAttrOfType<DenseElementsAttr>("broadcast_dimensions");
-    assert(dimAttr);
-    auto dimensions = dimAttr.getValues<int64_t>();
-    DenseSet<int64_t> dim_set(dimensions.begin(), dimensions.end());
-    // Tiled dims should be minor dims in `broadcast_dimensions`.
-    // Broadcasted dims after tiled dims are all tiled.
-    DenseMap<int64_t, int64_t> broadcast_dim_o2i;
-    for (auto en : llvm::enumerate(dimensions)) {
-      broadcast_dim_o2i[en.value()] = en.index();
-    }
-    int64_t rank = out_value.getType().cast<MemRefType>().getRank();
-    bool tile_start = false;
-    for (int64_t i = 0; i < rank; i++) {
-      auto may_indim = broadcast_dim_o2i.find(i);
-      if (!tile_start) {
-        tile_start = may_indim != broadcast_dim_o2i.end() &&
-                     in_tile.tileSizes.find(may_indim->second) !=
-                         in_tile.tileSizes.end();
-      }
-      if (tile_start) {
-        if (may_indim != broadcast_dim_o2i.end() &&
-            in_tile.tileSizes.find(may_indim->second) ==
-                in_tile.tileSizes.end()) {
-          // Non-tiled dim becomes minor than tiled dim.
-          LLVM_DEBUG(llvm::dbgs() << "forward propagation failed: " << *op);
-          return false;
-        } else {
-          out_tile.tileSizes[i] = ShapedType::kDynamicSize;
-        }
-      }
-    }
-    out_info.emplace_back(out_value, out_tile);
-  } else if (isa<lmhlo::BroadcastOp>(op)) {
-    if (input_index != 0) {
-      return true;
-    }
-    Value in_value = op->getOperand(0);
-    Value out_value = cast<lmhlo::LmhloOp>(op).getResultBuffer();
-    auto& in_tile = tile_plan[in_value];
-    TileInfo out_tile;
-    auto sizesAttr = op->getAttrOfType<DenseElementsAttr>("broadcast_sizes");
-    assert(sizesAttr);
-    auto sizes = sizesAttr.getValues<int64_t>();
-    int64_t length = sizes.end() - sizes.begin();
-    for (auto in_tile_pair : in_tile.tileSizes) {
-      out_tile.tileSizes[in_tile_pair.first + length] = in_tile_pair.second;
-    }
-    out_info.emplace_back(out_value, out_tile);
-  } else if (isa<lmhlo::DynamicReshapeOp>(op)) {
-    if (input_index != 0) {
-      return true;
-    }
-    SmallVector<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>> dim_eq;
-    Value in_value = op->getOperand(0);
-    Value out_value = cast<lmhlo::LmhloOp>(op).getResultBuffer();
-    int64_t in_rank = in_value.getType().cast<MemRefType>().getRank();
-    int64_t out_rank = out_value.getType().cast<MemRefType>().getRank();
-
-    if (!shapeAnalysis.extractContinuousDimEqualInfo(in_value, out_value,
-                                                     dim_eq)) {
-      return false;
-    }
-
-    auto& in_tile = tile_plan[in_value];
-    TileInfo out_tile;
-    int64_t in_tiles_left = in_tile.tileSizes.size();
-    int64_t in_non_tiles_left = in_rank - in_tiles_left;
-    DenseSet<int64_t> out_ranks_mapped;
-    for (auto equal : dim_eq) {
-      SmallVector<int64_t> ins = equal.first;
-      SmallVector<int64_t> outs = equal.second;
-      out_ranks_mapped.insert(outs.begin(), outs.end());
-      bool is_tiled = (in_tile.tileSizes.count(ins[0]) != 0);
-      for (int64_t i = 1; i < ins.size(); i++) {
-        if ((in_tile.tileSizes.count(ins[i]) != 0) != is_tiled) {
-          // Delinearized dims are not consistent about tiling, meaning some
-          // are tiled and some other are not tiled. This breaks propagation.
-          LLVM_DEBUG(llvm::dbgs() << "forward propagation failed: " << *op);
-          return false;
-        }
-      }
-      if (is_tiled) {
-        for (auto out : outs) {
-          out_tile.tileSizes[out] = ShapedType::kDynamicSize;
-        }
-        in_tiles_left -= ins.size();
-      } else {
-        in_non_tiles_left -= ins.size();
-      }
-    }
-
-    if (in_non_tiles_left == 0) {
-      if (in_tiles_left != 0) {
-        // Sometimes, we failed to map all equal relationships. Since we know
-        // all non-tiled dims in equal-map are processed, all dims not in
-        // equal-map are tiled.
-        for (int64_t i = 0; i < out_rank; i++) {
-          if (!out_ranks_mapped.contains(i)) {
-            out_tile.tileSizes[i] = ShapedType::kDynamicSize;
-          }
-        }
-      }
-    } else if (in_tiles_left != 0) {
-      // This means some dims not in i2o-map are tiled. But we cannot figure
-      // out who are they...
-      LLVM_DEBUG(llvm::dbgs() << "forward propagation failed: " << *op);
-      return false;
-    }
-    out_info.emplace_back(out_value, out_tile);
-  } else if (isa<lmhlo::DynamicGatherOp, lmhlo::GatherOp>(op)) {
-    Value in_value = op->getOperand(0);
-    Value out_value = cast<lmhlo::LmhloOp>(op).getResultBuffer();
-    int64_t in_rank = in_value.getType().cast<MemRefType>().getRank();
-    int64_t out_rank = out_value.getType().cast<MemRefType>().getRank();
-    auto gather = dyn_cast<lmhlo::GatherOp>(op);
-    auto d_gather = dyn_cast<lmhlo::DynamicGatherOp>(op);
-    auto dimension_numbers =
-        gather ? gather.dimension_numbers() : d_gather.dimension_numbers();
-    auto collapsed_slice_dims = dimension_numbers.getCollapsedSliceDims();
-    auto offset_dims = dimension_numbers.getOffsetDims();
-
-    SmallVector<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>> dim_eq;
-    // Non-collapsed dims.
-    // Map index of offset_dims to operand.
-    SmallVector<int64_t, 4> remapped_offset_dims;
-    DenseSet<int64_t> collapsed_set(collapsed_slice_dims.begin(),
-                                    collapsed_slice_dims.end());
-    for (int64_t i = 0; i < in_rank; i++) {
-      if (collapsed_set.contains(i)) {
-        continue;
-      }
-      remapped_offset_dims.push_back(i);
-    }
-    for (auto offset : llvm::enumerate(offset_dims)) {
-      dim_eq.emplace_back(
-          SmallVector<int64_t>({remapped_offset_dims[offset.index()]}),
-          SmallVector<int64_t>({offset.value()}));
-    }
-    // Collapsed dims.
-    SmallVector<int64_t> batch_dims;
-    DenseSet<int64_t> offset_dim_set(offset_dims.begin(), offset_dims.end());
-    for (int64_t i = 0; i < out_rank; ++i) {
-      if (!offset_dim_set.contains(i)) {
-        batch_dims.push_back(i);
-      }
-    }
-    SmallVector<int64_t> collapsed_slice_dims_vec(collapsed_slice_dims.begin(),
-                                                  collapsed_slice_dims.end());
-    dim_eq.emplace_back(collapsed_slice_dims_vec, batch_dims);
-
-    // TODO: the logic is the same with reshape. Rewrite as a function.
-    auto& in_tile = tile_plan[in_value];
-    TileInfo out_tile;
-    int64_t in_tiles_left = in_tile.tileSizes.size();
-    int64_t in_non_tiles_left = in_rank - in_tiles_left;
-    DenseSet<int64_t> out_ranks_mapped;
-    for (auto equal : dim_eq) {
-      SmallVector<int64_t> ins = equal.first;
-      SmallVector<int64_t> outs = equal.second;
-      out_ranks_mapped.insert(outs.begin(), outs.end());
-      bool is_tiled = (in_tile.tileSizes.count(ins[0]) != 0);
-      for (int64_t i = 1; i < ins.size(); i++) {
-        if ((in_tile.tileSizes.count(ins[i]) != 0) != is_tiled) {
-          // The grouped dims are not consistent about tiling, meaning some
-          // are tiled and some other are not tiled. This breaks propagation.
-          LLVM_DEBUG(llvm::dbgs() << "forward propagation failed: " << *op);
-          return false;
-        }
-      }
-      if (is_tiled) {
-        for (auto out : outs) {
-          out_tile.tileSizes[out] = ShapedType::kDynamicSize;
-        }
-        in_tiles_left -= ins.size();
-      } else {
-        in_non_tiles_left -= ins.size();
-      }
-    }
-
-    // TODO: gather op do not need the following logic?
-    if (in_non_tiles_left == 0) {
-      if (in_tiles_left != 0) {
-        // Sometimes, we failed to map all equal relationships. Since we know
-        // all non-tiled dims in equal-map are processed, all dims not in
-        // equal-map are tiled.
-        for (int64_t i = 0; i < out_rank; i++) {
-          if (!out_ranks_mapped.contains(i)) {
-            out_tile.tileSizes[i] = ShapedType::kDynamicSize;
-          }
-        }
-      }
-    } else if (in_tiles_left != 0) {
-      // This means some dims not in i2o-map are tiled. But we cannot figure
-      // out who are they...
-      LLVM_DEBUG(llvm::dbgs() << "forward propagation failed: " << *op);
-      return false;
-    }
-    out_info.emplace_back(out_value, out_tile);
-  } else {
-    LLVM_DEBUG(llvm::dbgs() << "forward propagation failed: " << *op);
-    return false;
-  }
-  return true;
-}
-
 bool StitchGpuFusionStrategy::tileCoverInfoPropagateO2I(
     ShapeAnalysis& shapeAnalysis, DenseMap<Value, TileInfo>& tile_plan,
     Operation* op, SmallVector<std::pair<Value, TileInfo>, 4>& in_info,
@@ -347,57 +120,112 @@ bool StitchGpuFusionStrategy::tileCoverInfoPropagateO2I(
     Value out_value = cast<lmhlo::LmhloOp>(op).getResultBuffer();
     int64_t in_rank = in_value.getType().cast<MemRefType>().getRank();
     int64_t out_rank = out_value.getType().cast<MemRefType>().getRank();
-    if (!shapeAnalysis.extractContinuousDimEqualInfo(in_value, out_value,
-                                                     dim_eq)) {
-      return false;
-    }
 
-    auto& out_tile = tile_plan[out_value];
-    TileInfo in_tile;
+    if (auto analysis_deprecated =
+            dynamic_cast<ShapeAnalysisDeprecated*>(&shapeAnalysis)) {
+      auto& shapeAnalysis = *analysis_deprecated;
+      if (!shapeAnalysis.extractContinuousDimEqualInfo(in_value, out_value,
+                                                       dim_eq)) {
+        return false;
+      }
 
-    int64_t out_tiles_left = out_tile.tileSizes.size();
-    int64_t out_non_tiles_left = out_rank - out_tiles_left;
-    DenseSet<int64_t> in_ranks_mapped;
-    for (auto equal : dim_eq) {
-      SmallVector<int64_t> ins = equal.first;
-      SmallVector<int64_t> outs = equal.second;
-      in_ranks_mapped.insert(ins.begin(), ins.end());
-      bool is_tiled = (out_tile.tileSizes.count(outs[0]) != 0);
-      // Make sure tile-info of all out dims are consistent.
-      for (int64_t i = 1; i < outs.size(); i++) {
-        if ((out_tile.tileSizes.count(outs[i]) != 0) != is_tiled) {
-          LLVM_DEBUG(llvm::dbgs() << "tile backprop failed: " << *op);
-          return false;
+      auto& out_tile = tile_plan[out_value];
+      TileInfo in_tile;
+
+      int64_t out_tiles_left = out_tile.tileSizes.size();
+      int64_t out_non_tiles_left = out_rank - out_tiles_left;
+      DenseSet<int64_t> in_ranks_mapped;
+      for (auto equal : dim_eq) {
+        SmallVector<int64_t> ins = equal.first;
+        SmallVector<int64_t> outs = equal.second;
+        in_ranks_mapped.insert(ins.begin(), ins.end());
+        bool is_tiled = (out_tile.tileSizes.count(outs[0]) != 0);
+        // Make sure tile-info of all out dims are consistent.
+        for (int64_t i = 1; i < outs.size(); i++) {
+          if ((out_tile.tileSizes.count(outs[i]) != 0) != is_tiled) {
+            LLVM_DEBUG(llvm::dbgs() << "tile backprop failed: " << *op);
+            return false;
+          }
+        }
+        if (is_tiled) {
+          for (auto in : ins) {
+            in_tile.tileSizes[in] = ShapedType::kDynamicSize;
+          }
+          out_tiles_left -= outs.size();
+        } else {
+          out_non_tiles_left -= outs.size();
         }
       }
-      if (is_tiled) {
-        for (auto in : ins) {
-          in_tile.tileSizes[in] = ShapedType::kDynamicSize;
+
+      if (out_non_tiles_left == 0) {
+        if (out_tiles_left != 0) {
+          // Sometimes, we failed to map all equal relationships. Since we know
+          // all non-tiled dims in equal-map are processed when reach here, all
+          // dims not in equal-map are tiled.
+          for (int64_t i = 0; i < in_rank; i++) {
+            if (!in_ranks_mapped.contains(i)) {
+              in_tile.tileSizes[i] = ShapedType::kDynamicSize;
+            }
+          }
+        }
+      } else if (out_tiles_left != 0) {
+        // This means some dims not in o2i-map are tiled. But we cannot figure
+        // out who are they...
+        LLVM_DEBUG(llvm::dbgs() << "tile backprop failed: " << *op);
+        return false;
+      }
+      in_info.emplace_back(in_value, in_tile);
+    } else {
+      bool last_is_tiled = false;
+      auto& out_tile = tile_plan[out_value];
+      // Partition output dims according to tile information to make sure
+      // each group is continuous and is either tiled or non-tiled.
+      SmallVector<int> out_group_end_dims;
+      auto outShape = out_value.getType().cast<MemRefType>().getShape();
+      for (int i = 0; i < out_rank; ++i) {
+        if (out_tile.tileSizes.count(i) == 0) {
+          if (!last_is_tiled) continue;
+          // ignore size-1 dim between tiled dims.
+          if (outShape[i] == 1) continue;
+          last_is_tiled = false;
+          out_group_end_dims.push_back(i);
+        } else {
+          if (last_is_tiled) continue;
+          last_is_tiled = true;
+          if (i != 0) out_group_end_dims.push_back(i);
         }
-        out_tiles_left -= outs.size();
-      } else {
-        out_non_tiles_left -= outs.size();
       }
-    }
 
-    if (out_non_tiles_left == 0) {
-      if (out_tiles_left != 0) {
-        // Sometimes, we failed to map all equal relationships. Since we know
-        // all non-tiled dims in equal-map are processed when reach here, all
-        // dims not in equal-map are tiled.
-        for (int64_t i = 0; i < in_rank; i++) {
-          if (!in_ranks_mapped.contains(i)) {
-            in_tile.tileSizes[i] = ShapedType::kDynamicSize;
+      TileInfo in_tile;
+      int last_map_input_dim = 0;
+      int last_map_output_dim = 0;
+      for (int out_group_end_dim : out_group_end_dims) {
+        bool matched = false;
+        for (int i = last_map_input_dim; i <= in_rank; ++i) {
+          if (shapeAnalysis.isProductEqual(out_value, last_map_output_dim,
+                                           out_group_end_dim, in_value,
+                                           last_map_input_dim, i)) {
+            last_map_input_dim = i;
+            last_map_output_dim = out_group_end_dim;
+            if (out_tile.tileSizes.count(last_map_output_dim) != 0) {
+              for (int inIdx = last_map_input_dim; inIdx < i; ++inIdx) {
+                in_tile.tileSizes[inIdx] = ShapedType::kDynamicSize;
+              }
+            }
+            matched = true;
+            break;
           }
         }
+        if (!matched) return false;
       }
-    } else if (out_tiles_left != 0) {
-      // This means some dims not in o2i-map are tiled. But we cannot figure
-      // out who are they...
-      LLVM_DEBUG(llvm::dbgs() << "tile backprop failed: " << *op);
-      return false;
+
+      if (last_is_tiled) {
+        for (int inIdx = last_map_input_dim; inIdx < in_rank; ++inIdx) {
+          in_tile.tileSizes[inIdx] = ShapedType::kDynamicSize;
+        }
+      }
+      in_info.emplace_back(in_value, in_tile);
     }
-    in_info.emplace_back(in_value, in_tile);
   } else if (isa<lmhlo::TransposeOp>(op)) {
     auto transpose = cast<lmhlo::TransposeOp>(op);
     auto permutation = transpose.permutation().getValues<int64_t>();
@@ -562,7 +390,7 @@ bool StitchGpuFusionStrategy::findFusionPatternTypeAndSubroot(
             Value shape = getEffectiveShape(fusion_pattern, result);
             return isRank2RowReduction(op)
                        ? shapeAnalysis.isShapeEqual(ref_shape, shape)
-                       : shapeAnalysis.HasSameNumElements(ref_shape, shape);
+                       : shapeAnalysis.isSameNumElements(ref_shape, shape);
           })) {
         fusion_type = FusionType::kStitch;
       }
@@ -579,7 +407,7 @@ bool StitchGpuFusionStrategy::findFusionPatternTypeAndSubroot(
             Value shape = getEffectiveShape(fusion_pattern, result);
             return isRank2ColReduction(op)
                        ? shapeAnalysis.isShapeEqual(ref_shape, shape)
-                       : shapeAnalysis.HasSameNumElements(ref_shape, shape);
+                       : shapeAnalysis.isSameNumElements(ref_shape, shape);
           })) {
         return false;
       }
@@ -606,7 +434,7 @@ bool StitchGpuFusionStrategy::findFusionPatternTypeAndSubroot(
         Value ref_shape = getEffectiveShape(fusion_pattern, results[0]);
         if (!llvm::all_of(results, [&](Value result) {
               Value shape = getEffectiveShape(fusion_pattern, result);
-              return shapeAnalysis.HasSameNumElements(ref_shape, shape);
+              return shapeAnalysis.isSameNumElements(ref_shape, shape);
             })) {
           return false;
         }
@@ -674,66 +502,110 @@ bool StitchGpuFusionStrategy::tileXroots(ShapeAnalysis& shapeAnalysis,
   const auto& dominant = fusion_pattern.getDominantOp()->getOperand(0);
   auto& dominant_tile = tile_plan[dominant];
   const auto& results = fusion_pattern.getResults();
+
+  SmallVector<int> dominant_non_tiled_dims;
+  for (int i = 0; i < dominant.getType().cast<MemRefType>().getRank(); ++i) {
+    if (dominant_tile.tileSizes.count(i) == 0) {
+      dominant_non_tiled_dims.push_back(i);
+    } else {
+      // Suppose that non-tiled dims are the continuous most-major dims.
+      // TODO(disc): change this once we support sub-roots other than
+      // row-reduction.
+      break;
+    }
+  }
   for (auto res : results) {
     bool irregular = false;
     Operation* op = fusion_pattern.findLastWriter(res);
     if (subroots_set.contains(op)) {
       continue;
     }
-    // Find equal dims between res and a sub-root.
-    SmallVector<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>> equal;
-    shapeAnalysis.extractContinuousDimEqualInfo(dominant, res, equal);
-    int64_t dominant_non_tiled_dim_matched_n = 0;
-    DenseSet<int64_t> res_non_tiled_dims;
-    for (auto eq_dim : equal) {
-      auto lhs = eq_dim.first;
-      bool non_tiled = (dominant_tile.tileSizes.count(lhs[0]) == 0);
-      for (int64_t i = 1; i < lhs.size(); i++) {
-        if ((dominant_tile.tileSizes.count(lhs[i]) == 0) != non_tiled) {
-          // Both tiled and non-tiled dims are mapped to the same dim of res,
-          // thus cannot determine the tile plan for res;
-          irregular_xroots.insert(op);
-          irregular = true;
+    if (auto analysis_deprecated =
+            dynamic_cast<ShapeAnalysisDeprecated*>(&shapeAnalysis)) {
+      auto& shapeAnalysis = *analysis_deprecated;
+      // Find equal dims between res and a sub-root.
+      SmallVector<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>> equal;
+      shapeAnalysis.extractContinuousDimEqualInfo(dominant, res, equal);
+      int64_t dominant_non_tiled_dim_matched_n = 0;
+      DenseSet<int64_t> res_non_tiled_dims;
+      for (auto eq_dim : equal) {
+        auto lhs = eq_dim.first;
+        bool non_tiled = (dominant_tile.tileSizes.count(lhs[0]) == 0);
+        for (int64_t i = 1; i < lhs.size(); i++) {
+          if ((dominant_tile.tileSizes.count(lhs[i]) == 0) != non_tiled) {
+            // Both tiled and non-tiled dims are mapped to the same dim of res,
+            // thus cannot determine the tile plan for res;
+            irregular_xroots.insert(op);
+            irregular = true;
+            break;
+          }
+        }
+        if (irregular) {
           break;
         }
+        if (non_tiled) {
+          res_non_tiled_dims.insert(eq_dim.second.begin(), eq_dim.second.end());
+          dominant_non_tiled_dim_matched_n += lhs.size();
+        }
       }
       if (irregular) {
-        break;
+        continue;
+      } else if (dominant_non_tiled_dim_matched_n != 1) {
+        // Check constraint: the element-number of non-tiled dims are the same
+        // between sub-roots and regular xroots. If not every non-tiled dim is
+        // matched between sub-root and current op, it is an irregular xroot.
+        // Note we know that sub-root is reduction and has only 1 non-tiled dim.
+        irregular_xroots.insert(op);
+        continue;
+      } else if (*std::max_element(res_non_tiled_dims.begin(),
+                                   res_non_tiled_dims.end()) !=
+                 res_non_tiled_dims.size() - 1) {
+        // Check constraint: the tiled dims are all minor dims for all regular
+        // xroot ops. This means the non-tiled dims should all be the major
+        // dims.
+        irregular_xroots.insert(op);
+        continue;
+      }
+
+      auto& plan = tile_plan[res];
+      int64_t rank = res.getType().cast<MemRefType>().getRank();
+      for (int64_t i = 0; i < rank; i++) {
+        if (res_non_tiled_dims.count(i) == 0) {
+          // Note that we only log whether a dimension is tiled or not. We do
+          // not care about the tiling size. Thus we set all tiled dims with
+          // with the value `kDynamicSize`.
+          plan.tileSizes[i] = ShapedType::kDynamicSize;
+        }
       }
-      if (non_tiled) {
-        res_non_tiled_dims.insert(eq_dim.second.begin(), eq_dim.second.end());
-        dominant_non_tiled_dim_matched_n += lhs.size();
+      regular_xroots.insert(op);
+    } else {
+      int last_non_tied_dim = -1;
+      SmallVector<int> res_non_tiled_dims;
+      int result_rank = res.getType().cast<MemRefType>().getRank();
+      do {
+        if (shapeAnalysis.isProductEqual(dominant, dominant_non_tiled_dims, res,
+                                         res_non_tiled_dims)) {
+          break;
+        }
+        ++last_non_tied_dim;
+        res_non_tiled_dims.push_back(last_non_tied_dim);
+      } while (last_non_tied_dim < result_rank);
+
+      if (last_non_tied_dim >= result_rank) {
+        irregular = true;
+        irregular_xroots.insert(op);
+        continue;
       }
-    }
-    if (irregular) {
-      continue;
-    } else if (dominant_non_tiled_dim_matched_n != 1) {
-      // Check constraint: the element-number of non-tiled dims are the same
-      // between sub-roots and regular xroots. If not every non-tiled dim is
-      // matched between sub-root and current op, it is an irregular xroot. Note
-      // we know that sub-root is reduction and has only 1 non-tiled dim.
-      irregular_xroots.insert(op);
-      continue;
-    } else if (*std::max_element(res_non_tiled_dims.begin(),
-                                 res_non_tiled_dims.end()) !=
-               res_non_tiled_dims.size() - 1) {
-      // Check constraint: the tiled dims are all minor dims for all regular
-      // xroot ops. This means the non-tiled dims should all be the major dims.
-      irregular_xroots.insert(op);
-      continue;
-    }
 
-    auto& plan = tile_plan[res];
-    int64_t rank = res.getType().cast<MemRefType>().getRank();
-    for (int64_t i = 0; i < rank; i++) {
-      if (res_non_tiled_dims.count(i) == 0) {
-        // Note that we only log whether a dimension is tiled or not. We do not
-        // care about the tiling size. Thus we set all tiled dims with with the
-        // value `kDynamicSize`.
+      auto& plan = tile_plan[res];
+      for (int i = last_non_tied_dim + 1; i < result_rank; ++i) {
+        // Note that we only log whether a dimension is tiled or not. We do
+        // not care about the tiling size. Thus we set all tiled dims with
+        // with the value `kDynamicSize`.
         plan.tileSizes[i] = ShapedType::kDynamicSize;
       }
+      regular_xroots.insert(op);
     }
-    regular_xroots.insert(op);
   }
 
   // Check constraint: all the external-only-roots should be regular. This is
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
index b8759511621..9599771fcb4 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"               // TF:local_config_mlir
 #include "mlir/Transforms/RegionUtils.h"  // TF:llvm-project
+#include "tensorflow/compiler/mlir/disc/disc_util.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/disc_shape_optimization_utils.h"
 #include "tensorflow/compiler/mlir/disc/transforms/fusion_utils.h"
@@ -510,8 +511,13 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
     SmallVector<Block*, 4> blocks;
     CollectBlocksInsideFunction(func, blocks);
 
-    ShapeAnalysis shapeAnalysis(func);
-    shapeAnalysis.run();
+    std::unique_ptr<ShapeAnalysis> shapeAnalysisPtr;
+    if (useShapeConstraintIR()) {
+      shapeAnalysisPtr.reset(new ShapeConstraintIRAnalysis(func));
+    } else {
+      shapeAnalysisPtr.reset(new ShapeAnalysisDeprecated{func});
+      static_cast<ShapeAnalysisDeprecated*>(shapeAnalysisPtr.get())->run();
+    }
 
     // process each block and do fusion within a block.
     tensorflow::ReadInt64FromEnvVar("disc_debug_max_fusion_numbers_", INT_MIN,
@@ -519,7 +525,7 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
     FusionPipeline pipeline = makeFusionPipeline();
     int64_t fusion_pattern_number = 0;
     for (Block* block : blocks) {
-      FusionPlanner planner(pipeline, block, &shapeAnalysis);
+      FusionPlanner planner(pipeline, block, shapeAnalysisPtr.get());
       llvm::Optional<FusionPlan> plan = planner.Run();
       if (!plan) {
         emitError(func.getLoc(),
@@ -561,7 +567,7 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
     func.walk([&](FusionOp op) {
       StringRef fusionName = getFusionName(op);
       if (!fusionName.empty()) return;
-      FusionPattern pattern(op, &shapeAnalysis);
+      FusionPattern pattern(op, shapeAnalysisPtr.get());
       auto signature = generateSignatureForFusion(pattern);
       if (!nameSet.count(signature)) {
         nameVec.push_back(signature);
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
index 51547f195c9..8ecfefe4d1e 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "tensorflow/compiler/mlir/disc/IR/lhlo_disc_ops.h"
+#include "tensorflow/compiler/mlir/disc/disc_util.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/codegen_utils.h"
 #include "tensorflow/compiler/mlir/disc/transforms/fusion_utils.h"
@@ -57,10 +58,14 @@ LogicalResult elemwiseLowerHelper(OpBuilder& b, Location loc, Operation* op,
 
   Value result_memref = cast<lmhlo::LmhloOp>(op).getResultBuffer();
   Value memref = result_memref;
-  if (shape_analysis != nullptr) {
+  // TODO(disc): we acutally don't need this after shape constraint refactor
+  // Remove this part of code after we finish the benchamrk after the refactor.
+  if (auto analysis_deprecated =
+          dynamic_cast<const ShapeAnalysisDeprecated*>(shape_analysis)) {
     auto parent_fusion = op->getParentOfType<mhlo::FusionOp>();
-    Value leader_memref = shape_analysis->GetLeaderValueWithSameShapeInFusion(
-        parent_fusion, result_memref);
+    Value leader_memref =
+        analysis_deprecated->GetLeaderValueWithSameShapeInFusion(parent_fusion,
+                                                                 result_memref);
     if (leader_memref != nullptr) {
       memref = leader_memref;
     }
@@ -128,10 +133,14 @@ LogicalResult miscLowerHelper(OpBuilder& b, Location loc, Operation* opaque_op,
   if (!op) return failure();
   Value result_memref = cast<lmhlo::LmhloOp>(&*op).getResultBuffer();
   Value memref = result_memref;
-  if (shape_analysis) {
+  // TODO(disc): we acutally don't need this after shape constraint refactor
+  // Remove this part of code after we finish the benchamrk after the refactor.
+  if (auto analysis_deprecated =
+          dynamic_cast<const ShapeAnalysisDeprecated*>(shape_analysis)) {
     lmhlo::FusionOp fusion = opaque_op->getParentOfType<lmhlo::FusionOp>();
-    Value leader_memref = shape_analysis->GetLeaderValueWithSameShapeInFusion(
-        fusion, result_memref);
+    Value leader_memref =
+        analysis_deprecated->GetLeaderValueWithSameShapeInFusion(fusion,
+                                                                 result_memref);
     if (leader_memref != nullptr) {
       memref = leader_memref;
     }
@@ -3397,7 +3406,9 @@ LogicalResult HandleGpuFusionOp(OpBuilder& b, Operation* fusion,
   assert(fusion_op);
 
   FusionPattern fusion_pattern(fusion_op, shape_analysis);
-  {  // Update in-fusio shape equality information.
+  if (auto analysis_deprecated =
+          dynamic_cast<ShapeAnalysisDeprecated*>(shape_analysis)) {
+    // Update in-fusion shape equality information.
     DenseSet<Value> values_in_fusion;
     auto operands = fusion_pattern.getOperands();
     values_in_fusion.insert(operands.begin(), operands.end());
@@ -3405,7 +3416,7 @@ LogicalResult HandleGpuFusionOp(OpBuilder& b, Operation* fusion,
     values_in_fusion.insert(internal_results.begin(), internal_results.end());
     auto results = fusion_pattern.getResults();
     values_in_fusion.insert(results.begin(), results.end());
-    shape_analysis->buildEqualShapesInFusion(fusion, values_in_fusion);
+    analysis_deprecated->buildEqualShapesInFusion(fusion, values_in_fusion);
   }
   LLVM_DEBUG(createPrintFusionParams(fusion_op, fusion_pattern));
 
@@ -3903,12 +3914,23 @@ class DiscLhloLegalizeRootsToParallelLoops
   void runOnOperation() override {
     func::FuncOp func = getOperation();
 
-    ShapeAnalysis shape_analysis(func);
-    if (failed(shape_analysis.run())) {
-      func->emitError("failed to do shape analysis");
-      signalPassFailure();
+    // skip shape constraint graph
+    if (func.getName() == SymbolicDimMgr::getShapeConstraintGraphFunctionName())
       return;
+
+    std::unique_ptr<ShapeAnalysis> shapeAnalysisPtr;
+    if (useShapeConstraintIR()) {
+      shapeAnalysisPtr.reset(new ShapeConstraintIRAnalysis(func));
+    } else {
+      shapeAnalysisPtr.reset(new ShapeAnalysisDeprecated{func});
+      if (failed(static_cast<ShapeAnalysisDeprecated*>(shapeAnalysisPtr.get())
+                     ->run())) {
+        func->emitError("failed to do shape analysis");
+        signalPassFailure();
+        return;
+      }
     }
+    auto& shape_analysis = *shapeAnalysisPtr;
 
     OpBuilder b(func);
     SmallVector<Operation*, 4> gpu_non_fusion_worklist;
diff --git a/tao_compiler/mlir/disc/transforms/shape_utils.cc b/tao_compiler/mlir/disc/transforms/shape_utils.cc
index f980130a1eb..c0f684db1b0 100644
--- a/tao_compiler/mlir/disc/transforms/shape_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/shape_utils.cc
@@ -49,7 +49,7 @@ LogicalResult SymbolDim::Merge(SymbolDim* other) {
   return success();
 }
 
-Type ShapeAnalysis::getRefinedType(Value value) {
+Type ShapeAnalysisDeprecated::getRefinedType(Value value) {
   SymbolShape* symbolShape = getShape(value);
 
   // not-shaped type
@@ -60,7 +60,7 @@ Type ShapeAnalysis::getRefinedType(Value value) {
   return RankedTensorType::get(dims, ty.getElementType());
 }
 
-LogicalResult ShapeAnalysis::run() {
+LogicalResult ShapeAnalysisDeprecated::run() {
   // Make sure only run once.
   if (initialized) {
     op_->emitError() << "re-initialized shape analysis";
@@ -71,7 +71,7 @@ LogicalResult ShapeAnalysis::run() {
   return buildShapeMap();
 }
 
-LogicalResult ShapeAnalysis::buildShapeMap() {
+LogicalResult ShapeAnalysisDeprecated::buildShapeMap() {
   // Do not merge the region loops. We want to compute shape-map, dim-value-map
   // and decompose one-by-one.
   for (auto& region : op_->getRegions()) {
@@ -94,7 +94,7 @@ LogicalResult ShapeAnalysis::buildShapeMap() {
   return success();
 }
 
-void ShapeAnalysis::postProcessingDimValues() {
+void ShapeAnalysisDeprecated::postProcessingDimValues() {
   // Update dimValue roots.
   for (auto sym2val : dimSymbol2DimValue_) {
     dimSymbol2DimValue_[sym2val.first] = getRootDimValue(sym2val.second);
@@ -169,7 +169,7 @@ void ShapeAnalysis::postProcessingDimValues() {
   dimValueMulDecompose_ = std::move(dimValMulDecomp);
 }
 
-LogicalResult ShapeAnalysis::buildSymbolShape(Value value) {
+LogicalResult ShapeAnalysisDeprecated::buildSymbolShape(Value value) {
   if (shapeMap_.find(value) != shapeMap_.end()) {
     return success();
   }
@@ -191,7 +191,7 @@ LogicalResult ShapeAnalysis::buildSymbolShape(Value value) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildRegionShapeMap(Region* region) {
+LogicalResult ShapeAnalysisDeprecated::buildRegionShapeMap(Region* region) {
   // Only SCF is supported a.t.m.
   if (region->getBlocks().size() != 1) {
     return region->getParentOp()->emitError(
@@ -203,7 +203,7 @@ LogicalResult ShapeAnalysis::buildRegionShapeMap(Region* region) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildBlockShapeMap(Block* block) {
+LogicalResult ShapeAnalysisDeprecated::buildBlockShapeMap(Block* block) {
   // mapping block arguments
   for (Value value : block->getArguments()) {
     if (failed(buildSymbolShape(value))) {
@@ -226,7 +226,8 @@ LogicalResult ShapeAnalysis::buildBlockShapeMap(Block* block) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildOperationShapeValueMap(Operation* op) {
+LogicalResult ShapeAnalysisDeprecated::buildOperationShapeValueMap(
+    Operation* op) {
   auto updateShapeValueEqual = [&](Value lhs, Value rhs) {
     if (!isShapeValueEqual(lhs, rhs)) {
       if (!shapeValueEqual_.isEquivalent(ValueWrapper(lhs),
@@ -269,7 +270,7 @@ LogicalResult ShapeAnalysis::buildOperationShapeValueMap(Operation* op) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildOperationShapeMap(Operation* op) {
+LogicalResult ShapeAnalysisDeprecated::buildOperationShapeMap(Operation* op) {
   // build shapes for the results of op
   for (Value result : op->getResults()) {
     if (failed(buildSymbolShape(result))) {
@@ -284,7 +285,7 @@ LogicalResult ShapeAnalysis::buildOperationShapeMap(Operation* op) {
   return applyOpConstraint(op);
 }
 
-LogicalResult ShapeAnalysis::buildRegionDimValueMap(Region* region) {
+LogicalResult ShapeAnalysisDeprecated::buildRegionDimValueMap(Region* region) {
   // Only SCF is supported a.t.m.
   if (region->getBlocks().size() != 1) {
     return region->getParentOp()->emitError(
@@ -297,7 +298,7 @@ LogicalResult ShapeAnalysis::buildRegionDimValueMap(Region* region) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildBlockDimValueMap(Block* block) {
+LogicalResult ShapeAnalysisDeprecated::buildBlockDimValueMap(Block* block) {
   // TODO: deal with memref.load and store directory rather than infer from
   // reshape.
   auto extractLmhloValueOfDim = [&](Value valueOfDims, int64_t index,
@@ -612,8 +613,9 @@ LogicalResult ShapeAnalysis::buildBlockDimValueMap(Block* block) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildDimValueMap(Value operand, int64_t dim,
-                                              DimValue dimValue) {
+LogicalResult ShapeAnalysisDeprecated::buildDimValueMap(Value operand,
+                                                        int64_t dim,
+                                                        DimValue dimValue) {
   auto symbol = getDim(operand, dim);
   if (symbol == nullptr) {
     return failure();
@@ -631,7 +633,7 @@ LogicalResult ShapeAnalysis::buildDimValueMap(Value operand, int64_t dim,
   return success();
 }
 
-SymbolShape* ShapeAnalysis::getShape(Value value) {
+SymbolShape* ShapeAnalysisDeprecated::getShape(Value value) {
   auto it = shapeMap_.find(value);
   if (it == shapeMap_.end()) return nullptr;
 
@@ -641,7 +643,7 @@ SymbolShape* ShapeAnalysis::getShape(Value value) {
   return &it->second;
 }
 
-SymbolDim* ShapeAnalysis::getRootDim(SymbolDim* symbolDim) {
+SymbolDim* ShapeAnalysisDeprecated::getRootDim(SymbolDim* symbolDim) {
   assert(symbolDim != nullptr);
   SymbolDim* parentSymbolDim = symbolDim;
   do {
@@ -653,7 +655,7 @@ SymbolDim* ShapeAnalysis::getRootDim(SymbolDim* symbolDim) {
   return parentSymbolDim;
 }
 
-SymbolDim* ShapeAnalysis::getDim(Value value, int64_t dim) {
+SymbolDim* ShapeAnalysisDeprecated::getDim(Value value, int64_t dim) {
   auto it = shapeMap_.find(value);
   if (it == shapeMap_.end()) return nullptr;
 
@@ -666,11 +668,11 @@ SymbolDim* ShapeAnalysis::getDim(Value value, int64_t dim) {
   return symbolDim;
 }
 
-DimValue ShapeAnalysis::getRootDimValue(DimValue dimValue) {
+DimValue ShapeAnalysisDeprecated::getRootDimValue(DimValue dimValue) {
   return dimValueEquivalence_.getOrInsertLeaderValue(dimValue);
 }
 
-DimValue ShapeAnalysis::getDimValue(SymbolDim* symbolDim) {
+DimValue ShapeAnalysisDeprecated::getDimValue(SymbolDim* symbolDim) {
   if (symbolDim == nullptr) {
     return DimValue(Value(nullptr));
   }
@@ -687,7 +689,7 @@ DimValue ShapeAnalysis::getDimValue(SymbolDim* symbolDim) {
   return rootDimValue;
 }
 
-DimValue ShapeAnalysis::getDimValue(Value operand, int64_t dim) {
+DimValue ShapeAnalysisDeprecated::getDimValue(Value operand, int64_t dim) {
   auto ty = operand.getType().dyn_cast<ShapedType>();
   if (!ty.hasRank() || dim >= ty.getRank()) {
     return DimValue(Value(nullptr));
@@ -711,7 +713,8 @@ DimValue ShapeAnalysis::getDimValue(Value operand, int64_t dim) {
   return getDimValue(symbolDim);
 }
 
-LogicalResult ShapeAnalysis::mapDimEqual(SymbolDim* lhs, SymbolDim* rhs) {
+LogicalResult ShapeAnalysisDeprecated::mapDimEqual(SymbolDim* lhs,
+                                                   SymbolDim* rhs) {
   if (lhs == nullptr || rhs == nullptr) {
     return success();
   }
@@ -731,7 +734,8 @@ LogicalResult ShapeAnalysis::mapDimEqual(SymbolDim* lhs, SymbolDim* rhs) {
   }
 }
 
-LogicalResult ShapeAnalysis::mapShapeEqual(SymbolShape* lhs, SymbolShape* rhs) {
+LogicalResult ShapeAnalysisDeprecated::mapShapeEqual(SymbolShape* lhs,
+                                                     SymbolShape* rhs) {
   if (!lhs || !rhs || lhs->rank() != rhs->rank()) return failure();
 
   for (auto&& en : llvm::zip(lhs->getSymbolDims(), rhs->getSymbolDims())) {
@@ -740,7 +744,7 @@ LogicalResult ShapeAnalysis::mapShapeEqual(SymbolShape* lhs, SymbolShape* rhs) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::mapShapeEqual(Value lhs, Value rhs) {
+LogicalResult ShapeAnalysisDeprecated::mapShapeEqual(Value lhs, Value rhs) {
   SymbolShape* lhsShape = getShape(lhs);
   SymbolShape* rhsShape = getShape(rhs);
 
@@ -777,23 +781,24 @@ LogicalResult ShapeAnalysis::mapShapeEqual(Value lhs, Value rhs) {
   return mapShapeEqual(lhsShape, rhsShape);
 }
 
-LogicalResult ShapeAnalysis::mapDimEqual(Value lhs, int64_t lhsIdx, Value rhs,
-                                         int64_t rhsIdx) {
+LogicalResult ShapeAnalysisDeprecated::mapDimEqual(Value lhs, int64_t lhsIdx,
+                                                   Value rhs, int64_t rhsIdx) {
   SymbolDim* lhsDim = getDim(lhs, lhsIdx);
   SymbolDim* rhsDim = getDim(rhs, rhsIdx);
   return mapDimEqual(lhsDim, rhsDim);
 }
 
 // TODO: map const and constOp equal.
-LogicalResult ShapeAnalysis::mapDimValueEqual(DimValue lhs, DimValue rhs) {
+LogicalResult ShapeAnalysisDeprecated::mapDimValueEqual(DimValue lhs,
+                                                        DimValue rhs) {
   if (!dimValueEquivalence_.isEquivalent(lhs, rhs)) {
     dimValueEquivalence_.unionSets(lhs, rhs);
   }
   return success();
 }
 
-bool ShapeAnalysis::isDimEqual(Value lhs, int64_t lhsDim, Value rhs,
-                               int64_t rhsDim) {
+bool ShapeAnalysisDeprecated::isDimEqual(Value lhs, int64_t lhsDim, Value rhs,
+                                         int64_t rhsDim) {
   auto lhs_ty = lhs.getType().dyn_cast<ShapedType>();
   auto rhs_ty = rhs.getType().dyn_cast<ShapedType>();
   if (!lhs_ty.hasRank() || !rhs_ty.hasRank()) {
@@ -813,7 +818,7 @@ bool ShapeAnalysis::isDimEqual(Value lhs, int64_t lhsDim, Value rhs,
          (getDimValue(lhs, lhsDim) == getDimValue(rhs, rhsDim));
 }
 
-SmallVector<std::vector<DimValue>> ShapeAnalysis::getDimDecompose(
+SmallVector<std::vector<DimValue>> ShapeAnalysisDeprecated::getDimDecompose(
     Value value, int64_t index) {
   auto dim_val = getDimValue(value, index);
   auto iter = dimValueMulDecompose_.find(dim_val);
@@ -822,7 +827,7 @@ SmallVector<std::vector<DimValue>> ShapeAnalysis::getDimDecompose(
              : iter->second;
 }
 
-bool ShapeAnalysis::isShapeEqual(Value lhs, Value rhs) {
+bool ShapeAnalysisDeprecated::isShapeEqual(Value lhs, Value rhs) {
   SymbolShape* lhsShape = getShape(lhs);
   SymbolShape* rhsShape = getShape(rhs);
   if (!lhsShape || !lhsShape) return false;
@@ -833,15 +838,15 @@ bool ShapeAnalysis::isShapeEqual(Value lhs, Value rhs) {
   return true;
 }
 
-bool ShapeAnalysis::isShapeValueEqual(Value lhs, Value rhs) {
+bool ShapeAnalysisDeprecated::isShapeValueEqual(Value lhs, Value rhs) {
   return shapeValueEqual_.isEquivalent(ValueWrapper(lhs), ValueWrapper(rhs));
 }
 
-bool ShapeAnalysis::isDimValueEqual(DimValue lhs, DimValue rhs) {
+bool ShapeAnalysisDeprecated::isDimValueEqual(DimValue lhs, DimValue rhs) {
   return dimValueEquivalence_.isEquivalent(lhs, rhs);
 }
 
-LogicalResult ShapeAnalysis::applyOpConstraint(Operation* op) {
+LogicalResult ShapeAnalysisDeprecated::applyOpConstraint(Operation* op) {
   if (op->getDialect() == op->getContext()->getLoadedDialect("tensor")) {
     return applyTensorOpConstraint(op);
   } else if (op->getDialect() == op->getContext()->getLoadedDialect("mhlo") ||
@@ -857,14 +862,14 @@ LogicalResult ShapeAnalysis::applyOpConstraint(Operation* op) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::applyTensorOpConstraint(Operation* op) {
+LogicalResult ShapeAnalysisDeprecated::applyTensorOpConstraint(Operation* op) {
   if (isa<tensor::CastOp>(op)) {
     return mapShapeEqual(op->getResult(0), op->getOperand(0));
   }
   return success();
 }
 
-LogicalResult ShapeAnalysis::applyMhloOpConstraint(Operation* op) {
+LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
   if (op->hasTrait<mlir::OpTrait::SameOperandsAndResultType>() ||
       op->hasTrait<mlir::OpTrait::SameOperandsAndResultShape>()) {
     Value ref;
@@ -1086,7 +1091,7 @@ LogicalResult ShapeAnalysis::applyMhloOpConstraint(Operation* op) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::applyLmhloOpConstraint(Operation* op) {
+LogicalResult ShapeAnalysisDeprecated::applyLmhloOpConstraint(Operation* op) {
   auto mapValueEquivalent = [&](Value lhs, Value rhs,
                                 EquivalenceClasses<ValueWrapper>& impl) {
     if (!impl.isEquivalent(ValueWrapper(lhs), ValueWrapper(rhs))) {
@@ -1240,7 +1245,7 @@ LogicalResult ShapeAnalysis::applyLmhloOpConstraint(Operation* op) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::applyDimEqualWithDimValue() {
+LogicalResult ShapeAnalysisDeprecated::applyDimEqualWithDimValue() {
   std::unordered_map<DimValue, SmallVector<SymbolDim*, 4>, DimValueHash>
       dimValue2DimSymbols;
   for (auto symbol_value : dimSymbol2DimValue_) {
@@ -1262,7 +1267,7 @@ LogicalResult ShapeAnalysis::applyDimEqualWithDimValue() {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildRegionMulDecompose(Region* region) {
+LogicalResult ShapeAnalysisDeprecated::buildRegionMulDecompose(Region* region) {
   if (region->getBlocks().size() != 1) {
     return region->getParentOp()->emitError(
         "only single block region is supported");
@@ -1273,7 +1278,7 @@ LogicalResult ShapeAnalysis::buildRegionMulDecompose(Region* region) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildBlockMulDecompose(Block* block) {
+LogicalResult ShapeAnalysisDeprecated::buildBlockMulDecompose(Block* block) {
   // mapping each op inside the block
   WalkResult result = block->walk([&](Operation* op) {
     if (auto muli = dyn_cast_or_null<arith::MulIOp>(op)) {
@@ -1358,7 +1363,8 @@ LogicalResult ShapeAnalysis::buildBlockMulDecompose(Block* block) {
   }
 }
 
-LogicalResult ShapeAnalysis::visitSymbolShapes(SymbolShapeVisitor visitor) {
+LogicalResult ShapeAnalysisDeprecated::visitSymbolShapes(
+    SymbolShapeVisitor visitor) {
   for (auto it = shapeMap_.begin(), stop = shapeMap_.end(); it != stop; ++it) {
     Value value = it->first;
     SymbolShape& symbolShape = it->second;
@@ -1378,7 +1384,7 @@ LogicalResult ShapeAnalysis::visitSymbolShapes(SymbolShapeVisitor visitor) {
   return success();
 }
 
-LogicalResult ShapeAnalysis::buildSymbolDimInstances(
+LogicalResult ShapeAnalysisDeprecated::buildSymbolDimInstances(
     DenseMap<SymbolDim*, SmallVector<Value>>& symbolDim2Instances,
     DenseMap<Value, SmallVector<Value>>& shapedValue2Dims) {
   // Instantiate each symbol dim and group all instances for each symbol dim.
@@ -1398,7 +1404,7 @@ LogicalResult ShapeAnalysis::buildSymbolDimInstances(
   return visitSymbolShapes(visitor);
 }
 
-LogicalResult ShapeAnalysis::buildSymbolDimInstancesDominantMap(
+LogicalResult ShapeAnalysisDeprecated::buildSymbolDimInstancesDominantMap(
     DenseMap<SymbolDim*, SmallVector<Value>>& instanceMap,
     DenseMap<SymbolDim*, DenseMap<Value, Value>>& dominantMap) {
   DominanceInfo dominanceInfo(op_);
@@ -1437,7 +1443,7 @@ LogicalResult ShapeAnalysis::buildSymbolDimInstancesDominantMap(
   return success();
 }
 
-void ShapeAnalysis::dumpSymbol2InstancesDominant(
+void ShapeAnalysisDeprecated::dumpSymbol2InstancesDominant(
     Symbol2InstancesDominantType symbol2InstancesDominant) {
   llvm::dbgs() << "symbol2InstancesDominant: "
                << symbol2InstancesDominant.size() << "\n";
@@ -1511,7 +1517,7 @@ void ShapeAnalysis::dumpSymbol2InstancesDominant(
 // value if they are the same (e.g. the operands for allocOp or reinterpret_cast
 // itself if it's external buffer) and following passes can simply reply on
 // normal CSE & canonicalization pass to simplify index computation.
-LogicalResult ShapeAnalysis::buildTieShapeOps() {
+LogicalResult ShapeAnalysisDeprecated::buildTieShapeOps() {
   // create a dim op for each dimension of a shaped Value and group all such dim
   // ops by SymbolDim. Examples:
   //   %1 = mhlo.abs(%0) : (tensor<?xf32>) -> tensor<?xf32>
@@ -1576,7 +1582,7 @@ LogicalResult ShapeAnalysis::buildTieShapeOps() {
   return visitSymbolShapes(visitor);
 }
 
-bool ShapeAnalysis::HasSameNumElements(Value lhs, Value rhs) {
+bool ShapeAnalysisDeprecated::isSameNumElements(Value lhs, Value rhs) {
   if (valueWithSameElements_.isEquivalent(ValueWrapper(lhs),
                                           ValueWrapper(rhs))) {
     return true;
@@ -1598,7 +1604,7 @@ bool ShapeAnalysis::HasSameNumElements(Value lhs, Value rhs) {
   return (lhs_mapped == lhs_rank) && (rhs_mapped == rhs_rank);
 }
 
-bool ShapeAnalysis::extractContinuousDimEqualInfo(
+bool ShapeAnalysisDeprecated::extractContinuousDimEqualInfo(
     Value lhs, Value rhs,
     SmallVector<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>>& equal) {
   equal.clear();
@@ -1777,7 +1783,7 @@ bool ShapeAnalysis::extractContinuousDimEqualInfo(
   return true;
 }
 
-void ShapeAnalysis::buildEqualShapesInFusion(
+void ShapeAnalysisDeprecated::buildEqualShapesInFusion(
     const Operation* fusion, const DenseSet<Value>& values_in_fusion) {
   auto& equal_in_fusion = valueWithEqualShapeInFusion_[fusion];
   DenseMap<SymbolShape*, SmallVector<Value>> value_classes;
@@ -1796,7 +1802,8 @@ void ShapeAnalysis::buildEqualShapesInFusion(
   }
 }
 
-Value ShapeAnalysis::GetLeaderValueWithSameShapeGlobal(Value val) const {
+Value ShapeAnalysisDeprecated::GetLeaderValueWithSameShapeGlobal(
+    Value val) const {
   if (valueWithEqualShape_.findLeader(ValueWrapper(val)) ==
       valueWithEqualShape_.member_end()) {
     return nullptr;
@@ -1804,7 +1811,7 @@ Value ShapeAnalysis::GetLeaderValueWithSameShapeGlobal(Value val) const {
   return valueWithEqualShape_.getLeaderValue(ValueWrapper(val)).getValue();
 }
 
-Value ShapeAnalysis::GetLeaderValueWithSameShapeInFusion(
+Value ShapeAnalysisDeprecated::GetLeaderValueWithSameShapeInFusion(
     const Operation* fusion, Value val) const {
   if (fusion == nullptr) {
     return nullptr;
@@ -1822,7 +1829,7 @@ Value ShapeAnalysis::GetLeaderValueWithSameShapeInFusion(
 
 // update with view map
 
-bool ShapeAnalysis::getConstIntValue(Operation* op, int64_t& result) {
+bool ShapeAnalysisDeprecated::getConstIntValue(Operation* op, int64_t& result) {
   if (auto constOp = dyn_cast_or_null<arith::ConstantIndexOp>(op)) {
     result = constOp.value();
     return true;
@@ -1892,5 +1899,117 @@ void OpListShapeAnalysis::PropagateEquality(
   } while (!converged);
 }
 
+bool ShapeAnalysis::isSameNumElements(Value lhs, Value rhs) {
+  if (lhs == rhs) return true;
+
+  auto lhsTy = lhs.getType().dyn_cast<ShapedType>();
+  auto rhsTy = rhs.getType().dyn_cast<ShapedType>();
+
+  if (!lhsTy || !rhsTy || !lhsTy.hasRank() || !rhsTy.hasRank()) return false;
+
+  return isProductEqual(lhs, 0, lhsTy.getRank(), rhs, 0, rhsTy.getRank());
+}
+
+bool ShapeAnalysis::isProductEqual(Value lhs, int lhsFrom, int lhsTo, Value rhs,
+                                   int rhsFrom, int rhsTo) {
+  SmallVector<int> lhsDimIdxs, rhsDimIdxs;
+  lhsDimIdxs.reserve(lhsTo - lhsFrom);
+  rhsDimIdxs.reserve(rhsTo - rhsFrom);
+  for (int i = lhsFrom; i < lhsTo; ++i) lhsDimIdxs.push_back(i);
+  for (int i = rhsFrom; i < rhsTo; ++i) rhsDimIdxs.push_back(i);
+
+  return isProductEqual(lhs, lhsDimIdxs, rhs, rhsDimIdxs);
+}
+
+ShapeConstraintIRAnalysis::ShapeConstraintIRAnalysis(Operation* op)
+    : op_(op), mgr_(op->getParentOfType<ModuleOp>()) {
+  mgr_.load();
+  ModuleOp m = op_->getParentOfType<ModuleOp>();
+  op_->walk([&](Operation* op) {
+    if (!isa<memref::AllocOp, memref::ReinterpretCastOp>(op)) return;
+    Value result = op->getResult(0);
+    auto resultTy = result.getType().dyn_cast<MemRefType>();
+    // Only support ranked memref types.
+    if (!resultTy) return;
+    // Ealry return if the memref is static or does not have symbolic dim attr
+    auto attrs = op->getAttrOfType<ArrayAttr>(
+        disc_shape::SymbolicDimOp::getSymbolicDimAttrName());
+    if (!attrs) return;
+    auto& symbols = memrefValue2SymDims_[result];
+    for (const auto& attr : attrs) {
+      auto symOp = mgr_.symbolTable().lookup<disc_shape::SymbolicDimOp>(
+          attr.cast<FlatSymbolRefAttr>().getValue());
+      assert(symOp && "symbolic op is not found");
+      symbols.push_back(symOp);
+    }
+  });
+}
+
+ShapeConstraintIRAnalysis::~ShapeConstraintIRAnalysis() { mgr_.save(); }
+
+bool ShapeConstraintIRAnalysis::isShapeEqual(Value lhs, Value rhs) {
+  if (lhs == rhs) return true;
+
+  auto lhsTy = lhs.getType().dyn_cast<ShapedType>();
+  auto rhsTy = rhs.getType().dyn_cast<ShapedType>();
+
+  if (!lhsTy || !rhsTy || !lhsTy.hasRank() || !rhsTy.hasRank()) return false;
+
+  if (lhsTy.hasStaticShape() && rhsTy.hasStaticShape()) {
+    return lhsTy.getShape() == rhsTy.getShape();
+  }
+
+  auto lhsIt = memrefValue2SymDims_.find(lhs);
+  auto rhsIt = memrefValue2SymDims_.find(rhs);
+
+  if (lhsIt == memrefValue2SymDims_.end() ||
+      rhsIt == memrefValue2SymDims_.end() ||
+      lhsIt->second.size() != rhsIt->second.size())
+    return false;
+
+  SmallVector<disc_shape::SymbolicDimOp> lhsSyms;
+  SmallVector<disc_shape::SymbolicDimOp> rhsSyms;
+  for (auto sym : lhsIt->second) {
+    lhsSyms.push_back(mgr_.getRootSymbolicDim(sym));
+  }
+  for (auto sym : rhsIt->second) {
+    rhsSyms.push_back(mgr_.getRootSymbolicDim(sym));
+  }
+  return lhsSyms == rhsSyms;
+}
+
+bool ShapeConstraintIRAnalysis::isProductEqual(Value lhs,
+                                               ArrayRef<int> lhsDimIdxs,
+                                               Value rhs,
+                                               ArrayRef<int> rhsDimIdxs) {
+  SymbolicDimProduct lhsProd;
+  SymbolicDimProduct rhsProd;
+
+  auto buildSymbolicDimProduct = [&](SymbolicDimProduct& prod, Value value,
+                                     ArrayRef<int> dimIdxs) {
+    auto ty = value.getType().dyn_cast<ShapedType>();
+    auto it = memrefValue2SymDims_.find(value);
+    if (!ty || !ty.hasRank()) return false;
+
+    for (int idx : dimIdxs) {
+      if (ty.getShape()[idx] == ShapedType::kDynamicSize) {
+        if (it == memrefValue2SymDims_.end() || it->second.size() <= idx)
+          return false;
+        prod.symbols.push_back(it->second[idx]);
+      } else {
+        prod.factor *= ty.getShape()[idx];
+      }
+    }
+    return true;
+  };
+
+  if (!buildSymbolicDimProduct(lhsProd, lhs, lhsDimIdxs) ||
+      !buildSymbolicDimProduct(rhsProd, rhs, rhsDimIdxs)) {
+    return false;
+  }
+
+  return mgr_.isSymbolicDimProductEqual(lhsProd, rhsProd);
+}
+
 }  // namespace disc_ral
 }  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/shape_utils.h b/tao_compiler/mlir/disc/transforms/shape_utils.h
index 13583ff94cf..468a75de20e 100644
--- a/tao_compiler/mlir/disc/transforms/shape_utils.h
+++ b/tao_compiler/mlir/disc/transforms/shape_utils.h
@@ -16,6 +16,7 @@
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "tensorflow/compiler/mlir/disc/transforms/disc_shape_optimization_utils.h"
 
 #ifndef TENSORFLOW_COMPILER_MLIR_XLA_SHAPE_UTILS_H_
 #define TENSORFLOW_COMPILER_MLIR_XLA_SHAPE_UTILS_H_
@@ -160,11 +161,76 @@ struct DimValueContainerHash {
   }
 };
 
+// A helper class to query and manipulate shape constraint IR on buffer level.
+class ShapeAnalysis {
+ public:
+  virtual ~ShapeAnalysis() = default;
+
+  // Returns true if the two value have the same symbolic shape.
+  virtual bool isShapeEqual(Value lhs, Value rhs) = 0;
+
+  // Suppose:
+  //  lhsDimIdxs = {ld0, ld1, ...}
+  //  rhsDimIdxs = {rd0, rd1, ...}
+  // Returns true if `lhs.shape[ld0] * lhs.shape[ld1] * ... ==
+  // rhs.shape[rd0] * rhs.shape[rd1] * ...`
+  virtual bool isProductEqual(Value lhs, ArrayRef<int> lhsDimIdxs, Value rhs,
+                              ArrayRef<int> rhsDimIdxs) = 0;
+
+  // Returns true if:
+  //  lhs.shape[lhsFrom] * ... lhs.shape[lhsTo-1] ==
+  //  rhs.shape[rhsFrom] * ... rhs.shape[rhsTo-1]
+  virtual bool isProductEqual(Value lhs, int lhsFrom, int lhsTo, Value rhs,
+                              int rhsFrom, int rhsTo);
+
+  // Returns true if the two value have the same number elements.
+  virtual bool isSameNumElements(Value lhs, Value rhs);
+};
+
+// A subclass to impement `ShapeAnalysis` on buffer level.
+// The implementation is based on shape constraint ir.
+class ShapeConstraintIRAnalysis : public ShapeAnalysis {
+ public:
+  // Build shape related analysis on the provided `op`.
+  // This generally can be divided into two steps:
+  // 1, load exsiting shape constraint ir (e.g. symbolic dim ops)
+  // 2, build mapping between memref values and symbolic dim ops.
+  explicit ShapeConstraintIRAnalysis(Operation* op);
+  // auto-save updated shape constriant ir when destroying.
+  ~ShapeConstraintIRAnalysis();
+
+  // Returns the `SymbolicDimMgr` this object holds.
+  SymbolicDimMgr& symbolicDimMgr() { return mgr_; }
+  const SymbolicDimMgr& symbolicDimMgr() const { return mgr_; }
+
+  // Returns true if the two value have the same symbolic shape.
+  bool isShapeEqual(Value lhs, Value rhs) override;
+
+  // Suppose:
+  //  lhsDimIdxs = {ld0, ld1, ...}
+  //  rhsDimIdxs = {rd0, rd1, ...}
+  // Returns true if `lhs.shape[ld0] * lhs.shape[ld1] * ... ==
+  // rhs.shape[rd0] * rhs.shape[rd1] * ...`
+  bool isProductEqual(Value lhs, ArrayRef<int> lhsDimIdxs, Value rhs,
+                      ArrayRef<int> rhsDimIdxs) override;
+
+ private:
+  // The operation this analysis runs on.
+  Operation* op_;
+
+  // The `SymbolicDimMgr` this analysis holds.
+  SymbolicDimMgr mgr_;
+
+  // Map a ranked memref value to an array of symbolicDims, each represents one
+  // dimension size of the memref value.
+  DenseMap<Value, SmallVector<disc_shape::SymbolicDimOp>> memrefValue2SymDims_;
+};
+
 // Shape analysis for propagating and analyzing known shape information in
 // compilation time in a given operation.
-class ShapeAnalysis {
+class ShapeAnalysisDeprecated : public ShapeAnalysis {
  public:
-  explicit ShapeAnalysis(Operation* op) : op_(op) {}
+  explicit ShapeAnalysisDeprecated(Operation* op) : op_(op) {}
 
   LogicalResult run();
 
@@ -175,7 +241,7 @@ class ShapeAnalysis {
   SymbolShape* getShape(Value value);
   DimValue getDimValue(Value operand, int64_t dim);
   bool isDimEqual(Value lhs, int64_t lhsDim, Value rhs, int64_t rhsDim);
-  bool isShapeEqual(Value lhs, Value rhs);
+  bool isShapeEqual(Value lhs, Value rhs) override;
   bool isShapeValueEqual(Value lhs, Value rhs);
 
   // Insert tie_shape ops to explicit tie dimension equality in the IR level.
@@ -184,7 +250,12 @@ class ShapeAnalysis {
 
   // Returns true if `lhs` and `rhs` are supposed to have same number of
   // elements.
-  bool HasSameNumElements(Value lhs, Value rhs);
+  bool isSameNumElements(Value lhs, Value rhs) override;
+
+  bool isProductEqual(Value lhs, ArrayRef<int> lhsDimIdxs, Value rhs,
+                      ArrayRef<int> rhsDimIdxs) override {
+    return false;
+  }
 
   // Extract continuous equal dims between `lhs` and `rhs`, with the
   // consideration of dim-linearize (e.g., caused by reshape). Note that when
@@ -311,7 +382,7 @@ class OpListShapeAnalysis {
 
   // Returns true if `lhs` and `rhs` are supposed to have same number of
   // elements.
-  bool HasSameNumElements(Value lhs, Value rhs) {
+  bool isSameNumElements(Value lhs, Value rhs) {
     return same_num_elements_impl_.isEquivalent(ValueWrapper(lhs),
                                                 ValueWrapper(rhs));
   }


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
index bae2b0938aa..64c60927c22 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
@@ -239,3 +239,56 @@ func @main(%arg0 : i32) -> (tensor<i32>, tensor<2x1xi32>) {
   return %0, %1 : tensor<i32>, tensor<2x1xi32>
 }
 
+// -----
+
+// test shape constraint attr attached in op.
+// test symbols dim ops 2/3 are not removed (used in kDiscSymbolicDimAttr)
+
+module {
+  // CHECK-LABEL: @main
+  func @main(%arg0: tensor<?x?xf32, [@S0, @S1]>, %arg1: tensor<?x?xf32, [@S0, @S1]>) -> tensor<?x?xf32, [@S0, @S1]> {
+    %0 = mhlo.add %arg0, %arg1 {kDiscSymbolicDimAttr = [@S2, @S3]} : tensor<?x?xf32, [@S0, @S1]>
+    return %0 : tensor<?x?xf32, [@S0, @S1]>
+  }
+  // Test symbols dim ops 2/3 are not removed (used in kDiscSymbolicDimAttr)
+  // CHECK-DAG: "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S0", value = -1 : i64}
+  // CHECK-DAG: "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S1", value = -1 : i64}
+  // CHECK-DAG: "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S2", value = -1 : i64}
+  // CHECK-DAG: "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S3", value = -1 : i64}
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S0", value = -1 : i64} : () -> ()
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S1", value = -1 : i64} : () -> ()
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S2", value = -1 : i64} : () -> ()
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S3", value = -1 : i64} : () -> ()
+  func @shape_constraint_graph() {
+    return
+  }
+}
+
+// -----
+
+// test shape constraint attr attached in op.
+// test symbols dim ops 2/3 are removed.
+
+module {
+  // CHECK-LABEL: @main
+  // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32, [@[[S0:.*]], @[[S1:.*]]]>, %[[ARG1:.*]]: tensor<2xindex>) -> tensor<?x?xf32, [@[[S0]], @[[S1]]]>
+  func @main(%arg0: tensor<?x?xf32, [@S0, @S1]>, %arg1: tensor<2xindex>) -> tensor<?x?xf32, [@S0, @S1]> {
+    // CHECK:  "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[ARG1]])
+    // CHECK-SAME: kDiscSymbolicDimAttr = [@[[S0]], @[[S1]]]
+    %0 = "mhlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {kDiscSymbolicDimAttr = [@S2, @S3], broadcast_dimensions = dense<[0,1]> : tensor<2xi64>} : (tensor<?x?xf32, [@S0, @S1]>, tensor<2xindex>) -> tensor<?x?xf32, [@S2, @S3]>
+    // CHECK-NOT: tensor.cast
+    %1 = tensor.cast %0 : tensor<?x?xf32, [@S2, @S3]> to tensor<?x?xf32, [@S0, @S1]>
+    return %1 : tensor<?x?xf32, [@S0, @S1]>
+  }
+  // CHECK-DAG: "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S0", value = -1 : i64}
+  // CHECK-DAG: "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S1", value = -1 : i64}
+  // CHECK-NOT: "disc_shape.SymbolicDim"()
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S0", value = -1 : i64} : () -> ()
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S1", value = -1 : i64} : () -> ()
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S2", value = -1 : i64} : () -> ()
+  "disc_shape.SymbolicDim"() {knownNegativeOne = false, knownNonNegative = true, knownNonSizeOne = false, knownNonSizeZero = false, sym_name = "S3", value = -1 : i64} : () -> ()
+  func @shape_constraint_graph() {
+    return
+  }
+}
+