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:

Fix patch:
diff --git a/tao_compiler/mlir/disc/BUILD b/tao_compiler/mlir/disc/BUILD
index f424898e3af..b6e145e9bf4 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -687,6 +687,7 @@ cc_library(
     name = "disc_fusion",
     srcs = ["transforms/lhlo_fusion.cc"],
     deps = [
+        ":disc_shape_optimization_utils",
         ":fusion_utils",
         ":pass_details",
         ":placement_utils",
@@ -1532,6 +1533,7 @@ cc_library(
     name = "disc_algebra_simplifier",
     srcs = ["transforms/disc_algebra_simplifier.cc"],
     deps = [
+        ":disc_shape_optimization_utils",
         ":mhlo_disc",
         ":pass_details",
         "//tensorflow/compiler/mlir/hlo:hlo",
@@ -1547,6 +1549,77 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "disc_shape_optimization_utils",
+    srcs = [
+        "transforms/disc_shape_optimization_utils.cc"
+    ],
+    hdrs = [
+        "transforms/disc_shape_optimization_utils.h",
+    ],
+    deps = [
+        ":disc_shape",
+        ":disc_util",
+        ":pass_details",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Shape",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "disc_shape_optimization",
+    srcs = [
+        "transforms/disc_shape_optimization.cc"
+    ],
+    deps = [
+        ":mhlo_disc",
+        ":disc_shape_optimization_utils",
+        ":disc_util",
+        ":pass_details",
+        "//tensorflow/compiler/mlir/hlo:hlo",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Shape",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "disc_strip_shape_constraint_ops",
+    srcs = [
+        "transforms/disc_strip_shape_constraint_ops.cc"
+    ],
+    deps = [
+        ":disc_shape_optimization_utils",
+        ":disc_util",
+        ":pass_details",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Shape",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "all_passes",
     hdrs = [
@@ -1582,11 +1655,13 @@ cc_library(
         ":disc_parallel_loop_tiling",
         ":disc_remove_dead_buffer",
         ":disc_remove_shape_constraints",
+        ":disc_shape_optimization",
         ":disc_shape_simplifier",
         ":disc_shape_to_std",
         ":disc_specialize_fusion_with_speculation",
         ":disc_std_bufferize",
         ":disc_stitch_fusion",
+        ":disc_strip_shape_constraint_ops",
         ":disc_tensor_to_std",
         ":disc_to_llvm",
         ":disc_transpose_simplifier",
diff --git a/tao_compiler/mlir/disc/IR/disc_shape_ops.cc b/tao_compiler/mlir/disc/IR/disc_shape_ops.cc
index 98f862dda15..fb72cc45683 100644
--- a/tao_compiler/mlir/disc/IR/disc_shape_ops.cc
+++ b/tao_compiler/mlir/disc/IR/disc_shape_ops.cc
@@ -216,6 +216,8 @@ struct IdentityTieShapeOp : public OpRewritePattern<TieShapeOp> {
 
   LogicalResult matchAndRewrite(TieShapeOp op,
                                 PatternRewriter& rewriter) const override {
+    // Do not touch tie_shape op with symbolic dim ref attrs.
+    if (op->hasAttr(SymbolicDimOp::getSymbolicDimAttrName())) return success();
     Value operand = op.value();
     auto operandTy = operand.getType().dyn_cast<RankedTensorType>();
     if (!operandTy) return failure();
@@ -296,6 +298,93 @@ void TieShapeOp::getCanonicalizationPatterns(RewritePatternSet& results,
 
 LogicalResult TieShapeOp::verify() { return Verify(*this); }
 
+//===----------------------------------------------------------------------===//
+// SymbolicDimOp
+//===----------------------------------------------------------------------===//
+
+void SymbolicDimOp::getCanonicalizationPatterns(RewritePatternSet& results,
+                                                MLIRContext* context) {}
+
+LogicalResult SymbolicDimOp::verify() { return Verify(*this); }
+
+int64_t SymbolicDimOp::getDimSize() {
+  if (auto attr = (*this)->getAttrOfType<IntegerAttr>("value"))
+    return attr.getInt();
+  return ShapedType::kDynamicSize;
+}
+
+void SymbolicDimOp::setDimSize(int64_t val) {
+  OpBuilder b(*this);
+  (*this)->setAttr("value", b.getI64IntegerAttr(val));
+  if (val == -1) {
+    setKnownNegativeOne(true);
+  } else if (val >= 0) {
+    setKnownNonNegative(true);
+    if (val != 0) setKnownNonSizeZero(true);
+    if (val != 1) setKnownNonSizeOne(true);
+  }
+}
+
+bool SymbolicDimOp::isDynamic() {
+  return getDimSize() == ShapedType::kDynamicSize;
+}
+
+void SymbolicDimOp::setKnownNonNegative(bool flag) {
+  OpBuilder b(*this);
+  (*this)->setAttr("knownNonNegative", b.getBoolAttr(flag));
+}
+
+void SymbolicDimOp::setKnownNegativeOne(bool flag) {
+  OpBuilder b(*this);
+  (*this)->setAttr("knownNegativeOne", b.getBoolAttr(flag));
+  if (flag) {
+    setKnownNonSizeOne(true);
+    setKnownNonSizeZero(true);
+  }
+}
+
+void SymbolicDimOp::setKnownNonSizeOne(bool flag) {
+  OpBuilder b(*this);
+  (*this)->setAttr("knownNonSizeOne", b.getBoolAttr(flag));
+}
+
+void SymbolicDimOp::setKnownNonSizeZero(bool flag) {
+  OpBuilder b(*this);
+  (*this)->setAttr("knownNonSizeZero", b.getBoolAttr(flag));
+}
+
+LogicalResult SymbolicDimOp::Merge(SymbolicDimOp other) {
+  if (!isDynamic() && !other.isDynamic() && getDimSize() != other.getDimSize())
+    return failure();
+  if (isDynamic() && !other.isDynamic()) setDimSize(other.getDimSize());
+
+  bool knownNonNegativeFlag = knownNonNegative() || other.knownNonNegative();
+  bool knownNegativeOneFlag = knownNegativeOne() || other.knownNegativeOne();
+  bool knownNonSizeOneFlag =
+      knownNonSizeOne() || other.knownNonSizeOne() || knownNegativeOneFlag;
+  bool knownNonSizeZeroFlag =
+      knownNonSizeZero() || other.knownNonSizeZero() || knownNegativeOneFlag;
+
+  if (knownNonNegativeFlag && knownNegativeOneFlag) return failure();
+
+  setKnownNonSizeZero(knownNonSizeZeroFlag);
+  setKnownNonSizeOne(knownNonSizeOneFlag);
+  setKnownNegativeOne(knownNegativeOneFlag);
+  setKnownNonNegative(knownNonNegativeFlag);
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// DimOp
+//===----------------------------------------------------------------------===//
+LogicalResult DimOp::verify() { return Verify(*this); }
+
+//===----------------------------------------------------------------------===//
+// TieProductEqualOp
+//===----------------------------------------------------------------------===//
+LogicalResult TieProductEqualOp::verify() { return Verify(*this); }
+
 }  // namespace disc_shape
 }  // namespace mlir
 
diff --git a/tao_compiler/mlir/disc/IR/disc_shape_ops.td b/tao_compiler/mlir/disc/IR/disc_shape_ops.td
index 6142f3fc6ac..82c127acc32 100644
--- a/tao_compiler/mlir/disc/IR/disc_shape_ops.td
+++ b/tao_compiler/mlir/disc/IR/disc_shape_ops.td
@@ -19,6 +19,7 @@ limitations under the License.
 #define DISC_SHAPE_OPS
 
 include "mlir/IR/OpBase.td"
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def DISCShape_Dialect : Dialect {
@@ -73,4 +74,91 @@ def DISCShape_TieShapeOp: DISCShape_Op<"tie_shape", [NoSideEffect]> {
   let hasCanonicalizer = 1;
 }
 
+def DISCShape_SymbolicDimOp:
+      DISCShape_Op<"SymbolicDim", [IsolatedFromAbove, Symbol]> {
+  let summary = "representing a symbolic dimension size";
+  let description = [{
+    Representing a symbolic dimension size.
+  }];
+  let arguments = (ins
+    SymbolNameAttr:$sym_name,
+    DefaultValuedAttr<I64Attr, "-1">:$value,
+    DefaultValuedAttr<BoolAttr, "false">:$knownNonNegative,
+    DefaultValuedAttr<BoolAttr, "false">:$knownNegativeOne,
+    DefaultValuedAttr<BoolAttr, "false">:$knownNonSizeOne,
+    DefaultValuedAttr<BoolAttr, "false">:$knownNonSizeZero
+    // TODO(disc): add likely values
+  );
+  let extraClassDeclaration = [{
+    // Returns the dimension value this SymbolicDim op represents.
+    int64_t getDimSize();
+
+    // Sets the dimension value this SymbolicDim op represents to `val`
+    void setDimSize(int64_t val);
+
+    // Returns true if this SymbolicDim is not known at compile-time.
+    bool isDynamic();
+
+    // Returns true if we know that the dimension value this SymbolicDim op
+    // represents is size one.
+    bool knownSizeOne() { return getDimSize() == 1; }
+
+    // Sets `knownNonNegative` to the value of `flag`
+    void setKnownNonNegative(bool flag);
+
+    // Sets `knownNegativeOne` to the value of `flag`
+    void setKnownNegativeOne(bool flag);
+
+    // Sets `knownNonSizeOne` to the value of `flag`
+    void setKnownNonSizeOne(bool flag);
+
+    // Sets `knownNonSizeZero` to the value of `flag`
+    void setKnownNonSizeZero(bool flag);
+
+    // Try to merge two SymbolicDim ops.
+    // Returns failure if fail to merge.
+    LogicalResult Merge(SymbolicDimOp other);
+
+    // Returns the name of attribute used for store the symbolicDim op reference attribute.
+    static StringRef getSymbolicDimAttrName() { return "kDiscSymbolicDimAttr"; }
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
+def DISCShape_DimOp:
+      DISCShape_Op<"dim", [NoSideEffect]> {
+  let summary = "binding to a symbolic dimension size";
+  let description = [{
+    Binding to a symbolic dimension size.
+  }];
+  let arguments = (ins
+    FlatSymbolRefAttr:$name
+  );
+  let results = (outs Index);
+}
+
+def SymbolicDimPredicateStructuredInterface : OpInterface<"PredicateOp"> {
+  let cppNamespace = "::mlir::disc_shape";
+
+  let methods = [
+  ];
+}
+
+class DISCShape_PredicateOp<string mnemonic, list<Trait> traits> :
+    DISCShape_Op<mnemonic, !listconcat(
+      [SymbolicDimPredicateStructuredInterface], traits)> {
+}
+
+def DISCShape_TieProductEqualOp: DISCShape_Op<"tie_product_equal", [AttrSizedOperandSegments]> {
+  let summary = "mark two groups x and y satisfy: product(x) == product(y)";
+  let description = [{
+    Mark two groups x and y satisfy: product(x) == product(y)
+  }];
+  let arguments = (ins
+    Arg<Variadic<Index>, "">:$lhs,
+    Arg<Variadic<Index>, "">:$rhs
+  );
+}
+
 #endif // DISC_RAL_OPS
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index 590094422dd..68276a5524d 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -224,8 +224,18 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(createCSEPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
 
-  // propagate some known shape information.
-  pm.addPass(disc_ral::createDiscShapeSimplifierPass());
+  bool enable_shape_constraint_ir = false;
+  tensorflow::ReadBoolFromEnvVar("DISC_ENABLE_SHAPE_CONSTRAINT_IR",
+                                 enable_shape_constraint_ir,
+                                 &enable_shape_constraint_ir);
+  if (!enable_shape_constraint_ir) {
+    // propagate some known shape information.
+    pm.addPass(disc_ral::createDiscShapeSimplifierPass());
+  } else {
+    pm.addNestedPass<FuncOp>(disc_ral::createDiscConvertShapeToStandardPass());
+    // shape-related optimization
+    pm.addPass(disc_ral::createDiscShapeOptimizationPass());
+  }
 
   pm.addNestedPass<FuncOp>(disc_ral::createDiscConvertTensorToStandardPass());
   pm.addNestedPass<FuncOp>(disc_ral::createDiscConvertHloToStandardPass());
@@ -236,12 +246,24 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(disc_ral::createDiscAlgebraSimplifierPass());
   pm.addNestedPass<FuncOp>(disc_ral::createDiscSplitLargeOpsPass());
   pm.addNestedPass<FuncOp>(disc_ral::createDiscDotRewriterPass());
+  if (enable_shape_constraint_ir) {
+    // shape-related optimization
+    pm.addPass(disc_ral::createDiscShapeOptimizationPass());
+  }
 
   // Either merge dots to batched dot or merge dots sharing the same operand.
   pm.addNestedPass<FuncOp>(disc_ral::createDiscDotMergePass());
+  if (enable_shape_constraint_ir) {
+    // shape-related optimization
+    pm.addPass(disc_ral::createDiscShapeOptimizationPass());
+  }
 
   if (gpu_enabled) {
     pm.addNestedPass<FuncOp>(mhlo::createHloCanonicalizeReductionPass());
+    if (enable_shape_constraint_ir) {
+      // shape-related optimization
+      pm.addPass(disc_ral::createDiscShapeOptimizationPass());
+    }
   }
 
   pm.addPass(disc_ral::createDiscMarkShapeCalcOpPass());
@@ -263,8 +285,16 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   tensorflow::ReadBoolFromEnvVar("TAO_MLIR_ENABLE_AMP", false, &enable_fp16);
   pm.addNestedPass<FuncOp>(
       disc_ral::createDiscElementTypeConverterPass(enable_fp16));
+  if (enable_shape_constraint_ir) {
+    // shape-related optimization
+    pm.addPass(disc_ral::createDiscShapeOptimizationPass());
+  }
 
   pm.addNestedPass<FuncOp>(disc_ral::createDiscConvRewriter());
+  if (enable_shape_constraint_ir) {
+    // shape-related optimization
+    pm.addPass(disc_ral::createDiscShapeOptimizationPass());
+  }
   // Run CSE after conv rewriter pass to eliminate some redundant transpose ops.
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<FuncOp>(createCSEPass());
@@ -275,9 +305,20 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
     // side. This pass ensures this property.
     pm.addNestedPass<FuncOp>(disc_ral::createDiscGpuConvPaddingLegalization());
   }
+  if (enable_shape_constraint_ir) {
+    // shape-related optimization
+    pm.addPass(disc_ral::createDiscShapeOptimizationPass());
+    pm.addNestedPass<FuncOp>(disc_ral::createDiscAlgebraSimplifierPass());
+  }
 
-  // Create tie_shape ops to explicitly express dim size equality info.
-  pm.addPass(disc_ral::createDiscShapeSimplifierPass("main", true));
+  if (!enable_shape_constraint_ir) {
+    // Create tie_shape ops to explicitly express dim size equality info.
+    pm.addPass(disc_ral::createDiscShapeSimplifierPass("main", true));
+  } else {
+    // shape-related optimization
+    // Create tie_shape ops to explicitly express dim size equality info.
+    pm.addPass(disc_ral::createDiscShapeOptimizationPass("main", true));
+  }
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<FuncOp>(createCSEPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
@@ -480,6 +521,11 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
   pm.addPass(createStripDebugInfoPass());
 
+  if (enable_shape_constraint_ir) {
+    // remove disc_shape.SymbolicDim ops and related ops.
+    pm.addPass(disc_ral::createStripShapeConstraintOpsPass());
+  }
+
   // Host side codegen: std -> binary
   pm.addPass(disc_ral::createDiscToLLVMPass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
diff --git a/tao_compiler/mlir/disc/transforms/conv_rewriter.cc b/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
index 37cda249026..a27a3cbcc40 100644
--- a/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
+++ b/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
@@ -324,10 +324,12 @@ struct DiscConvRewriterPass
     OpBuilder b(params.conv);
     b.setInsertionPointAfter(params.conv);
     auto new_tp = GetTransposeOutputType(params.output, in2OutAttr, b);
+    auto originalOutputTy = params.output.getType();
     params.output.setType(new_tp);
 
     auto transpose_op =
         InsertTranspose(params.conv, params.output, out2InAttr, b);
+    transpose_op->getResult(0).setType(originalOutputTy);
     params.output.replaceAllUsesWith(transpose_op->getResult(0));
     transpose_op->setOperand(0, params.output);
   }
diff --git a/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc b/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc
index 1e1b7b05c03..8afb3550435 100644
--- a/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc
@@ -20,6 +20,7 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
+#include "tensorflow/compiler/mlir/disc/transforms/disc_shape_optimization_utils.h"
 
 #define DEBUG_TYPE "disc-algebra-simplifier"
 
@@ -51,6 +52,7 @@ struct ExpandPowOp : public OpRewritePattern<mhlo::PowOp> {
     for (int i = 1; i < exponential; ++i)
       newResult =
           rewriter.create<mhlo::MulOp>(loc, newResult, op->getOperand(0));
+    newResult.setType(op->getResult(0).getType());
     rewriter.replaceOp(op, newResult);
     return success();
   }
@@ -115,10 +117,185 @@ struct ExpandPowOp : public OpRewritePattern<mhlo::PowOp> {
   }
 };
 
+// convert:
+//   %1 = mhlo.dynamic_broadcast_in_dim(%0, %target_shape) :
+//     (tensor<?x?xf32, [@S0, @S1]>, ...) -> tensor<?x?xf32, [@S0, @S1]>
+//   use(%1)
+// to:
+//   use(%0)
+template <typename OpTy>
+struct IdentityBroadCastInDimOpCanonicalizationPattern
+    : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter& rewriter) const override {
+    auto fromTy =
+        op->getOperand(0).getType().template dyn_cast<RankedTensorType>();
+    auto toTy =
+        op->getResult(0).getType().template dyn_cast<RankedTensorType>();
+    if (!fromTy || !toTy) return failure();
+    // identity bcast op if fromTy & toTy have the same static shape
+    if (fromTy.hasStaticShape() && toTy.hasStaticShape() && fromTy == toTy) {
+      rewriter.replaceOp(op, op->getOperands());
+      return success();
+    }
+
+    // (partial) dynamic shape cases.
+    // Try to check if the input and out have the same symbolic dim shape.
+    auto fromSymbols = getRankedValueSymbolicDimRefs(op->getOperand(0));
+    auto toSymbols = getRankedValueSymbolicDimRefs(op->getResult(0));
+    if (!fromSymbols || !toSymbols ||
+        (*fromSymbols).size() != (*toSymbols).size())
+      return failure();
+
+    for (const auto& z : llvm::zip(*fromSymbols, *toSymbols))
+      if (std::get<0>(z) != std::get<1>(z)) return failure();
+
+    rewriter.replaceOp(op, op->getOperands());
+    return success();
+  }
+};
+
+// Convert an tf.expand like reshape op followed by a dynamic-broadcast-in-dim
+// op to a single dynamic-broadcast-in-dim op. The dim size of expanded dims
+// should be one.
+//
+// An example is like following:
+//
+//  %0 = ...: tensor<?x?xf32>
+//  %d0 = tensor.dim %0, %c0 : tensor<?x?xf32>
+//  %d1 = tensor.dim %0, %c1 : tensor<?x?xf32>
+//  %new_shape = tensor.from_elements %d0, %d1, %c1 : tensor<3xindex>
+//  %1 = "mhlo.dynamic_reshape"(%0, %new_shape) : (tensor<?x?xf32>,
+//  tensor<3xindex>) -> tensor<?x?x1xf32>
+//  %2 = "mhlo.dynamic_broadcast_in_dim"(%1, %...) {broadcast_dimensions =
+//  dense<[0, 1, 2]> : tensor<3xi64>} : (tensor<?x?x1xf32>, tensor<3xindex>)
+//  ->tensor<?x?x?xf32>
+//
+// This pattern will be converted to:
+//
+//  %2 = "mhlo.dynamic_broadcast_in_dim"(%0, %...) {broadcast_dimensions =
+//  dense<[0, 1]> : tensor<3xi64>} : (tensor<?x?xf32>, tensor<3xindex>) ->
+//  tensor<?x?x?xf32>
+//
+// Note that the reshape op can be either `mhlo::DynamicReshapeOp` or
+// `mhlo::ReshapeOp`. The expanded dim can be in any axis.
+template <typename T>
+LogicalResult tryToExtractDimMap(const SmallVectorImpl<T>& in,
+                                 const SmallVectorImpl<T>& out,
+                                 RankedTensorType outTy,
+                                 DenseMap<int64_t, int64_t>& dimMap) {
+  int64_t inRank = static_cast<int64_t>(in.size());
+  int64_t outRank = static_cast<int64_t>(out.size());
+
+  int64_t outIdx = 0;
+  for (int64_t inIdx = 0; inIdx < inRank; ++inIdx, ++outIdx) {
+    while (in[inIdx] != out[outIdx]) {
+      if (outTy.getShape()[outIdx] != 1) return failure();
+      if (++outIdx >= outRank) return failure();
+    }
+    dimMap[inIdx] = outIdx;
+  }
+  return success();
+}
+
+void newBcastOp(mhlo::BroadcastInDimOp op, PatternRewriter& rewriter,
+                SmallVector<Value>& newOperands,
+                DenseIntElementsAttr bcastAttr) {
+  rewriter.replaceOpWithNewOp<mhlo::BroadcastInDimOp>(
+      op, op.getType(), newOperands[0], bcastAttr);
+}
+
+void newBcastOp(mhlo::DynamicBroadcastInDimOp op, PatternRewriter& rewriter,
+                SmallVector<Value>& newOperands,
+                DenseIntElementsAttr bcastAttr) {
+  rewriter.replaceOpWithNewOp<mhlo::DynamicBroadcastInDimOp>(
+      op, op.getType(), newOperands[0], newOperands[1], bcastAttr);
+}
+
+template <typename OpTy>
+struct BroadCastInDimOfReshapeOpCanonicalizationPattern
+    : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter& rewriter) const override {
+    Value bcastOut = op->getResult(0);
+    Value reshapeOut = op->getOperand(0);
+    auto reshapeOp = reshapeOut.getDefiningOp();
+    if (!reshapeOp || !isa<mhlo::ReshapeOp, mhlo::DynamicReshapeOp>(reshapeOp))
+      return failure();
+    Value reshapeInput = reshapeOp->getOperand(0);
+
+    auto bcastOutTy = bcastOut.getType().template dyn_cast<RankedTensorType>();
+    auto reshapeOutTy =
+        reshapeOut.getType().template dyn_cast<RankedTensorType>();
+    auto reshapeInputTy =
+        reshapeInput.getType().template dyn_cast<RankedTensorType>();
+    if (!bcastOutTy || !reshapeOutTy || !reshapeInputTy) return failure();
+
+    // early stop if not a tf.expand like reshape
+    if (reshapeInputTy.getRank() >= reshapeOutTy.getRank()) return failure();
+
+    bool matched = false;
+    // map one dimension of the input of the reshape to one dimension of the
+    // output of the reshape.
+    DenseMap<int64_t, int64_t> dimMap;
+
+    if (reshapeOutTy.hasStaticShape() && reshapeInputTy.hasStaticShape()) {
+      // static shape casse
+      if (failed(
+              tryToExtractDimMap(llvm::to_vector<4>(reshapeInputTy.getShape()),
+                                 llvm::to_vector<4>(reshapeOutTy.getShape()),
+                                 reshapeOutTy, dimMap))) {
+        return failure();
+      }
+      matched = true;
+    } else {
+      // (partial) dynamic shape cases.
+      // Try to check if the input and out have the compatible symbolic dim
+      // shape.
+      auto reshapeInputSymbols = getRankedValueSymbolicDimRefs(reshapeInput);
+      auto reshapeOutSymbols = getRankedValueSymbolicDimRefs(reshapeOut);
+      if (!reshapeInputSymbols || !reshapeOutSymbols) return failure();
+
+      if (failed(tryToExtractDimMap(*reshapeInputSymbols, *reshapeOutSymbols,
+                                    reshapeOutTy, dimMap))) {
+        return failure();
+      }
+      matched = true;
+    }
+
+    if (!matched) return failure();
+
+    SmallVector<int64_t> newBcastDims;
+    auto oldBcastDims = op.broadcast_dimensions().template getValues<int64_t>();
+    for (size_t d = 0, rank = dimMap.size(); d < rank; ++d)
+      newBcastDims.push_back(oldBcastDims[dimMap[d]]);
+
+    RankedTensorType ty =
+        RankedTensorType::get({static_cast<int64_t>(newBcastDims.size())},
+                              rewriter.getIntegerType(64));
+    SmallVector<Value> newOperands{reshapeInput};
+    for (Value operand : op->getOperands().drop_front())
+      newOperands.push_back(operand);
+
+    newBcastOp(op, rewriter, newOperands,
+               DenseIntElementsAttr::get(ty, newBcastDims));
+    return success();
+  }
+};
+
 void populateDiscAlgebraSimplifierPatterns(RewritePatternSet& patterns) {
   // clang-format off
   patterns.insert<
-    ExpandPowOp
+    BroadCastInDimOfReshapeOpCanonicalizationPattern<mhlo::BroadcastInDimOp>,
+    BroadCastInDimOfReshapeOpCanonicalizationPattern<mhlo::DynamicBroadcastInDimOp>,
+    ExpandPowOp,
+    IdentityBroadCastInDimOpCanonicalizationPattern<mhlo::BroadcastInDimOp>,
+    IdentityBroadCastInDimOpCanonicalizationPattern<mhlo::BroadcastOp>,
+    IdentityBroadCastInDimOpCanonicalizationPattern<mhlo::DynamicBroadcastInDimOp>
   >(patterns.getContext());
   // clang-format on
 }
diff --git a/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc b/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
index 357324d51d4..8bfebac62ee 100644
--- a/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
@@ -431,6 +431,7 @@ bool DotShareOperandMergeConverter::applyMerging(DotCluster& cluster,
           loc, merged_dot, GetI64ElementsAttr(start, &builder),
           GetI64ElementsAttr(limit, &builder),
           GetI64ElementsAttr(strides, &builder));
+      slice.getResult().setType(op->getResult(0).getType());
       op->replaceAllUsesWith(slice);
     }
   } else {
@@ -491,6 +492,7 @@ bool DotShareOperandMergeConverter::applyMerging(DotCluster& cluster,
       auto dyn_slice = builder.create<mhlo::RealDynamicSliceOp>(
           loc, slice_type, merged_dot, start_indices, limit_indices,
           strides_indices);
+      dyn_slice.getResult().setType(op->getResult(0).getType());
 
       op->replaceAllUsesWith(dyn_slice);
     }
diff --git a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
index 59554ae4af7..5647f91657a 100644
--- a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
@@ -201,6 +201,10 @@ struct TieShapeOpConverter : public BaseOpConversion<TieShapeOp> {
 
     Value castedValue = disc_ral::CastMemRefTo(rewriter, loc, memref, memrefTy,
                                                operands.drop_front());
+    StringRef attrName = disc_shape::SymbolicDimOp::getSymbolicDimAttrName();
+    if (op->hasAttr(attrName)) {
+      castedValue.getDefiningOp()->setAttr(attrName, op->getAttr(attrName));
+    }
     rewriter.replaceOp(op, {castedValue});
     return success();
   }
diff --git a/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc b/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc
index a6c169fc614..c764e9263a9 100644
--- a/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc
@@ -53,6 +53,10 @@ LogicalResult ReinterpretCastOpConverter::matchAndRewrite(
   if (!allocTy.hasStaticShape()) {
     alloc->setOperands(op.sizes());
   }
+  StringRef attrName = disc_shape::SymbolicDimOp::getSymbolicDimAttrName();
+  if (op->hasAttr(attrName)) {
+    alloc->setAttr(attrName, op->getAttr(attrName));
+  }
   rewriter.replaceOp(op, alloc->getResults());
   return success();
 }
diff --git a/tao_compiler/mlir/disc/transforms/disc_passes.td b/tao_compiler/mlir/disc/transforms/disc_passes.td
index 5abd4fe20e1..7f9916c982d 100644
--- a/tao_compiler/mlir/disc/transforms/disc_passes.td
+++ b/tao_compiler/mlir/disc/transforms/disc_passes.td
@@ -379,3 +379,26 @@ def DiscAlgebraSimplifierPass : Pass<"disc-algebra-simplifier", "mlir::func::Fun
   let summary = "Apply some basic algebra simplification optimizations";
   let constructor = "createDiscAlgebraSimplifierPass()";
 }
+
+def DiscShapeOptimizationPass : Pass<"disc-shape-optimization", "ModuleOp"> {
+  let summary = "apply some shape-related optimizations.";
+  let constructor = "createDiscShapeOptimizationPass()";
+  let dependentDialects = [
+      "disc_shape::DISCShapeDialect",
+      "shape::ShapeDialect"
+  ];
+  let options = [
+    Option<"entry_func_name_", "entry-func-name", "std::string",
+           /*default=*/"\"main\"", "Name of the entry function.">,
+    Option<"keep_tie_shape_", "keep-tie-shape", "bool",
+           /*default=*/"false", "whether to keep tie shape op after this pass.">,
+  ];
+}
+
+def DiscStripShapeConstraintOpsPass : Pass<"disc-strip-shape-constraint-ops", "ModuleOp"> {
+  let summary = "strip disc shape constraint ops.";
+  let constructor = "createStripShapeConstraintOpsPass()";
+  let dependentDialects = [
+      "disc_shape::DISCShapeDialect",
+  ];
+}
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc b/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
new file mode 100644
index 00000000000..2b1f59823ca
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
@@ -0,0 +1,1474 @@
+/* Copyright 2022 The BladeDISC Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements the logic to do some shape optimizations on tensor
+// level.
+#include <unordered_set>
+#include <utility>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Debug.h"
+#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // TF:llvm-project
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/MLIRContext.h"  // TF:llvm-project
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Pass/Pass.h"  // TF:local_config_mlir
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/Passes.h"  // TF:llvm-project
+#include "tensorflow/compiler/mlir/disc/IR/disc_shape_ops.h"
+#include "tensorflow/compiler/mlir/disc/IR/hlo_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/disc_shape_optimization_utils.h"
+
+namespace mlir {
+namespace disc_ral {
+
+using ::mlir::func::FuncOp;
+
+namespace {
+
+/////////////////////////// Stage #1 BEGIN ////////////////////////////////////
+
+// Insert disc_shape.tie_shape for each value having RankedTensorType type.
+// before
+//  ```
+//    %0 = ... : tensor<?x10xf32>
+//    use(%0)
+//  ```
+// after
+//  ```
+//    %0 = ... : tensor<?x10xf32>
+//    %0_d0 = tensor.dim %0, %c0 : tensor<?x10xf32>
+//    %0_d1 = tensor.dim %0, %c1 : tensor<?x10xf32>
+//    %0_new = disc_shape.tie_shape(%0, %0_d0, %0_d1) : tensor<?x10xf32>
+//    use(%0_new)
+//  ```
+LogicalResult insertTieShapeOnValue(Value value, OpBuilder& b,
+                                    const Location& loc) {
+  // Only insert tie_shape ops for non-zero ranked tensor type
+  auto ty = value.getType().dyn_cast<RankedTensorType>();
+  if (!ty || ty.getRank() == 0) return success();
+
+  DenseSet<Operation*> dimOps;
+  SmallVector<Value> dimSizes;
+  for (int dim = 0, rank = ty.getRank(); dim < rank; ++dim) {
+    auto dimOp = b.create<tensor::DimOp>(loc, value, dim);
+    dimOps.insert(dimOp);
+    dimSizes.push_back(dimOp.getResult());
+  }
+
+  Value newValue = b.create<disc_shape::TieShapeOp>(loc, ty, value, dimSizes);
+  auto users = llvm::to_vector<4>(value.getUsers());
+  for (Operation* user : users) {
+    // skip those dim ops used to fetch the dim size values of original shaped
+    // value.
+    if (dimOps.find(user) != dimOps.end()) continue;
+    if (user == newValue.getDefiningOp()) continue;
+    user->replaceUsesOfWith(value, newValue);
+  }
+  return success();
+}
+
+// forward declaration
+LogicalResult insertTieShapeOnRegion(Region* region);
+
+LogicalResult insertTieShapeOnOperation(Operation* op) {
+  if (isa<disc_shape::TieShapeOp>(op)) return success();
+  // recursively visit the regions of the op.
+  if (!isa<mhlo::ReduceOp, mhlo::ReduceWindowOp>(op)) {
+    for (Region& region : op->getRegions())
+      if (failed(insertTieShapeOnRegion(&region)))
+        return op->emitError("fail to insert tie shape for op's region\n");
+  }
+
+  OpBuilder b(op);
+  b.setInsertionPointAfter(op);
+  for (Value v : op->getResults()) {
+    if (failed(insertTieShapeOnValue(v, b, op->getLoc())))
+      return op->emitError("fail to insert tie shape for op's result\n");
+  }
+
+  return success();
+}
+
+LogicalResult insertTieShapeOnBlock(Block* block) {
+  // mapping block arguments
+  OpBuilder b(block, block->begin());
+  Location loc = block->getParentOp()->getLoc();
+  for (Value value : block->getArguments()) {
+    if (failed(insertTieShapeOnValue(value, b, loc))) {
+      return block->getParentOp()->emitError(
+          "failed to insert tie_shape op for block arg");
+    }
+  }
+
+  // mapping each op inside the block
+  // save a snapshot before visiting in case new ops are inserted during
+  // visiting.
+  SmallVector<Operation*> op_list;
+  for (Operation& op : *block) op_list.push_back(&op);
+  for (Operation* op : op_list) {
+    if (failed(insertTieShapeOnOperation(op))) return failure();
+  }
+  return success();
+}
+
+LogicalResult insertTieShapeOnRegion(Region* region) {
+  for (Block& block : *region) {
+    if (failed(insertTieShapeOnBlock(&block))) return failure();
+  }
+  return success();
+}
+
+// convert:
+//   %shape = shape.shape_of %0 : tensor<?x?xf32> -> tensor<2xindex>
+// to:
+//   %d0 = tensor.dim %0, %c0 : tensor<?x?xf32>
+//   %d1 = tensor.dim %0, %c1 : tensor<?x?xf32>
+//   %shape = tensor.from_elements %d0, %d1 : tensor<2xindex>
+struct ExpandShapeOfOpPattern : public OpRewritePattern<shape::ShapeOfOp> {
+  using OpRewritePattern<shape::ShapeOfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(shape::ShapeOfOp op,
+                                PatternRewriter& rewriter) const override {
+    auto ty = op.getResult().getType().dyn_cast<RankedTensorType>();
+    if (!ty || !ty.hasStaticShape() || !ty.getElementType().isIndex())
+      return failure();
+
+    SmallVector<Value> dimSizes;
+    for (int dim = 0, rank = ty.getShape()[0]; dim < rank; ++dim)
+      dimSizes.push_back(
+          rewriter.create<tensor::DimOp>(op.getLoc(), op.getArg(), dim));
+
+    rewriter.replaceOpWithNewOp<tensor::FromElementsOp>(op, dimSizes);
+    return success();
+  }
+};
+
+/// Fold dim of an operation that implements the InferShapedTypeOpInterface
+template <typename OpTy>
+struct DimOfShapedTypeOpInterface : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy dimOp,
+                                PatternRewriter& rewriter) const override {
+    OpResult dimValue = dimOp.source().template dyn_cast<OpResult>();
+    if (!dimValue) return failure();
+    auto shapedTypeOp =
+        dyn_cast<InferShapedTypeOpInterface>(dimValue.getOwner());
+    if (!shapedTypeOp) return failure();
+
+    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
+    if (!dimIndex) return failure();
+
+    SmallVector<Value> reifiedResultShapes;
+    if (failed(shapedTypeOp.reifyReturnTypeShapes(
+            rewriter, shapedTypeOp->getOperands(), reifiedResultShapes)))
+      return failure();
+
+    if (reifiedResultShapes.size() != shapedTypeOp->getNumResults())
+      return failure();
+
+    Value resultShape = reifiedResultShapes[dimValue.getResultNumber()];
+    auto resultShapeType = resultShape.getType().dyn_cast<RankedTensorType>();
+    if (!resultShapeType ||
+        (!resultShapeType.getElementType().isa<IndexType>() &&
+         !resultShapeType.getElementType().isa<IntegerType>()))
+      return failure();
+
+    Location loc = dimOp.getLoc();
+    Value newValue = rewriter.create<tensor::ExtractOp>(
+        loc, resultShape,
+        rewriter.createOrFold<arith::ConstantIndexOp>(loc, *dimIndex));
+    if (!newValue.getType().isa<IndexType>())
+      newValue = rewriter.create<arith::IndexCastOp>(
+          loc, rewriter.getIndexType(), newValue);
+
+    rewriter.replaceOp(dimOp, newValue);
+    return success();
+  }
+};
+
+// materialize shape computation IR by resolving tensor.dim + op having
+// interface InferShapedTypeOpInterface: before
+//  ```
+//    %arg0 = ... : tensor<?x10xf32>
+//    %0 = mhlo.transpose(%arg0) : tensor<?x?xf32>
+//    %0_d0 = tensor.dim %0, %c0 : tensor<?x?xf32>
+//    %0_d1 = tensor.dim %0, %c1 : tensor<?x?xf32>
+//    %0_new = disc_shape.tie_shape(%0, %0_d0, %0_d1) : tensor<?x?xf32>
+//    use(%0_new)
+//  ```
+// after materializling mhlo.transpose shape interface
+//  ```
+//    %arg0 = ... : tensor<?x10xf32>
+//    %0 = mhlo.transpose(%arg0) {permutaion = [1, 0]} : tensor<?x?xf32>
+//    %0_d0 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+//    %0_d1 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+//    %0_new = disc_shape.tie_shape(%0, %0_d0, %0_d1) : tensor<?x?xf32>
+//    use(%0_new)
+//  ```
+LogicalResult materializeShapeComputation(ModuleOp m, FuncOp main) {
+  // Currently we call inline before all disc passes and thus we do not need to
+  // worry about function call ops. Re-visit this once we change the strategy.
+  if (failed(insertTieShapeOnRegion(&main.getBody()))) {
+    return failure();
+  }
+
+  RewritePatternSet patterns(m.getContext());
+  // clang-format off
+  patterns.add<
+      ExpandShapeOfOpPattern,
+      DimOfShapedTypeOpInterface<tensor::DimOp>
+  >(patterns.getContext());
+  // clang-format on
+
+  if (failed(
+          applyPatternsAndFoldGreedily(m->getRegions(), std::move(patterns)))) {
+    return m.emitError() << "fail to materialize shape computation\n";
+  }
+  return success();
+}
+
+/////////////////////////// Stage #1 END //////////////////////////////////////
+
+/////////////////////////// Stage #2 BEGIN ////////////////////////////////////
+
+using PassPipelineRunner =
+    std::function<LogicalResult(OpPassManager&, ModuleOp)>;
+
+// convert:
+//   %1 = disc_shape.tie_shape %0, %d0, %d1, ... : (tensor<?x?xf32>, ...) ->
+//   tensor<?x?xf32> %dim_size = tensor.dim %1[%c0] : tensor<2xindex>
+//   use(%dim_size)
+// to:
+//   use(%d0)
+struct DimOfTieShapeOpCanonicalizationPattern
+    : public OpRewritePattern<tensor::DimOp> {
+  using OpRewritePattern<tensor::DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::DimOp op,
+                                PatternRewriter& rewriter) const override {
+    auto tieShapeOp = op.source().getDefiningOp<disc_shape::TieShapeOp>();
+    if (!tieShapeOp) return failure();
+    Optional<int64_t> dimIndex = op.getConstantIndex();
+    if (!dimIndex) return failure();
+    rewriter.replaceOp(op, tieShapeOp->getOperand(1 + *dimIndex));
+    return success();
+  }
+};
+
+// convert:
+//   %1 = disc_shape.tie_shape %0, %d0, %d1, ...
+//         : (tensor<?x?xf32>, ...) -> tensor<?x?xf32>
+//   %2 = tensor.extract %1[...] : ...
+// to:
+//   %2 = tensor.extract %0[...] : ...
+struct ExtractElementOfTieShapeOpCanonicalizationPattern
+    : public OpRewritePattern<tensor::ExtractOp> {
+  using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractOp op,
+                                PatternRewriter& rewriter) const override {
+    auto tieShapeOp = op.tensor().getDefiningOp<disc_shape::TieShapeOp>();
+    if (!tieShapeOp) return failure();
+    rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
+        op, tieShapeOp->getOperand(0), op.indices());
+    return success();
+  }
+};
+
+// Adds shape optimization related patterns.
+void populateShapeOptimizationPatterns(MLIRContext* context,
+                                       RewritePatternSet* patterns) {
+  // clang-format off
+  patterns->insert<
+      DimOfTieShapeOpCanonicalizationPattern,
+      ExtractElementOfTieShapeOpCanonicalizationPattern
+  >(patterns->getContext());
+  // clang-format on
+}
+
+// Adds canonicalization patterns to the list of patterns.
+void addCanonicalizationPatterns(MLIRContext* context,
+                                 RewritePatternSet* patterns) {
+  for (RegisteredOperationName op : context->getRegisteredOperations())
+    op.getCanonicalizationPatterns(*patterns, context);
+}
+
+LogicalResult runCanonicalizer(ModuleOp m, PassPipelineRunner runner) {
+  MLIRContext* context = m.getContext();
+  RewritePatternSet patterns(context);
+  populateShapeOptimizationPatterns(context, &patterns);
+  addCanonicalizationPatterns(context, &patterns);
+
+  SmallVector<std::string> disablePatterns = {
+      "mlir::disc_shape::{anonymous}::IdentityTieShapeOp"};
+  FrozenRewritePatternSet frozenSet(std::move(patterns), disablePatterns);
+
+  if (failed(applyPatternsAndFoldGreedily(m->getRegions(),
+                                          std::move(frozenSet)))) {
+    return m.emitError() << "fail to run canonicalizer\n";
+  }
+
+  OpPassManager dynamicPM("builtin.func");
+  dynamicPM.addPass(createCSEPass());
+  return runner(dynamicPM, m);
+}
+
+// Returns true if the type is possible to be a shape tensor type.
+// Here shape tensor type is defined as follow:
+// - rank-1 static-shaped tensor type
+// - element type of the tensor is int or index
+// - number of elements of the tensor < 32, supposing that the
+//   higiest possible rank is smaller than 32.
+bool isCandidateShapeTensorType(Type ty) {
+  // Try to check if it's a candidate shape tensor.
+  auto tensorTy = ty.dyn_cast<RankedTensorType>();
+
+  return (tensorTy && tensorTy.getRank() == 1 && tensorTy.hasStaticShape() &&
+          tensorTy.getElementType().isIntOrIndex() &&
+          tensorTy.getShape()[0] < 32);
+}
+
+class ShapeComputationIRAnalysis {
+ public:
+  explicit ShapeComputationIRAnalysis(FuncOp func, SymbolicDimMgr& mgr);
+
+  // Analyzes the shape computation IR and the shape constraint IR in the target
+  // module Returns failure if failed.
+  LogicalResult run();
+
+  // Returns the func op this analysis object runs on.
+  FuncOp getFunc() { return funcOp_; }
+
+  // Each dim size value (as one shape operand of one disc_shape.tie_shape) is
+  // mapped to one SymbolicDim after the analysis. This function group all such
+  // dim size values by the SymbolicDim, and return the results.
+  DenseMap<SymbolicDimOp, SmallVector<Value>> getSymbolicDimSSAValueInstance();
+
+  Type getRefinedType(Value value);
+
+  SymbolicDimOp value2SymbolicDimOp(Value value);
+
+ private:
+  LogicalResult runOnRegion(Region* region);
+  LogicalResult runOnBlock(Block* block);
+  LogicalResult runOnOperation(Operation* op);
+
+  LogicalResult buildSymbolicShape(Value value);
+  LogicalResult buildSymbolicShapeForResultsOfOp(Operation* op);
+
+  LogicalResult applyOpConstraint(Operation* op);
+  LogicalResult applyIndexOpConstraint(Operation* op);
+  LogicalResult applyShapeTensorOpConstraint(Operation* op);
+  LogicalResult applyRankedTensorOpConstraint(Operation* op);
+  LogicalResult applyMhloOpConstraint(Operation* op);
+  LogicalResult applyMhloElemOpConstraint(Operation* op);
+  LogicalResult applyMhloDotLikeOpConstraint(Operation* op);
+  LogicalResult applyMhloBcastOpConstraint(Operation* op);
+  LogicalResult applyMhloConcatOpConstraint(Operation* op);
+  LogicalResult applyMhloReshapeLikeOpConstraint(Operation* op);
+
+  LogicalResult applyTieShapeOpConstraint(Operation* op);
+  LogicalResult applyTieShapeOfReshapePatternConstraint(Operation* op);
+  LogicalResult tryToBuildProductEqualityForReshape(
+      Operation* reshapeOp, const SmallVectorImpl<SymbolicDimExpr>& inExprs,
+      int inStartIdx, const SmallVectorImpl<SymbolicDimExpr>& outExprs,
+      int outStartIdx);
+
+  LogicalResult mapRankedValueShapeEqual(Value lhs, Value rhs);
+
+ private:
+  bool initialized_ = false;
+  FuncOp funcOp_;
+  SymbolicDimMgr& mgr_;
+
+  // Map scalar int/index SSA value to a symbolicDim
+  DenseMap<Value, SymbolicDimOp> value2SymDim_;
+
+  // Map a shape tensor value (1D ranked int/index tensor) to an array of
+  // symbolicDims, each for one component of the shape tensor.
+  DenseMap<Value, SmallVector<SymbolicDimOp>> shapeTensor2SymDims_;
+
+  // Map a ranked tensor value to an array of symbolicDims, each represents one
+  // dimension size of the tensor.
+  DenseMap<Value, SmallVector<SymbolicDimOp>> rankedTensor2SymDims_;
+
+  // Map an index value to its defining symbolic expr;
+  DenseMap<Value, SymbolicDimExpr> value2DefiningExpr_;
+};
+
+ShapeComputationIRAnalysis::ShapeComputationIRAnalysis(FuncOp func,
+                                                       SymbolicDimMgr& mgr)
+    : funcOp_(func), mgr_(mgr) {}
+
+LogicalResult ShapeComputationIRAnalysis::run() {
+  // Make sure only run once.
+  if (initialized_) {
+    return funcOp_->emitError()
+           << "re-initialized shape analysis is not supported\n";
+  }
+  initialized_ = true;
+  return runOnRegion(&funcOp_.getBody());
+}
+
+LogicalResult ShapeComputationIRAnalysis::runOnRegion(Region* region) {
+  // Only SCF is supported a.t.m.
+  if (region->getBlocks().size() != 1) {
+    return region->getParentOp()->emitError(
+        "only single block region is supported");
+  }
+  for (Block& block : *region) {
+    if (failed(runOnBlock(&block))) return failure();
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::runOnBlock(Block* block) {
+  // mapping block arguments
+  for (Value value : block->getArguments()) {
+    if (failed(buildSymbolicShape(value))) {
+      return block->getParentOp()->emitError(
+          "failed to build shape for block arg");
+    }
+  }
+
+  // mapping each op inside the block
+  // save a snapshot before visiting in case new ops are inserted during
+  // visiting.
+  SmallVector<Operation*> op_list;
+  for (Operation& op : *block) op_list.push_back(&op);
+  for (Operation* op : op_list) {
+    if (failed(runOnOperation(op))) return failure();
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::runOnOperation(Operation* op) {
+  LLVM_DEBUG(llvm::dbgs() << "runOnOperation: " << *op << "\n");
+  if (failed(buildSymbolicShapeForResultsOfOp(op)))
+    return op->emitError() << "fail to buildSymbolicShapeForResultsOfOp\n";
+
+  // TODO(disc): visit the regions of op once we support funcitonal control
+  // flow.
+
+  // apply op's shape constraint
+  return applyOpConstraint(op);
+}
+
+LogicalResult ShapeComputationIRAnalysis::buildSymbolicShapeForResultsOfOp(
+    Operation* op) {
+  // build shapes for the results of op
+  for (Value result : op->getResults()) {
+    if (failed(buildSymbolicShape(result))) {
+      return op->emitError("failed to build shape for op's result");
+    }
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::buildSymbolicShape(Value value) {
+  Type ty = value.getType();
+  if (ty.isIntOrIndex()) {
+    SymbolicDimOp sym = mgr_.newSymbolicDim();
+    value2SymDim_[value] = sym;
+    value2DefiningExpr_[value] = SymbolicDimExpr(value);
+  } else if (auto tensorTy = ty.dyn_cast<RankedTensorType>()) {
+    SmallVector<SymbolicDimOp> symbols =
+        mgr_.getOrCreateSymbolicDimsForRankedValue(value);
+    rankedTensor2SymDims_[value] = std::move(symbols);
+    // Try to check if it's a candidate shape tensor.
+    if (isCandidateShapeTensorType(ty)) {
+      SmallVector<SymbolicDimOp> symbols;
+      for (int i = 0, d = tensorTy.getShape()[0]; i < d; ++i)
+        symbols.push_back(mgr_.newSymbolicDim());
+      shapeTensor2SymDims_[value] = std::move(symbols);
+    }
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::mapRankedValueShapeEqual(Value lhs,
+                                                                   Value rhs) {
+  auto& lhsDims = rankedTensor2SymDims_[lhs];
+  auto& rhsDims = rankedTensor2SymDims_[rhs];
+
+  if (lhsDims.size() != rhsDims.size()) {
+    return mlir::emitError(lhs.getLoc())
+           << "miss match rank:\n\t"
+           << "lhs = " << lhs << "\n\trhs = " << rhs << "\n";
+  }
+
+  for (const auto& en : llvm::zip(lhsDims, rhsDims)) {
+    if (failed(mgr_.mapSymbolicDimEqual(std::get<0>(en), std::get<1>(en))))
+      return mlir::emitError(lhs.getLoc())
+             << "fail to merge:\n\t"
+             << "lhs = " << lhs << "\n\trhs = " << rhs << "\n";
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyIndexOpConstraint(
+    Operation* op) {
+  if (op->getNumResults() == 0) return success();
+
+  // Only handle scalar int/index op
+  Type ty = op->getResult(0).getType();
+  if (!ty.isIntOrIndex()) return success();
+
+  if (isa<arith::IndexCastOp>(op)) {
+    Value in = op->getOperand(0);
+    Value out = op->getResult(0);
+    if (failed(mgr_.mapSymbolicDimEqual(value2SymDim_[in], value2SymDim_[out])))
+      return op->emitError() << "fail to merge dim\n";
+    value2DefiningExpr_[out] = value2DefiningExpr_[in];
+  } else if (auto dimOp = dyn_cast<tensor::DimOp>(op)) {
+    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
+    if (!dimIndex) return success();
+    value2SymDim_[op->getResult(0)].setKnownNonNegative(true);
+    if (failed(mgr_.mapSymbolicDimEqual(
+            value2SymDim_[op->getResult(0)],
+            rankedTensor2SymDims_[dimOp.source()][*dimIndex])))
+      return op->emitError() << "fail to merge dim\n";
+  } else if (isa<arith::ConstantIndexOp, arith::ConstantIntOp>(op)) {
+    int64_t val = op->getAttrOfType<IntegerAttr>("value").getInt();
+    LLVM_DEBUG(llvm::dbgs() << "applyIndexOpConstraint arith const op val = "
+                            << val << "\n");
+    Value out = op->getResult(0);
+    if (failed(mgr_.mapSymbolicDimEqual(value2SymDim_[out],
+                                        mgr_.newConstantSymbolicDim(val))))
+      return op->emitError() << "fail to merge dim\n";
+    value2DefiningExpr_[out] = SymbolicDimExpr(val, out.getContext());
+  } else if (auto extractOp = dyn_cast<tensor::ExtractOp>(op)) {
+    if (!isCandidateShapeTensorType(extractOp.tensor().getType()))
+      return success();
+    auto indexOp = dyn_cast_or_null<arith::ConstantIndexOp>(
+        op->getOperand(1).getDefiningOp());
+    if (!indexOp) return success();
+    int64_t index = indexOp.getValue().cast<IntegerAttr>().getInt();
+    auto& shapeTensorDims = shapeTensor2SymDims_[extractOp.tensor()];
+    if (index >= shapeTensorDims.size())
+      return op->emitError() << "miss match shape tensor size\n";
+    if (failed(mgr_.mapSymbolicDimEqual(value2SymDim_[op->getResult(0)],
+                                        shapeTensorDims[index])))
+      return op->emitError() << "fail to merge dim\n";
+  } else if (auto mulOp = dyn_cast<arith::MulIOp>(op)) {
+    // TODO: propagete attributes like knownNonNegative/...
+    Value lhs = op->getOperand(0);
+    Value rhs = op->getOperand(1);
+    Value out = op->getResult(0);
+    value2DefiningExpr_[out] = SymbolicDimExpr::buildMulExpr(
+        value2DefiningExpr_[lhs], value2DefiningExpr_[rhs]);
+  }
+
+  // TODO: add support for arith::addi/subi/divi...
+
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyShapeTensorOpConstraint(
+    Operation* op) {
+  if (isa<tensor::FromElementsOp>(op)) {
+    auto& symbols = shapeTensor2SymDims_[op->getResult(0)];
+    if (symbols.size() != op->getOperands().size())
+      op->emitError() << "miss match dim size and num operands\n";
+    for (const auto& z : llvm::zip(op->getOperands(), symbols)) {
+      if (failed(mgr_.mapSymbolicDimEqual(value2SymDim_[std::get<0>(z)],
+                                          std::get<1>(z))))
+        return op->emitError() << "fail to merge dim\n";
+    }
+  } else if (isa<shape::BroadcastOp>(op)) {
+    // TODO(disc): support more than two operands
+    if (op->getNumOperands() != 2) return success();
+    auto& lhsDims = shapeTensor2SymDims_[op->getOperand(0)];
+    auto& rhsDims = shapeTensor2SymDims_[op->getOperand(1)];
+    auto& outDims = shapeTensor2SymDims_[op->getResult(0)];
+
+    SymbolicDimOp sizeOneDim = mgr_.newConstantSymbolicDim(1);
+    int lhsRank = static_cast<int>(lhsDims.size());
+    int rhsRank = static_cast<int>(rhsDims.size());
+    int outRank = static_cast<int>(outDims.size());
+    for (int d = outRank - 1; d >= 0; --d) {
+      int reverseD = outRank - 1 - d;
+      SymbolicDimOp lhsSymbol =
+          ((reverseD < lhsRank) ? lhsDims[lhsRank - 1 - reverseD] : sizeOneDim);
+      SymbolicDimOp rhsSymbol =
+          ((reverseD < rhsRank) ? rhsDims[rhsRank - 1 - reverseD] : sizeOneDim);
+      SymbolicDimOp outSymbol = outDims[d];
+
+      auto getRoot = [&](SymbolicDimOp sym) {
+        return mgr_.getRootSymbolicDim(sym);
+      };
+
+      if (getRoot(lhsSymbol) == getRoot(rhsSymbol))
+        if (failed(mgr_.mapSymbolicDimEqual(lhsSymbol, outSymbol)))
+          return op->emitError() << "fail to merge dim\n";
+      if (getRoot(lhsSymbol) == getRoot(sizeOneDim))
+        if (failed(mgr_.mapSymbolicDimEqual(rhsSymbol, outSymbol)))
+          return op->emitError() << "fail to merge dim\n";
+      if (getRoot(rhsSymbol) == getRoot(sizeOneDim))
+        if (failed(mgr_.mapSymbolicDimEqual(lhsSymbol, outSymbol)))
+          return op->emitError() << "fail to merge dim\n";
+    }
+  }
+
+  // TODO: add support for arith::addi/subi/...
+
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyMhloElemOpConstraint(
+    Operation* op) {
+  if (op->hasTrait<mlir::OpTrait::SameOperandsAndResultType>() ||
+      op->hasTrait<mlir::OpTrait::SameOperandsAndResultShape>()) {
+    Value ref;
+    if (op->getNumOperands() > 0) ref = op->getOperands().front();
+    if (op->getNumResults() > 0) ref = op->getResults().front();
+    if (!ref) return success();
+    for (Value operand : op->getOperands())
+      if (failed(mapRankedValueShapeEqual(operand, ref)))
+        return op->emitError()
+               << "fail to merge symbolic dim of operand of element op\n";
+    for (Value result : op->getResults())
+      if (failed(mapRankedValueShapeEqual(result, ref)))
+        return op->emitError()
+               << "fail to merge symbolic dim of result of element op\n";
+  } else if (op->hasTrait<mlir::OpTrait::SameTypeOperands>() ||
+             op->hasTrait<mlir::OpTrait::SameOperandsShape>()) {
+    if (op->getNumOperands() == 0) return success();
+    Value ref = op->getOperands().front();
+    for (Value operand : op->getOperands())
+      if (failed(mapRankedValueShapeEqual(operand, ref)))
+        return op->emitError()
+               << "fail to merge symbolic dim between operands of element op\n";
+  } else if (auto clamp = dyn_cast<mhlo::ClampOp>(op)) {
+    auto operandTy = clamp.operand().getType().dyn_cast<RankedTensorType>();
+    auto minTy = clamp.min().getType().dyn_cast<RankedTensorType>();
+    auto maxTy = clamp.max().getType().dyn_cast<RankedTensorType>();
+    if (!operandTy || !minTy || !maxTy) return success();
+
+    if (minTy.getRank() != 0) {
+      if (failed(mapRankedValueShapeEqual(clamp.operand(), clamp.min())))
+        return op->emitError()
+               << "fail to merge the symbolic dim of operand and "
+                  "min of mhlo::ClampOp\n";
+    }
+    if (maxTy.getRank() != 0) {
+      if (failed(mapRankedValueShapeEqual(clamp.operand(), clamp.max())))
+        return op->emitError()
+               << "fail to merge the symbolic dim of operand and "
+                  "max of mhlo::ClampOp\n";
+    }
+  } else if (auto select = dyn_cast<mhlo::SelectOp>(op)) {
+    auto predTy = select.pred().getType().dyn_cast<RankedTensorType>();
+    auto trueTy = select.on_true().getType().dyn_cast<RankedTensorType>();
+    auto falseTy = select.on_false().getType().dyn_cast<RankedTensorType>();
+    auto resultTy = select.getResult().getType().dyn_cast<RankedTensorType>();
+    if (!predTy || !trueTy || !falseTy || !resultTy) return success();
+
+    if (predTy.getRank() != 0) {
+      if (failed(mapRankedValueShapeEqual(select.pred(), select.getResult())))
+        return op->emitError() << "fail to merge the symbolic dim of pred and "
+                                  "result of mhlo::SelectOp\n";
+    }
+    if (trueTy.getRank() != 0) {
+      if (failed(
+              mapRankedValueShapeEqual(select.on_true(), select.getResult())))
+        return op->emitError()
+               << "fail to merge the symbolic dim of on_true and "
+                  "result of mhlo::SelectOp\n";
+    }
+    if (falseTy.getRank() != 0) {
+      if (failed(
+              mapRankedValueShapeEqual(select.on_false(), select.getResult())))
+        return op->emitError()
+               << "fail to merge the symbolic dim of on_false and "
+                  "result of mhlo::SelectOp\n";
+    }
+  }
+
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyMhloDotLikeOpConstraint(
+    Operation* op) {
+  if (auto dot_general = dyn_cast<mhlo::DotGeneralOp>(op)) {
+    Value lhs = op->getOperand(0);
+    Value rhs = op->getOperand(1);
+    auto lhsTy = lhs.getType().dyn_cast<RankedTensorType>();
+    auto rhsTy = rhs.getType().dyn_cast<RankedTensorType>();
+    if (!lhsTy || !rhsTy) return success();
+    auto& lhsDims = rankedTensor2SymDims_[lhs];
+    auto& rhsDims = rankedTensor2SymDims_[rhs];
+    if (lhsTy.getRank() != lhsDims.size() || rhsTy.getRank() != rhsDims.size())
+      return op->emitError("lhs or rhs mismatch rank\n");
+    auto dim_numbers = dot_general.dot_dimension_numbers();
+    // Contracting dimensions.
+    auto lhs_contracting_dims = dim_numbers.getLhsContractingDimensions();
+    auto rhs_contracting_dims = dim_numbers.getRhsContractingDimensions();
+    assert(lhs_contracting_dims.size() == rhs_contracting_dims.size());
+    for (int64_t i = 0; i < lhs_contracting_dims.size(); i++) {
+      int64_t lhs_dim = lhs_contracting_dims[i];
+      int64_t rhs_dim = rhs_contracting_dims[i];
+      if (failed(mgr_.mapSymbolicDimEqual(lhsDims[lhs_dim], rhsDims[rhs_dim])))
+        return op->emitError() << "fail to merge dim\n";
+    }
+    // Batching dimensions.
+    auto lhs_batching_dims = dim_numbers.getLhsBatchingDimensions();
+    auto rhs_batching_dims = dim_numbers.getRhsBatchingDimensions();
+    assert(lhs_batching_dims.size() == rhs_batching_dims.size());
+    for (int64_t i = 0; i < lhs_batching_dims.size(); i++) {
+      int64_t lhs_dim = lhs_batching_dims[i];
+      int64_t rhs_dim = rhs_batching_dims[i];
+      if (failed(mgr_.mapSymbolicDimEqual(lhsDims[lhs_dim], rhsDims[rhs_dim])))
+        return op->emitError() << "fail to merge dim\n";
+    }
+  } else if (auto dot = dyn_cast<mhlo::DotOp>(op)) {
+    Value lhs = op->getOperand(0);
+    Value rhs = op->getOperand(1);
+    auto lhsTy = lhs.getType().dyn_cast<RankedTensorType>();
+    auto rhsTy = rhs.getType().dyn_cast<RankedTensorType>();
+    if (!lhsTy || !rhsTy) return success();
+    auto& lhsDims = rankedTensor2SymDims_[lhs];
+    auto& rhsDims = rankedTensor2SymDims_[rhs];
+    if (lhsTy.getRank() != lhsDims.size() || rhsTy.getRank() != rhsDims.size())
+      return op->emitError("lhs or rhs mismatch rank\n");
+
+    if (failed(mgr_.mapSymbolicDimEqual(lhsDims[1], rhsDims[0])))
+      return op->emitError() << "fail to merge dim\n";
+  } else if (auto einsum = dyn_cast<mhlo::EinsumOp>(op)) {
+    auto lhsTy = einsum.lhs().getType().dyn_cast<RankedTensorType>();
+    auto rhsTy = einsum.rhs().getType().dyn_cast<RankedTensorType>();
+    auto outTy = einsum.getResult().getType().dyn_cast<RankedTensorType>();
+    if (!lhsTy || !rhsTy || !outTy) return success();
+
+    auto& lhsDims = rankedTensor2SymDims_[einsum.lhs()];
+    auto& rhsDims = rankedTensor2SymDims_[einsum.rhs()];
+    auto& outDims = rankedTensor2SymDims_[einsum.getResult()];
+
+    StringRef equation = einsum.einsum_config();
+    llvm::SmallDenseMap<char, llvm::SmallDenseMap<EquationVariable, size_t>>
+        all_tokens;
+    if (!parseEinsumEquation(equation, all_tokens, nullptr, nullptr, nullptr)) {
+      return einsum.emitError("unexpected character in einsum equation");
+    }
+    for (auto token : all_tokens) {
+      SmallVector<int64_t> equalValues(3, -1);
+      for (auto item : token.second) {
+        if (item.first == kIsLhs) {
+          equalValues[0] = item.second;
+        } else if (item.first == kIsRhs) {
+          equalValues[1] = item.second;
+        } else {
+          // kIsResult
+          equalValues[2] = item.second;
+        }
+      }
+      if (equalValues[0] >= 0 && equalValues[1] >= 0) {
+        int64_t lhsIdx = equalValues[0];
+        int64_t rhsIdx = equalValues[1];
+        if (lhsIdx >= lhsDims.size() || rhsIdx >= rhsDims.size()) {
+          return op->emitError("lhs or rhs mismatch rank\n");
+        }
+        if (failed(mgr_.mapSymbolicDimEqual(lhsDims[lhsIdx], rhsDims[rhsIdx])))
+          return op->emitError() << "fail to merge dim\n";
+      }
+      if (equalValues[0] >= 0 && equalValues[2] >= 0) {
+        int64_t lhsIdx = equalValues[0];
+        int64_t outIdx = equalValues[2];
+        if (lhsIdx >= lhsDims.size() || outIdx >= outDims.size()) {
+          return op->emitError("lhs or rhs mismatch rank\n");
+        }
+        if (failed(mgr_.mapSymbolicDimEqual(lhsDims[lhsIdx], outDims[outIdx])))
+          return op->emitError() << "fail to merge dim\n";
+      }
+      if (equalValues[1] >= 0 && equalValues[2] >= 0) {
+        int64_t rhsIdx = equalValues[1];
+        int64_t outIdx = equalValues[2];
+        if (rhsIdx >= rhsDims.size() || outIdx >= outDims.size()) {
+          return op->emitError("lhs or rhs mismatch rank\n");
+        }
+        if (failed(mgr_.mapSymbolicDimEqual(rhsDims[rhsIdx], outDims[outIdx])))
+          return op->emitError() << "fail to merge dim\n";
+      }
+    }
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyMhloBcastOpConstraint(
+    Operation* op) {
+  if (isa<mhlo::DynamicBroadcastInDimOp>(op)) {
+    auto& shapeTensorDims = shapeTensor2SymDims_[op->getOperand(1)];
+    auto& outDims = rankedTensor2SymDims_[op->getResult(0)];
+
+    if (shapeTensorDims.size() != outDims.size())
+      return op->emitError() << "mismatch out rank and shape tensor size\n";
+    for (const auto& z : llvm::zip(shapeTensorDims, outDims)) {
+      if (failed(mgr_.mapSymbolicDimEqual(std::get<0>(z), std::get<1>(z))))
+        return op->emitError() << "fail to merge dim\n";
+    }
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyMhloConcatOpConstraint(
+    Operation* op) {
+  if (auto concat = dyn_cast<mhlo::ConcatenateOp>(op)) {
+    Value out = op->getResult(0);
+    int64_t axis = concat.dimension();
+    auto ty = out.getType().dyn_cast<RankedTensorType>();
+    if (!ty) return success();
+    auto& outDims = rankedTensor2SymDims_[out];
+    if (ty.getRank() != outDims.size())
+      return op->emitError() << "output rank mismatch\n";
+    for (Value operand : op->getOperands()) {
+      auto& inDims = rankedTensor2SymDims_[operand];
+      if (inDims.size() != outDims.size())
+        return op->emitError() << "input and output rank mismatch\n";
+      for (int64_t i = 0; i < ty.getRank(); ++i) {
+        if (i == axis) continue;
+        if (failed(mgr_.mapSymbolicDimEqual(inDims[i], outDims[i])))
+          return op->emitError() << "fail to merge dim\n";
+      }
+    }
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyMhloReshapeLikeOpConstraint(
+    Operation* op) {
+  if (auto dynReshape = dyn_cast<mhlo::DynamicReshapeOp>(op)) {
+    Value in = op->getOperand(0);
+    Value targetShape = op->getOperand(1);
+    Value out = op->getResult(0);
+    auto inTy = in.getType().dyn_cast<RankedTensorType>();
+    auto outTy = out.getType().dyn_cast<RankedTensorType>();
+    if (!inTy || !outTy) return success();
+
+    auto& shapeTensorDims = shapeTensor2SymDims_[targetShape];
+    auto& outDims = rankedTensor2SymDims_[out];
+    if (shapeTensorDims.size() != outDims.size())
+      return op->emitError() << "mismatch out rank and shape tensor size\n";
+    for (const auto& z : llvm::zip(shapeTensorDims, outDims)) {
+      if (failed(mgr_.mapSymbolicDimEqual(std::get<0>(z), std::get<1>(z))))
+        return op->emitError() << "fail to merge dim\n";
+    }
+    if (failed(mgr_.mapSymbolicDimProductEqual(
+            SymbolicDimProduct{outDims},
+            SymbolicDimProduct{rankedTensor2SymDims_[in]})))
+      return op->emitError() << "fail to map product equal between the operand "
+                                "and result of reshape op\n";
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyMhloOpConstraint(Operation* op) {
+  if (failed(applyMhloElemOpConstraint(op))) return failure();
+
+  if (failed(applyMhloDotLikeOpConstraint(op))) return failure();
+
+  if (failed(applyMhloBcastOpConstraint(op))) return failure();
+
+  if (failed(applyMhloConcatOpConstraint(op))) return failure();
+
+  if (failed(applyMhloReshapeLikeOpConstraint(op))) return failure();
+
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyRankedTensorOpConstraint(
+    Operation* op) {
+  if (isa<tensor::CastOp>(op)) {
+    auto inTy = op->getOperand(0).getType().dyn_cast<RankedTensorType>();
+    auto outTy = op->getResult(0).getType().dyn_cast<RankedTensorType>();
+    if (!inTy || !outTy) return success();
+    if (failed(mapRankedValueShapeEqual(op->getOperand(0), op->getResult(0))))
+      return op->emitError() << "fail to merge the symbolic dim of operand and "
+                                "result of tensor::CastOp\n";
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::tryToBuildProductEqualityForReshape(
+    Operation* reshapeOp, const SmallVectorImpl<SymbolicDimExpr>& inExprs,
+    int inStartIdx, const SmallVectorImpl<SymbolicDimExpr>& outExprs,
+    int outStartIdx) {
+  SmallVector<SymbolicDimExpr> subInExprs, subOutExprs;
+  SmallVector<SymbolicDimOp> subInDims, subOutDims;
+  auto& inDims = rankedTensor2SymDims_[reshapeOp->getOperand(0)];
+  auto& outDims = rankedTensor2SymDims_[reshapeOp->getResult(0)];
+  auto MergeExprs = [&](const SmallVector<SymbolicDimExpr>& exprs) {
+    SymbolicDimExpr result = SymbolicDimExpr(1, reshapeOp->getContext());
+    for (const auto& expr : exprs)
+      result = SymbolicDimExpr::buildMulExpr(result, expr);
+    return result;
+  };
+  for (int outIdx = outStartIdx; outIdx < outExprs.size(); ++outIdx) {
+    subOutExprs.push_back(outExprs[outIdx]);
+    subOutDims.push_back(outDims[outIdx]);
+    subInExprs.clear();
+    subInDims.clear();
+    SymbolicDimExpr outProduct = MergeExprs(subOutExprs);
+    for (int inIdx = inStartIdx; inIdx < inExprs.size(); ++inIdx) {
+      subInExprs.push_back(inExprs[inIdx]);
+      subInDims.push_back(inDims[inIdx]);
+      SymbolicDimExpr inProduct = MergeExprs(subInExprs);
+      if (SymbolicDimExpr::isEqual(outProduct, inProduct)) {
+        if (failed(mgr_.mapSymbolicDimProductEqual(
+                SymbolicDimProduct{subOutDims},
+                SymbolicDimProduct{subInDims}))) {
+          return reshapeOp->emitError()
+                 << "fail to map partial product equal between the operand "
+                    "and result of reshape op\n";
+        }
+        return tryToBuildProductEqualityForReshape(
+            reshapeOp, inExprs, inIdx + 1, outExprs, outIdx + 1);
+      }
+    }
+  }
+  return success();
+}
+
+// match:
+//   %0 = disc_shape.tie_shape %in, %in_d0, %in_d1, ...
+//   %out = mhlo.dynamic_reshape(%0, %target_shape) ...
+//   %1 = disc_shape.tie_shape %out, %out_d0, %out_d1, ...
+// and try to build product-equality between (%in_d0, %in_d1, ...) and
+// (%out_d0, %out_d1, ...).
+LogicalResult
+ShapeComputationIRAnalysis::applyTieShapeOfReshapePatternConstraint(
+    Operation* op) {
+  auto outTieShape = dyn_cast<disc_shape::TieShapeOp>(op);
+  if (!outTieShape) return success();
+  auto reshapeOp = dyn_cast_or_null<mhlo::DynamicReshapeOp>(
+      outTieShape->getOperand(0).getDefiningOp());
+  if (!reshapeOp) return success();
+  auto inTieShape = dyn_cast_or_null<disc_shape::TieShapeOp>(
+      reshapeOp->getOperand(0).getDefiningOp());
+  if (!inTieShape) return success();
+  SmallVector<SymbolicDimExpr> inExprs, outExprs;
+  llvm::for_each(inTieShape->getOperands().drop_front(),
+                 [&](Value v) { inExprs.push_back(value2DefiningExpr_[v]); });
+  llvm::for_each(outTieShape->getOperands().drop_front(),
+                 [&](Value v) { outExprs.push_back(value2DefiningExpr_[v]); });
+
+  return tryToBuildProductEqualityForReshape(reshapeOp, inExprs, 0, outExprs,
+                                             0);
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyTieShapeOpConstraint(
+    Operation* op) {
+  // supppose:
+  //   %out = disc_shape.tie_shape(%in, %d0, %d1, ...)
+  // we have
+  //   - %in and %out should have same shape
+  //   - the shape of %out == [d0, d1, ...]
+  if (auto tieShape = dyn_cast<disc_shape::TieShapeOp>(op)) {
+    if (failed(mapRankedValueShapeEqual(op->getOperand(0), op->getResult(0))))
+      return op->emitError()
+             << "fail to map the shape of input and output of tie shape op\n";
+
+    auto& resultDims = rankedTensor2SymDims_[op->getResult(0)];
+    if (resultDims.size() + 1 != op->getNumOperands())
+      return op->emitError()
+             << "miss match number shape operand and the rank of result\n";
+    for (const auto& en : llvm::enumerate(op->getOperands().drop_front())) {
+      if (failed(mgr_.mapSymbolicDimEqual(value2SymDim_[en.value()],
+                                          resultDims[en.index()])))
+        return op->emitError() << "fail to merge symbolic dim\n";
+    }
+
+    if (isCandidateShapeTensorType(op->getResult(0).getType())) {
+      auto& lhsShapeTensorDims = shapeTensor2SymDims_[op->getOperand(0)];
+      auto& rhsShapeTensorDims = shapeTensor2SymDims_[op->getResult(0)];
+      if (lhsShapeTensorDims.size() != rhsShapeTensorDims.size())
+        return op->emitError()
+               << "miss match input/output dim size for tie_shape op\n";
+      for (const auto& z : llvm::zip(lhsShapeTensorDims, rhsShapeTensorDims)) {
+        if (failed(mgr_.mapSymbolicDimEqual(std::get<0>(z), std::get<1>(z))))
+          return op->emitError() << "fail to merge dim\n";
+      }
+    }
+
+    if (failed(applyTieShapeOfReshapePatternConstraint(op)))
+      return op->emitError() << "fail to apply tie_shape + reshape like op "
+                                "pattern constraint\n";
+  }
+  return success();
+}
+
+LogicalResult ShapeComputationIRAnalysis::applyOpConstraint(Operation* op) {
+  if (failed(applyIndexOpConstraint(op)))
+    return op->emitError() << "fail to apply constraint for index op\n";
+
+  if (failed(applyShapeTensorOpConstraint(op)))
+    return op->emitError() << "fail to apply constraint for shape tensor op\n";
+
+  if (op->getDialect() == op->getContext()->getLoadedDialect("mhlo") ||
+      op->getDialect() == op->getContext()->getLoadedDialect("mhlo_disc")) {
+    if (failed(applyMhloOpConstraint(op)))
+      return op->emitError() << "fail to apply constraint for mhlo op\n";
+  }
+
+  if (failed(applyRankedTensorOpConstraint(op))) {
+    return op->emitError() << "fail to apply constraint for ranked tensor op\n";
+  }
+
+  if (failed(applyTieShapeOpConstraint(op))) {
+    return op->emitError() << "fail to apply constraint for tie_shape op\n";
+  }
+
+  return success();
+}
+
+DenseMap<SymbolicDimOp, SmallVector<Value>>
+ShapeComputationIRAnalysis::getSymbolicDimSSAValueInstance() {
+  DenseMap<SymbolicDimOp, SmallVector<Value>> instanceMap;
+  funcOp_.walk([&](disc_shape::TieShapeOp tieShapeOp) {
+    Value rankedValue = tieShapeOp->getOperand(0);
+    auto& symbolicDims = rankedTensor2SymDims_[rankedValue];
+    for (auto& en : llvm::enumerate(tieShapeOp->getOperands().drop_front())) {
+      SymbolicDimOp root = mgr_.getRootSymbolicDim(symbolicDims[en.index()]);
+      instanceMap[root].push_back(en.value());
+    }
+  });
+  return instanceMap;
+}
+
+Type ShapeComputationIRAnalysis::getRefinedType(Value value) {
+  auto ty = value.getType().dyn_cast<RankedTensorType>();
+  if (!ty) return value.getType();
+
+  SmallVector<int64_t> newShape;
+  SmallVector<Attribute> refAttrs;
+  bool noDynamicDim = true;
+  for (SymbolicDimOp sym : rankedTensor2SymDims_[value]) {
+    auto root = mgr_.getRootSymbolicDim(sym);
+    newShape.push_back(root.getDimSize());
+    if (newShape.back() == ShapedType::kDynamicSize) noDynamicDim = false;
+    refAttrs.push_back(SymbolRefAttr::get(value.getContext(), root.getName()));
+  }
+
+  if (noDynamicDim) {
+    // static shape ranked tensor type: no needs to add symbolic dim ref attrs.
+    value.setType(RankedTensorType::get(ty.getShape(), ty.getElementType()));
+    return RankedTensorType::get(newShape, ty.getElementType());
+  } else {
+    auto symbolicShapeAttr = ArrayAttr::get(value.getContext(), refAttrs);
+    value.setType(RankedTensorType::get(ty.getShape(), ty.getElementType(),
+                                        symbolicShapeAttr));
+    return RankedTensorType::get(newShape, ty.getElementType(),
+                                 symbolicShapeAttr);
+  }
+}
+
+SymbolicDimOp ShapeComputationIRAnalysis::value2SymbolicDimOp(Value value) {
+  auto it = value2SymDim_.find(value);
+  if (it == value2SymDim_.end()) return {};
+  return mgr_.getRootSymbolicDim(it->second);
+}
+
+DenseMap<Value, Value> buildSymbolDimInstancesDominantMap(
+    DenseMap<SymbolicDimOp, SmallVector<Value>>& instanceMap,
+    DominanceInfo& dominanceInfo) {
+  DenseMap<Value, Value> dominantMap;
+  for (auto& it : instanceMap) {
+    auto& instances = it.second;
+
+    // in normal cases, there should be only one root, aka, dominant value
+    // of all the values of instances
+    SmallVector<Value> roots;
+    for (Value v : instances) {
+      bool is_root = true;
+      for (Value other : instances) {
+        if (v == other) continue;
+        if (dominanceInfo.dominates(other, v.getDefiningOp())) {
+          is_root = false;
+          continue;
+        }
+      }
+      if (is_root) {
+        roots.push_back(v);
+      }
+    }
+    // we should let as much values as possible to be dominated by a same root
+    for (Value root : roots) {
+      for (Value v : instances) {
+        if (dominantMap.find(v) == dominantMap.end() &&
+            dominanceInfo.dominates(root, v.getDefiningOp())) {
+          dominantMap[v] = root;
+        }
+      }
+    }
+  }
+  return dominantMap;
+}
+
+LogicalResult useSameSSAValueIfSymbolicEqual(
+    ShapeComputationIRAnalysis& analysis, bool& changed) {
+  auto instanceMap = analysis.getSymbolicDimSSAValueInstance();
+  DominanceInfo dominanceInfo(analysis.getFunc());
+  auto dominantMap =
+      buildSymbolDimInstancesDominantMap(instanceMap, dominanceInfo);
+
+  for (auto& pair : dominantMap) {
+    Value v = pair.first;
+    Value dominant = pair.second;
+    if (v != dominant) {
+      changed = true;
+      v.replaceAllUsesWith(dominant);
+    }
+  }
+
+  return success();
+}
+
+LogicalResult refineTensorType(ShapeComputationIRAnalysis& analysis,
+                               bool& changed) {
+  auto updateIfNotSame = [&](Value value) {
+    Type refinedTy = analysis.getRefinedType(value);
+    if (refinedTy != value.getType()) {
+      changed = true;
+      value.setType(refinedTy);
+    }
+  };
+
+  FuncOp func = analysis.getFunc();
+
+  // apply refined type for each value
+  func.walk([&](Operation* op) {
+    for (Value operand : op->getOperands()) updateIfNotSame(operand);
+
+    for (Value result : op->getResults()) updateIfNotSame(result);
+  });
+
+  // apply refined function type
+  // 1, collect input types
+  SmallVector<Type, 4> refinedInputTypes;
+  for (Value arg : func.getArguments()) {
+    refinedInputTypes.push_back(analysis.getRefinedType(arg));
+    if (arg.getType() != refinedInputTypes.back()) changed = true;
+  }
+
+  // 2, collect output types
+  SmallVector<Type, 4> refinedOutputTypes;
+  assert(func.getBody().getBlocks().size() == 1);
+  Operation& op = func.getBody().front().getOperations().back();
+  for (Value operand : op.getOperands()) {
+    refinedOutputTypes.push_back(analysis.getRefinedType(operand));
+    if (operand.getType() != refinedOutputTypes.back()) changed = true;
+  }
+
+  // 3, refine function type to new type
+  auto newFuncTy = FunctionType::get(func.getContext(), refinedInputTypes,
+                                     refinedOutputTypes);
+  func.setType(newFuncTy);
+  return success();
+}
+
+LogicalResult tryToSimplifyCompareOp(ShapeComputationIRAnalysis& analysis,
+                                     bool& changed) {
+  FuncOp func = analysis.getFunc();
+
+  SmallVector<arith::CmpIOp> ops;
+  func.walk([&](arith::CmpIOp op) { ops.push_back(op); });
+
+  for (arith::CmpIOp op : ops) {
+    OpBuilder b(op);
+    auto lhsSym = analysis.value2SymbolicDimOp(op.getLhs());
+    auto rhsSym = analysis.value2SymbolicDimOp(op.getRhs());
+    if (!lhsSym || !rhsSym) continue;
+
+    Value pred;
+    Value truePred = b.create<arith::ConstantIntOp>(op.getLoc(), 1, 1);
+    Value falsePred = b.create<arith::ConstantIntOp>(op.getLoc(), 0, 1);
+    if (op.getPredicate() == arith::CmpIPredicate::eq) {
+      if (lhsSym == rhsSym) {
+        pred = truePred;
+      } else if (lhsSym.knownNonNegative() && rhsSym.knownNegativeOne()) {
+        pred = falsePred;
+      } else if (rhsSym.knownNonNegative() && lhsSym.knownNegativeOne()) {
+        pred = falsePred;
+      }
+      // TODO(disc): support other cases
+    } else if (op.getPredicate() == arith::CmpIPredicate::ne) {
+      if (lhsSym == rhsSym) {
+        pred = falsePred;
+      } else if (lhsSym.knownNonNegative() && rhsSym.knownNegativeOne()) {
+        pred = truePred;
+      } else if (rhsSym.knownNonNegative() && lhsSym.knownNegativeOne()) {
+        pred = truePred;
+      }
+      // TODO(disc): support other cases
+    }
+    // TODO(disc): support arith::CmpIPredicate::lt/...
+    if (pred) {
+      op.getResult().replaceAllUsesWith(pred);
+      changed = true;
+    }
+  }
+  return success();
+}
+
+LogicalResult simplifyAccordingToShapeConstraintInfo(
+    ShapeComputationIRAnalysis& analysis, bool& changed) {
+  if (failed(tryToSimplifyCompareOp(analysis, changed))) return failure();
+
+  // TODO(disc): add other possible simplifier here.
+  return success();
+}
+
+LogicalResult applyShapeComputationOptimization(
+    ShapeComputationIRAnalysis& analysis, bool& changed) {
+  // 1, using the same ssa value for all symbolic-equal dim instances.
+  if (failed(useSameSSAValueIfSymbolicEqual(analysis, changed)))
+    return analysis.getFunc()->emitError(
+        "useSameSSAValueIfSymbolicEqual failed\n");
+
+  // 2, After propagation some (partial) known dim size infos, refined
+  // the ranked tensor type.
+  if (failed(refineTensorType(analysis, changed)))
+    return analysis.getFunc()->emitError("refineTensorType failed\n");
+
+  // 3, simplify some expression after propagation shape constraint info.
+  // e.g. if symbolic dim %d is known not negative, then `arith.cmpi eq, %d,
+  // %c-1` could be replaced with a const.
+  if (failed(simplifyAccordingToShapeConstraintInfo(analysis, changed)))
+    return analysis.getFunc()->emitError("fail to simplify\n");
+
+  return success();
+}
+
+// It’s an iterative process and each iteration can be divided into five steps:
+// - Step #1: canonicalization.
+//   including normal canonicalization pattern, cse and a bunch of new rewrite
+//   patterns (e.g. scalarize shape computation IR whenever possible).
+// - Step #2: load existing shape constraint IR
+// - Step #3: global shape computation IR analysis.
+// - Step #4: gloabl shape computation IR optimization.
+// - Step #5: save updated shape constraint IR.
+LogicalResult optimizeShapeComputation(ModuleOp m, FuncOp main,
+                                       PassPipelineRunner runner) {
+  bool changed;
+  do {
+    changed = false;
+    if (failed(runCanonicalizer(m, runner))) {
+      return failure();
+    }
+
+    LLVM_DEBUG(
+        llvm::dbgs()
+        << "Module after runCanonicalizer in optimize-shape-computation:\n"
+        << m << "\n");
+
+    SymbolicDimMgr mgr(m);
+    if (failed(mgr.load())) {
+      return m.emitError() << "fail to load shape constraint IR\n";
+    }
+
+    ShapeComputationIRAnalysis analysis(main, mgr);
+    if (failed(analysis.run())) {
+      return m.emitError() << "fail to analysis shape computation IR\n";
+    }
+
+    if (failed(applyShapeComputationOptimization(analysis, changed))) {
+      return m.emitError() << "fail to optimize shape computation IR\n";
+    }
+
+    LLVM_DEBUG(
+        llvm::dbgs()
+        << "Module after apply-shape-opt in optimize-shape-computation:\n"
+        << m << "\n");
+
+    if (failed(mgr.save())) {
+      return m.emitError() << "fail to save shape constraint IR\n";
+    }
+
+    LLVM_DEBUG(llvm::dbgs()
+               << "Module after save-shape-ir in optimize-shape-computation:\n"
+               << m << "\n");
+  } while (changed);
+
+  if (failed(runCanonicalizer(m, runner))) {
+    return failure();
+  }
+  LLVM_DEBUG(llvm::dbgs() << "Module after optimizeShapeComputation:\n"
+                          << m << "\n");
+  return success();
+}
+
+/////////////////////////// Stage #2 END //////////////////////////////////////
+
+/////////////////////////// Stage #3 BEGIN ////////////////////////////////////
+
+struct ForwardTieShapeOperandToItsConsumers
+    : public OpRewritePattern<disc_shape::TieShapeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(disc_shape::TieShapeOp tieShapeOp,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOp(tieShapeOp, tieShapeOp->getOperand(0));
+    return success();
+  }
+};
+
+LogicalResult cleanUp(ModuleOp m, bool keep_tie_shape) {
+  if (!keep_tie_shape) {
+    RewritePatternSet patterns(m.getContext());
+    patterns.add<ForwardTieShapeOperandToItsConsumers>(patterns.getContext());
+
+    if (failed(applyPatternsAndFoldGreedily(m->getRegions(),
+                                            std::move(patterns)))) {
+      return m.emitError() << "fail to do cleanup\n";
+    }
+  } else {
+    // We will keep a `disc_shape.tie_shape` for each ranked tensor type value.
+    // And thus we can safely: 1, drop all symbolicDim reference attributes in
+    // ranked tensor type to prepare for bufferization (currently bufferizaiton
+    // pass does not support symbolicDim reference). 2, add an attribute to
+    // `disc_shape.tie_shape` op to record symbolicDim reference attribute.
+    // Example. convert from:
+    // ```
+    //   %1 = mhlo.abs(%0) : (tensor<?x?xf32, [@S0, @S1]>) -> tensor<?x?xf32,
+    //   [@S0, @S1]>
+    //   %2 = disc_shape.tie_shape(%1, %d0, %d1) : tensor<?x?xf32, [@S0, @S1]>
+    // ```
+    // to
+    // ```
+    //   %1 = mhlo.abs(%0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+    //   %2 = disc_shape.tie_shape(%1, %d0, %d1) {kDiscSymbolicDimAttr = [@S0,
+    //   @S1]} : tensor<?x?xf32>
+    // ```
+
+    // Part #1: attach the symbolicDim reference attribute to each
+    // `disc_shape.tie_shape` op.
+    SmallVector<disc_shape::TieShapeOp> tieShapeOps;
+    m.walk([&](disc_shape::TieShapeOp op) {
+      auto ty = op->getResult(0).getType().dyn_cast<RankedTensorType>();
+      if (!ty) return;
+      auto attrs = ty.getEncoding().dyn_cast_or_null<ArrayAttr>();
+      if (!attrs) return;
+      op->setAttr(disc_shape::SymbolicDimOp::getSymbolicDimAttrName(), attrs);
+      tieShapeOps.push_back(op);
+    });
+
+    for (disc_shape::TieShapeOp op : tieShapeOps) {
+      Operation* definingOp = op->getOperand(0).getDefiningOp();
+      if (!definingOp) continue;
+      if (definingOp->getBlock() != op->getBlock())
+        return op->emitError(
+            "tie_shape op and the defining op of its source tensor are not in "
+            "the same block\n");
+      // Try to move the definingOp as close to op as possible.
+      SmallVector<Operation*> notUseDefiningOpVec;
+      DenseSet<Value> consumersOfDefiningOpSet{definingOp->getResults().begin(),
+                                               definingOp->getResults().end()};
+      for (auto it = Block::iterator(definingOp), end = Block::iterator(op);
+           it != end; ++it) {
+        if (llvm::any_of(it->getOperands(), [&](Value val) {
+              return consumersOfDefiningOpSet.count(val);
+            })) {
+          for (Value val : it->getResults())
+            consumersOfDefiningOpSet.insert(val);
+        } else {
+          notUseDefiningOpVec.push_back(&*it);
+        }
+      }
+      for (Operation* op : notUseDefiningOpVec) op->moveBefore(definingOp);
+    }
+
+    if (failed(walkRankedTensorValue(
+            m, [&](Value value, RankedTensorType ty, ArrayAttr attrs) {
+              value.setType(
+                  RankedTensorType::get(ty.getShape(), ty.getElementType()));
+              return success();
+            }))) {
+      return failure();
+    }
+    if (failed(updateFunctionType(m))) return failure();
+  }
+  return success();
+}
+
+/////////////////////////// Stage #3 END //////////////////////////////////////
+
+struct DiscShapeOptimizationPass
+    : public DiscShapeOptimizationPassBase<DiscShapeOptimizationPass> {
+  DiscShapeOptimizationPass(const std::string& entry_func_name,
+                            bool keep_tie_shape)
+      : DiscShapeOptimizationPassBase<
+            DiscShapeOptimizationPass>::DiscShapeOptimizationPassBase() {
+    this->entry_func_name_ = entry_func_name;
+    this->keep_tie_shape_ = keep_tie_shape;
+  }
+
+  void runOnOperation() override;
+};
+
+void DiscShapeOptimizationPass::runOnOperation() {
+  ModuleOp m = getOperation();
+  FuncOp main = m.lookupSymbol<FuncOp>(entry_func_name_);
+  if (!main) {
+    m.emitError("entry func: " + entry_func_name_ + " not found");
+    signalPassFailure();
+    return;
+  }
+
+  // Stage #1: Explictily materialize shape computation IR on tensor level
+  if (failed(materializeShapeComputation(m, main))) {
+    signalPassFailure();
+    return;
+  }
+  LLVM_DEBUG(llvm::dbgs() << "Module after materialize shape computation:\n"
+                          << m << "\n");
+
+  // Stage #2: Optimize shape computation IR on tensor level
+  PassPipelineRunner runner = [this](OpPassManager& dynamicPM, ModuleOp m) {
+    return runPipeline(dynamicPM, m);
+  };
+  if (failed(optimizeShapeComputation(m, main, runner))) {
+    signalPassFailure();
+    return;
+  }
+  LLVM_DEBUG(llvm::dbgs() << "Module after shape optimizaiton:\n" << m << "\n");
+
+  // Stage #3: clean up
+  if (failed(cleanUp(m, keep_tie_shape_))) {
+    signalPassFailure();
+    return;
+  }
+  LLVM_DEBUG(llvm::dbgs() << "Module after cleanup:\n" << m << "\n");
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createDiscShapeOptimizationPass(
+    const std::string& entry_func_name, bool keep_tie_shape) {
+  return std::make_unique<DiscShapeOptimizationPass>(entry_func_name,
+                                                     keep_tie_shape);
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
new file mode 100644
index 00000000000..55bc06dc530
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
@@ -0,0 +1,718 @@
+/* Copyright 2022 The BladeDISC Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements the logic to do some shape optimizations on tensor
+// level.
+
+#include "tensorflow/compiler/mlir/disc/transforms/disc_shape_optimization_utils.h"
+
+#include <algorithm>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "tensorflow/compiler/mlir/disc/IR/disc_shape_ops.h"
+#include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
+
+namespace mlir {
+namespace disc_ral {
+
+using ::mlir::func::FuncOp;
+
+// Gives a consistent order of a list op SymbolicDim Ops
+bool compareSymbolicDimOpNames(StringRef lhs, StringRef rhs) {
+  // S -> unknown dimension size at compile time
+  // C -> constant dimension size at compile time
+  // Not known encoding schema, fallback branch
+  if (lhs.size() < 1 || lhs[0] != 'S' && lhs[0] != 'C') return lhs < rhs;
+  if (rhs.size() < 1 || rhs[0] != 'S' && rhs[0] != 'C') return lhs < rhs;
+
+  int64_t lhsIdx, rhsIdx;
+  if (lhs.substr(1).getAsInteger(10, lhsIdx) ||
+      rhs.substr(1).getAsInteger(10, rhsIdx))
+    return lhs < rhs;
+  return (lhs[0] < rhs[0]) || (lhs[0] == rhs[0] && lhsIdx < rhsIdx);
+}
+
+// Gives a consistent order of a list op SymbolicDimProducts
+bool compareSymbolicDimProduct(const SymbolicDimProduct& lhs,
+                               const SymbolicDimProduct& rhs) {
+  if (lhs.symbols.size() < rhs.symbols.size()) return true;
+  if (lhs.symbols.size() == rhs.symbols.size()) {
+    for (const auto& z : llvm::zip(lhs.symbols, rhs.symbols)) {
+      StringRef lhsName = const_cast<SymbolicDimOp&>(std::get<0>(z)).getName();
+      StringRef rhsName = const_cast<SymbolicDimOp&>(std::get<1>(z)).getName();
+      if (compareSymbolicDimOpNames(lhsName, rhsName)) return true;
+      if (lhsName != rhsName) return false;
+    }
+  }
+  return false;
+}
+
+llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
+                              const SymbolicDimProduct& product) {
+  os << "SymbolicDimProduct[\n\tfactor: " << product.factor << ",\n";
+  for (auto& s : product.symbols) os << "\tsymbol: " << s << "\n";
+  os << "]\n";
+  return os;
+}
+
+SymbolicDimMgr::SymbolicDimMgr(ModuleOp m) : m_(m) {}
+
+LogicalResult SymbolicDimMgr::load() {
+  m_.walk([&](disc_shape::SymbolicDimOp op) {
+    symbolDimUnionSet_[op] = op;
+    symbolNameSet_.insert(op.getName().str());
+  });
+  return loadShapeConstraintGraph();
+}
+
+std::string SymbolicDimMgr::getNextName() {
+  std::string name;
+  do {
+    name = (llvm::Twine("S") + llvm::Twine(nextSymbolicOpIdx_++)).str();
+  } while (!symbolNameSet_.insert(name).second);
+  return name;
+}
+
+SymbolicDimOp SymbolicDimMgr::newSymbolicDim() {
+  OpBuilder b(m_);
+  auto symbol = b.create<SymbolicDimOp>(m_.getLoc(), getNextName());
+  symbolDimUnionSet_[symbol] = symbol;
+  m_.push_back(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()))
+             .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)
+    current = symbolDimUnionSet_[current];
+  return current;
+}
+
+bool SymbolicDimMgr::isSymbolicDimEqual(SymbolicDimOp lhs, SymbolicDimOp rhs) {
+  SymbolicDimOp lhsRoot = getRootSymbolicDim(lhs);
+  SymbolicDimOp rhsRoot = getRootSymbolicDim(rhs);
+  return lhsRoot == rhsRoot;
+}
+
+LogicalResult SymbolicDimMgr::mapSymbolicDimEqual(SymbolicDimOp lhs,
+                                                  SymbolicDimOp rhs) {
+  SymbolicDimOp lhsRoot = getRootSymbolicDim(lhs);
+  SymbolicDimOp rhsRoot = getRootSymbolicDim(rhs);
+
+  if (lhsRoot != rhsRoot) {
+    if (compareSymbolicDimOpNames(lhsRoot.getName(), rhsRoot.getName())) {
+      if (failed(lhsRoot.Merge(rhsRoot))) return failure();
+      symbolDimUnionSet_[rhsRoot] = lhsRoot;
+    } else {
+      if (failed(rhsRoot.Merge(lhsRoot))) return failure();
+      symbolDimUnionSet_[lhsRoot] = rhsRoot;
+    }
+  }
+  return success();
+}
+
+// Suppose m and n are non-negative integers.
+int64_t gcd(int64_t m, int64_t n) {
+  if (!m) return n;
+  if (!n) return m;
+  return (m < n) ? gcd(m, n % m) : gcd(m % n, n);
+}
+
+SymbolicDimProduct SymbolicDimMgr::simplifySymbolicDimProduct(
+    const SymbolicDimProduct& x) {
+  SmallVector<SymbolicDimOp> copied;
+  // TODO(disc): expand op if it can be factorized into other symbol dims.
+  for (SymbolicDimOp op : x.symbols) copied.push_back(getRootSymbolicDim(op));
+
+  llvm::sort(copied, [&](SymbolicDimOp lhs, SymbolicDimOp rhs) {
+    return compareSymbolicDimOpNames(lhs.getName(), rhs.getName());
+  });
+  SymbolicDimProduct newX;
+  newX.factor = x.factor;
+  for (SymbolicDimOp op : copied) {
+    if (!op.isDynamic()) {
+      newX.factor *= op.getDimSize();
+    } else {
+      newX.symbols.push_back(op);
+    }
+  }
+  return newX;
+}
+
+std::pair<SymbolicDimProduct, SymbolicDimProduct>
+SymbolicDimMgr::simplifySymbolicDimProductPair(const SymbolicDimProduct& x,
+                                               const SymbolicDimProduct& y) {
+  // First do some basic clean up (e.g. folding const symbolic dim op into the
+  // fator field)
+  auto lhs = simplifySymbolicDimProduct(x);
+  auto rhs = simplifySymbolicDimProduct(y);
+
+  SymbolicDimProduct newLhs, newRhs;
+  int64_t gcdFactor = gcd(std::abs(lhs.factor), std::abs(rhs.factor));
+  // 0 * lhs_symbols = 0 * rhs_symbols, no more information, just return empty
+  // newLhs & newRhs
+  if (!gcdFactor) return std::make_pair(std::move(newLhs), std::move(newRhs));
+
+  // Canonicalization factor form: always let the smaller factor being positive
+  // number.
+  if (std::abs(lhs.factor) < std::abs(rhs.factor)) {
+    if (lhs.factor < 0) gcdFactor = -gcdFactor;
+  } else {
+    if (rhs.factor < 0) gcdFactor = -gcdFactor;
+  }
+
+  newLhs.factor = lhs.factor / gcdFactor;
+  newRhs.factor = rhs.factor / gcdFactor;
+
+  DenseMap<SymbolicDimOp, int> lhsSymbolMap;
+  DenseMap<SymbolicDimOp, int> rhsSymbolMap;
+  for (SymbolicDimOp op : lhs.symbols) ++lhsSymbolMap[op];
+  for (SymbolicDimOp op : rhs.symbols) ++rhsSymbolMap[op];
+
+  for (SymbolicDimOp op : lhs.symbols) {
+    auto it = rhsSymbolMap.find(op);
+    if (it != rhsSymbolMap.end() && op.knownNonSizeZero()) {
+      if (--it->second == 0) rhsSymbolMap.erase(it);
+      continue;
+    }
+    newLhs.symbols.push_back(op);
+  }
+
+  for (SymbolicDimOp op : rhs.symbols) {
+    auto it = lhsSymbolMap.find(op);
+    if (it != lhsSymbolMap.end() && op.knownNonSizeZero()) {
+      if (--it->second == 0) lhsSymbolMap.erase(it);
+      continue;
+    }
+    newRhs.symbols.push_back(op);
+  }
+
+  if (!newLhs.factor) newLhs.symbols.clear();
+  if (!newRhs.factor) newRhs.symbols.clear();
+
+  return std::make_pair(std::move(newLhs), std::move(newRhs));
+}
+
+llvm::Optional<SymbolicDimProduct> SymbolicDimMgr::symbolicDimProductDivide(
+    const SymbolicDimProduct& lhs, const SymbolicDimProduct& rhs) {
+  LLVM_DEBUG(llvm::dbgs() << "Try to check if x % y == 0?\nx = " << lhs
+                          << "y = " << rhs << "\n");
+  SymbolicDimProduct newLhs, newRhs;
+  std::tie(newLhs, newRhs) = simplifySymbolicDimProductPair(lhs, rhs);
+  LLVM_DEBUG(llvm::dbgs() << "Try to check if x % y == 0? after simplify:\nx = "
+                          << newLhs << "y = " << newRhs << "\n");
+
+  // early return if any is zero.
+  if (newLhs.factor == 0 || newRhs.factor == 0) return {};
+  // early return if the const factor is divisible.
+  if (newLhs.factor % newRhs.factor != 0) return {};
+  if (newLhs.symbols.size() < newRhs.symbols.size()) return {};
+
+  SymbolicDimProduct result;
+  result.factor = newLhs.factor / newRhs.factor;
+
+  DenseMap<SymbolicDimOp, int> symProcMap;
+  for (SymbolicDimOp sym : newRhs.symbols) ++symProcMap[sym];
+
+  for (SymbolicDimOp sym : newLhs.symbols) {
+    auto it = symProcMap.find(sym);
+    if (it == symProcMap.end()) {
+      result.symbols.push_back(sym);
+      continue;
+    }
+    if (--it->second == 0) {
+      symProcMap.erase(it);
+      continue;
+    }
+  }
+
+  if (!symProcMap.empty()) return {};
+  LLVM_DEBUG(llvm::dbgs() << "x % y == 0\nx = " << newLhs << "y = " << newRhs
+                          << "x / y = " << result << "\n");
+  return result;
+}
+
+// Try to check if:
+//   lhs = common_factors * lhs'
+//   rhs = common_factors * rhs'
+//   and we already know that product(lhs') == product(rhs')
+bool SymbolicDimMgr::isMultipleOfKnownSymbolicDimProductEqualPair(
+    const SymbolicDimProduct& lhs, const SymbolicDimProduct& rhs) {
+  for (auto& pairOutter : productEqualityMap_) {
+    SymbolicDimProduct& x = pairOutter.first;
+    auto factorX = symbolicDimProductDivide(lhs, x);
+    if (!factorX) continue;
+    for (auto& pairInner : pairOutter.second) {
+      if (!pairInner.second) continue;
+      SymbolicDimProduct& y = pairInner.first;
+      auto factorY = symbolicDimProductDivide(rhs, y);
+      if (!factorY || factorX != factorY) continue;
+      return true;
+    }
+  }
+
+  return false;
+}
+
+bool SymbolicDimMgr::isSymbolicDimProductEqual(const SymbolicDimProduct& lhs,
+                                               const SymbolicDimProduct& rhs) {
+  SymbolicDimProduct newLhs, newRhs;
+  std::tie(newLhs, newRhs) = simplifySymbolicDimProductPair(lhs, rhs);
+
+  // early return for identity case.
+  if (newLhs == newRhs) return true;
+
+  LogicalResult status = updateProductEqualityMap();
+  assert(!failed(status));
+  return isMultipleOfKnownSymbolicDimProductEqualPair(newLhs, newRhs);
+}
+
+LogicalResult SymbolicDimMgr::mapSymbolicDimProductEqual(
+    const SymbolicDimProduct& lhs, const SymbolicDimProduct& rhs) {
+  LLVM_DEBUG(llvm::dbgs() << "Try to map product equal: x = " << lhs
+                          << "\ny = " << rhs << "\n");
+  SymbolicDimProduct newLhs, newRhs;
+  std::tie(newLhs, newRhs) = simplifySymbolicDimProductPair(lhs, rhs);
+  LLVM_DEBUG(llvm::dbgs() << "Try to map product equal after simplify: x = "
+                          << newLhs << "\ny = " << newRhs << "\n");
+
+  // early return for identity case.
+  if (newLhs == newRhs) return success();
+
+  if (newLhs.factor == newRhs.factor && newLhs.symbols.size() == 1 &&
+      newRhs.symbols.size() == 1) {
+    return mapSymbolicDimEqual(newLhs.symbols[0], newRhs.symbols[0]);
+  }
+
+  productEqualityMap_[newLhs][newRhs] = productEqualityMap_[newRhs][newLhs] =
+      true;
+
+  return updateProductEqualityMap();
+}
+
+LogicalResult SymbolicDimMgr::updateProductEqualityMap() {
+  SymbolicDimProductMap newMap;
+  DenseSet<SymbolicDimProduct> productSet;
+  for (auto& pairOutter : productEqualityMap_) {
+    SymbolicDimProduct& x = pairOutter.first;
+    for (auto& pairInner : pairOutter.second) {
+      if (!pairInner.second) continue;
+      SymbolicDimProduct& y = pairInner.first;
+      SymbolicDimProduct newX, newY;
+      std::tie(newX, newY) = simplifySymbolicDimProductPair(x, y);
+      if (newX == newY) continue;
+      newMap[newX][newY] = newMap[newY][newX] = true;
+      productSet.insert(newX);
+      productSet.insert(newY);
+    }
+  }
+
+  bool changed;
+  do {
+    changed = false;
+    for (auto& x : productSet)
+      for (auto& y : productSet)
+        for (auto& z : productSet) {
+          if (x != z && newMap[x][y] && newMap[y][z] && !newMap[x][z]) {
+            newMap[x][z] = newMap[z][x] = true;
+            changed = true;
+          }
+        }
+  } while (changed);
+
+  productEqualityMap_ = std::move(newMap);
+
+  for (auto& x : productSet)
+    for (auto& y : productSet) {
+      if (!productEqualityMap_[x][y]) continue;
+      productEqualityMap_[x][y] = productEqualityMap_[y][x] = false;
+      if (!isMultipleOfKnownSymbolicDimProductEqualPair(x, y)) {
+        productEqualityMap_[x][y] = productEqualityMap_[y][x] = true;
+      }
+    }
+
+  DenseSet<SymbolicDimProduct> toRemove;
+  for (auto& x : productSet) {
+    if (llvm::all_of(productSet, [&](const SymbolicDimProduct& y) {
+          return !productEqualityMap_[x][y];
+        })) {
+      toRemove.insert(x);
+    }
+  }
+  for (auto& x : toRemove) {
+    productEqualityMap_.erase(x);
+    LLVM_DEBUG(llvm::dbgs() << "When updateProductEqualityMap try to remove "
+                               "redundant symbolicDimProduct "
+                            << x << "\n");
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "After updateProductEqualityMap:\n");
+  for (auto& pairOutter : productEqualityMap_) {
+    SymbolicDimProduct& x = pairOutter.first;
+    for (auto& pairInner : pairOutter.second) {
+      if (!pairInner.second) continue;
+      SymbolicDimProduct& y = pairInner.first;
+      LLVM_DEBUG(llvm::dbgs() << "Pair: x = " << x << "\ny = " << y << "\n");
+    }
+  }
+  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();
+      }))) {
+    return failure();
+  }
+
+  // Update function type
+  if (failed(updateFunctionType(m_))) return failure();
+
+  // 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();
+      }))) {
+    return failure();
+  }
+
+  // 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();
+  }
+
+  // canonicalize the name of symbolic dim ops
+  llvm::sort(usedSymbolNames,
+             [&](const std::string& lhs, const std::string& rhs) {
+               return compareSymbolicDimOpNames(lhs, rhs);
+             });
+  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();
+    } else {
+      nameMapping[en.value()] =
+          (llvm::Twine("S") + llvm::Twine(numNonConstDims++)).str();
+    }
+  }
+  for (SymbolicDimOp op : usedSymbolicOps)
+    op.setName(nameMapping[op.getName().str()]);
+
+  // 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();
+      }))) {
+    return failure();
+  }
+
+  // Update function type
+  if (failed(updateFunctionType(m_))) return failure();
+
+  // Save shape constraint graph function
+  return saveShapeConstraintGraph();
+}
+
+/* static */ StringRef SymbolicDimMgr::getShapeConstraintGraphFunctionName() {
+  return "shape_constraint_graph";
+}
+
+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();
+
+  OpBuilder moduleBuilder(m_);
+  auto funcTy = FunctionType::get(m_.getContext(), {}, {});
+  FuncOp func = moduleBuilder.create<FuncOp>(m_.getLoc(), funcName, funcTy);
+  Block* block = func.addEntryBlock();
+  OpBuilder b(block, block->begin());
+  auto returnOp = b.create<func::ReturnOp>(func.getLoc());
+  m_.push_back(func);
+  Location loc = func.getLoc();
+  b.setInsertionPoint(returnOp);
+
+  // save product equal predicate
+  // TODO(disc): return null if the symbolicOps not exists.
+  auto build_operands = [&](const SymbolicDimProduct& prod) {
+    SmallVector<Value> values;
+    if (prod.factor != 1) {
+      values.push_back(b.create<arith::ConstantIndexOp>(loc, prod.factor));
+    }
+    for (SymbolicDimOp sym : prod.symbols) {
+      values.push_back(b.create<disc_shape::DimOp>(
+          loc, b.getIndexType(), FlatSymbolRefAttr::get(sym)));
+    }
+    return values;
+  };
+  SmallVector<SymbolicDimProduct> sortedProductVec;
+  for (auto& p : productEqualityMap_) sortedProductVec.push_back(p.first);
+  llvm::sort(sortedProductVec, compareSymbolicDimProduct);
+  for (auto& x : sortedProductVec) {
+    for (auto& y : sortedProductVec) {
+      if (!compareSymbolicDimProduct(x, y)) continue;
+      if (!productEqualityMap_[x][y]) continue;
+      auto lhsOperands = build_operands(x);
+      auto rhsOperands = build_operands(y);
+      b.create<disc_shape::TieProductEqualOp>(loc, llvm::None, lhsOperands,
+                                              rhsOperands);
+    }
+  }
+  return success();
+}
+
+LogicalResult SymbolicDimMgr::loadShapeConstraintGraph() {
+  StringRef funcName = getShapeConstraintGraphFunctionName();
+  auto func = m_.lookupSymbol<FuncOp>(funcName);
+  if (!func) return success();
+
+  auto build_sym_product = [&](ValueRange range,
+                               SymbolicDimProduct& product) -> LogicalResult {
+    for (Value v : range) {
+      auto definingOp = v.getDefiningOp();
+      if (auto constOp = dyn_cast_or_null<arith::ConstantIndexOp>(definingOp)) {
+        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());
+        if (!sym) return dimOp->emitError() << "fail to find symbolic dim op\n";
+        product.symbols.push_back(sym);
+        continue;
+      }
+      return failure();
+    }
+    return success();
+  };
+  if (func.walk([&](disc_shape::TieProductEqualOp op) {
+            SymbolicDimProduct lhs, rhs;
+            if (failed(build_sym_product(op.lhs(), lhs)) ||
+                failed(build_sym_product(op.rhs(), rhs)) ||
+                failed(mapSymbolicDimProductEqual(lhs, rhs)))
+              return WalkResult::interrupt();
+            return WalkResult::advance();
+          })
+          .wasInterrupted()) {
+    return func->emitError() << "fail to load product equal constraint\n";
+  }
+  return success();
+}
+
+SmallVector<SymbolicDimOp>
+SymbolicDimMgr::getOrCreateSymbolicDimsForRankedValue(Value value) {
+  // TODO: load existing symbols from the attribute attached on the tensor type
+  SmallVector<SymbolicDimOp> symbols;
+  auto ty = value.getType().cast<RankedTensorType>();
+  if (auto attrs = value.getType()
+                       .cast<RankedTensorType>()
+                       .getEncoding()
+                       .dyn_cast_or_null<ArrayAttr>()) {
+    for (Attribute attr : attrs) {
+      auto sym = m_.lookupSymbol<SymbolicDimOp>(attr.cast<FlatSymbolRefAttr>());
+      assert(sym);
+      symbols.push_back(sym);
+    }
+    assert(ty.getRank() == symbols.size());
+  } else {
+    for (int64_t dim : ty.getShape()) {
+      symbols.push_back(dim == ShapedType::kDynamicSize
+                            ? newSymbolicDim()
+                            : newConstantSymbolicDim(dim));
+    }
+  }
+
+  return symbols;
+}
+
+llvm::Optional<SmallVector<FlatSymbolRefAttr>> getRankedValueSymbolicDimRefs(
+    Value value) {
+  auto ty = value.getType().dyn_cast<RankedTensorType>();
+  if (!ty) return {};
+  auto attrs = ty.getEncoding().dyn_cast_or_null<ArrayAttr>();
+  if (!attrs) return {};
+  if (attrs.size() != ty.getRank()) return {};
+  SmallVector<FlatSymbolRefAttr> symbols;
+  for (const auto& attr : attrs) {
+    auto symbol = attr.dyn_cast<FlatSymbolRefAttr>();
+    if (!symbol) return {};
+    symbols.push_back(symbol);
+  }
+  return symbols;
+}
+
+// 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.
+LogicalResult updateFunctionType(FuncOp func) {
+  if (func.getBody().getBlocks().size() == 0) return success();
+  if (func.getBody().getBlocks().size() > 1)
+    return func->emitError() << "not support multi-block function\n";
+  SmallVector<Type, 4> refinedInputTypes;
+  for (Value arg : func.getBody().front().getArguments()) {
+    refinedInputTypes.push_back(arg.getType());
+  }
+
+  // 2, collect output types
+  SmallVector<Type, 4> refinedOutputTypes;
+  Operation& op = func.getBody().front().getOperations().back();
+  for (Value operand : op.getOperands()) {
+    refinedOutputTypes.push_back(operand.getType());
+  }
+
+  // 3, refine function type to new type
+  auto newFuncTy = FunctionType::get(func.getContext(), refinedInputTypes,
+                                     refinedOutputTypes);
+  func.setType(newFuncTy);
+  return success();
+}
+
+// Walk each ranked tensor type values inside op.
+LogicalResult walkRankedTensorValue(Operation* op, Visitor visitor) {
+  if (op->walk([&](Operation* op) {
+          for (Value value : llvm::to_vector(op->getResults())) {
+            auto ty = value.getType().dyn_cast<RankedTensorType>();
+            if (!ty) continue;
+            auto attrs = ty.getEncoding().dyn_cast_or_null<ArrayAttr>();
+            if (!attrs) continue;
+            if (failed(visitor(value, ty, attrs)))
+              return WalkResult::interrupt();
+          }
+          return WalkResult::advance();
+        }).wasInterrupted()) {
+    return failure();
+  }
+
+  if (op->walk([&](Block* block) {
+          for (Value value : llvm::to_vector((block->getArguments()))) {
+            auto ty = value.getType().dyn_cast<RankedTensorType>();
+            if (!ty) continue;
+            auto attrs = ty.getEncoding().dyn_cast_or_null<ArrayAttr>();
+            if (!attrs) continue;
+            if (failed(visitor(value, ty, attrs)))
+              return WalkResult::interrupt();
+          }
+          return WalkResult::advance();
+        }).wasInterrupted()) {
+    return failure();
+  }
+  return success();
+}
+
+LogicalResult updateFunctionType(Operation* op) {
+  if (op->walk([&](FuncOp func) {
+          if (failed(updateFunctionType(func))) return WalkResult::interrupt();
+          return WalkResult::advance();
+        }).wasInterrupted())
+    return failure();
+  return success();
+}
+
+SymbolicDimExpr::SymbolicDimExpr(Value sym)
+    : expr(getAffineSymbolExpr(0, sym.getContext())), symbols(1, sym) {}
+
+SymbolicDimExpr::SymbolicDimExpr(int64_t val, MLIRContext* context)
+    : expr(getAffineConstantExpr(val, context)) {}
+
+llvm::Optional<int64_t> SymbolicDimExpr::getConstValue() {
+  if (auto cstExpr = expr.dyn_cast<AffineConstantExpr>())
+    return cstExpr.getValue();
+  return {};
+}
+
+template <typename Combiner>
+/* static */ SymbolicDimExpr SymbolicDimExpr::buildBinaryExpr(
+    const SymbolicDimExpr& lhs, const SymbolicDimExpr& rhs,
+    Combiner&& combiner) {
+  SymbolicDimExpr expr;
+  expr.symbols.append(lhs.symbols);
+  expr.symbols.append(rhs.symbols);
+  expr.expr = combiner(
+      lhs.expr, rhs.expr.shiftSymbols(rhs.symbols.size(), lhs.symbols.size()));
+  return expr;
+}
+
+/* static */ SymbolicDimExpr SymbolicDimExpr::buildMulExpr(
+    const SymbolicDimExpr& lhs, const SymbolicDimExpr& rhs) {
+  return buildBinaryExpr(lhs, rhs,
+                         [](AffineExpr a, AffineExpr b) { return a + b; });
+}
+
+/* static */ bool SymbolicDimExpr::isEqual(const SymbolicDimExpr& lhs,
+                                           const SymbolicDimExpr& rhs) {
+  // TODO(disc): canonicalize the order of symbols of lhs and rhs first
+  // TODO(disc): simplify expr of lhs and rhs first
+  return lhs.symbols == rhs.symbols && lhs.expr == rhs.expr;
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.h b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.h
new file mode 100644
index 00000000000..f29e21aad64
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.h
@@ -0,0 +1,266 @@
+// Copyright 2022 The BladeDISC Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <memory>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/AffineExpr.h"
+#include "tensorflow/compiler/mlir/disc/IR/disc_shape_ops.h"
+
+#ifndef TENSORFLOW_COMPILER_MLIR_DISC_TRANSFORMS_DISC_SHAPE_OPTIMIZATION_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_DISC_TRANSFORMS_DISC_SHAPE_OPTIMIZATION_UTILS_H_
+
+namespace mlir {
+namespace disc_ral {
+
+using disc_shape::SymbolicDimOp;
+
+// 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>> getRankedValueSymbolicDimRefs(
+    Value value);
+
+using Visitor = std::function<LogicalResult(Value value, RankedTensorType ty,
+                                            ArrayAttr attrs)>;
+// Walk each ranked tensor type values inside op.
+LogicalResult walkRankedTensorValue(Operation* op, Visitor visitor);
+
+// 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.
+LogicalResult updateFunctionType(func::FuncOp func);
+
+// Updates the type of all functions inside the op.
+LogicalResult updateFunctionType(Operation* op);
+
+// Represents a product of symbolic and concrete factors. This will allow us to
+// prove product equalities symbolically.
+struct SymbolicDimProduct {
+  // List all symbolic factors as they can not be aggregated.
+  llvm::SmallVector<SymbolicDimOp> symbols;
+
+  // Product of all const factors.
+  int64_t factor = 1;
+  bool empty() { return factor == 1 && symbols.empty(); }
+};
+
+// Returns true if two SymbolicDimProduct are equal
+inline bool operator==(const SymbolicDimProduct& lhs,
+                       const SymbolicDimProduct& rhs) {
+  return lhs.factor == rhs.factor && lhs.symbols == rhs.symbols;
+}
+
+// Returns true if two SymbolicDimProduct are not equal
+inline bool operator!=(const SymbolicDimProduct& lhs,
+                       const SymbolicDimProduct& rhs) {
+  return !(lhs == rhs);
+}
+
+// Dumps SymbolicDimProduct to the output stream
+llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
+                              const SymbolicDimProduct& product);
+
+}  // namespace disc_ral
+}  // namespace mlir
+
+namespace llvm {
+
+// To Support use SymbolicDimProduct as key of DenseMap/DenseSet
+template <>
+struct DenseMapInfo<mlir::disc_ral::SymbolicDimProduct> {
+  static inline mlir::disc_ral::SymbolicDimProduct getEmptyKey() {
+    return {
+        llvm::SmallVector<mlir::disc_shape::SymbolicDimOp>{
+            llvm::DenseMapInfo<mlir::disc_shape::SymbolicDimOp>::getEmptyKey()},
+        llvm::DenseMapInfo<int64_t>::getEmptyKey()};
+  }
+  static inline mlir::disc_ral::SymbolicDimProduct getTombstoneKey() {
+    return {
+        llvm::SmallVector<mlir::disc_shape::SymbolicDimOp>{llvm::DenseMapInfo<
+            mlir::disc_shape::SymbolicDimOp>::getTombstoneKey()},
+        llvm::DenseMapInfo<int64_t>::getTombstoneKey()};
+  }
+  static unsigned getHashValue(mlir::disc_ral::SymbolicDimProduct symProd) {
+    unsigned hash =
+        llvm::DenseMapInfo<size_t>::getHashValue(symProd.symbols.size());
+    for (mlir::disc_shape::SymbolicDimOp op : symProd.symbols)
+      hash = llvm::hash_combine(
+          hash,
+          llvm::DenseMapInfo<mlir::disc_shape::SymbolicDimOp>::getHashValue(
+              op));
+    return llvm::hash_combine(
+        hash, llvm::DenseMapInfo<int64_t>::getHashValue(symProd.factor));
+  }
+  static bool isEqual(const mlir::disc_ral::SymbolicDimProduct& lhs,
+                      const mlir::disc_ral::SymbolicDimProduct& rhs) {
+    return lhs == rhs;
+  }
+};
+
+}  // namespace llvm
+
+namespace mlir {
+namespace disc_ral {
+
+// Reprensets a symbolic expression of symbolicDims.
+class SymbolicDimExpr {
+ public:
+  SymbolicDimExpr() = default;
+  explicit SymbolicDimExpr(Value symbol);
+  SymbolicDimExpr(int64_t val, MLIRContext* context);
+
+  explicit operator bool() const { return static_cast<bool>(expr); }
+
+  // Returns null if not a const SymbolicDimExpr, or the const value.
+  llvm::Optional<int64_t> getConstValue();
+
+  static SymbolicDimExpr buildMulExpr(const SymbolicDimExpr& lhs,
+                                      const SymbolicDimExpr& rhs);
+  static bool isEqual(const SymbolicDimExpr& lhs, const SymbolicDimExpr& rhs);
+
+ private:
+  template <typename Combiner>
+  static SymbolicDimExpr buildBinaryExpr(const SymbolicDimExpr& lhs,
+                                         const SymbolicDimExpr& rhs,
+                                         Combiner&& combiner);
+
+ private:
+  // TODO(disc): AffineExpr is known not able to simplify even basic symbolic
+  // expression. For example, `affine_map<(d0, d1) -> (d0 + d1 - d0 - 1)>` can
+  // not be simplified to `affine_map<(d0, d1) -> (d1 - 1)>`
+  //
+  // We may change to a professional symbolic expression tool (e.g. SymPy in
+  // python, or GiNaC in C++) in the future.
+  AffineExpr expr;
+  SmallVector<Value> symbols;
+};
+
+// A class to manage shape-constraint related IR
+class SymbolicDimMgr {
+ public:
+  explicit SymbolicDimMgr(ModuleOp m);
+
+  // Loads pre-defined SymbolicDim ops from the module this mgr runs on.
+  LogicalResult load();
+
+  // Returns a new symbolicDim instance. The returned symbolicDim is owned by
+  // this mgr.
+  SymbolicDimOp newSymbolicDim();
+
+  // Returns a symbolicDim which have static dim size == `val`.
+  SymbolicDimOp newConstantSymbolicDim(int64_t val);
+
+  SmallVector<SymbolicDimOp> getOrCreateSymbolicDimsForRankedValue(Value value);
+
+  // All symbolic-equal dims form a group.
+  // Returns the root SymbolicDim of the symbolic-equal symbolic dim group that
+  // this SymbolicDim belongs to.
+  SymbolicDimOp getRootSymbolicDim(SymbolicDimOp symbol);
+
+  // Returns true if lhs and rhs are known to be equal.
+  bool isSymbolicDimEqual(SymbolicDimOp lhs, SymbolicDimOp rhs);
+
+  // Marks lhs and rhs have same size and try to merge lhs & rhs static known
+  // info. Returns failure if failed to merge lhs & rhs.
+  LogicalResult mapSymbolicDimEqual(SymbolicDimOp lhs, SymbolicDimOp rhs);
+
+  // Returns the simplified version of SymbolicDimProduct.
+  // This will try to fold some symbolicDim ops with const values.
+  SymbolicDimProduct simplifySymbolicDimProduct(const SymbolicDimProduct& x);
+
+  // Returns the simplified version of SymbolicDimProductPair.
+  // This will try to reduce some common symbolic ops if they are known not
+  // zero.
+  std::pair<SymbolicDimProduct, SymbolicDimProduct>
+  simplifySymbolicDimProductPair(const SymbolicDimProduct& x,
+                                 const SymbolicDimProduct& y);
+
+  // Returns null if x is not divided exactly by y, or else the result of x / y
+  // Suppose that all symbols are not zero, thus common symbolic dim factors can
+  // be elimiated safely.
+  // For example:
+  //   x = 6 * symbol_0 * symbol_1 * symbol_2
+  //   y = 3 * symbol_0 * symbol_1
+  //   x / y == 2 * symbol_2 (suppose symbol_0 and symbol_1 are not zero)
+  llvm::Optional<SymbolicDimProduct> symbolicDimProductDivide(
+      const SymbolicDimProduct& x, const SymbolicDimProduct& y);
+
+  // mark group [a0, b0, ...] and group [a1, b1, c1, ...] are group
+  // multiplication equal `a0 * b0 * ... = a1 * b1 * c1 * ...`
+  bool isSymbolicDimProductEqual(const SymbolicDimProduct& lhs,
+                                 const SymbolicDimProduct& rhs);
+
+  // mark `product([a0, b0, ...]) == product([a1, b1, c1, ...])`
+  LogicalResult mapSymbolicDimProductEqual(const SymbolicDimProduct& lhs,
+                                           const SymbolicDimProduct& rhs);
+
+  // Saves the updated shape constraint IR
+  LogicalResult save();
+
+  // Returns the name of the shape constraint graph
+  static StringRef getShapeConstraintGraphFunctionName();
+
+ private:
+  // Returns next unique name for a new SymbolicDim op.
+  std::string getNextName();
+
+  // Simplify the ProductEqualityMap
+  LogicalResult updateProductEqualityMap();
+
+  // The shape constraint graph stores predicates between symbolicDim ops.
+  // Saves/loads shape constraint graph.
+  LogicalResult saveShapeConstraintGraph();
+  LogicalResult loadShapeConstraintGraph();
+
+  // Try to check if:
+  //   lhs = common_factors * lhs'
+  //   rhs = common_factors * rhs'
+  //   and we already know that product(lhs') == product(rhs')
+  bool isMultipleOfKnownSymbolicDimProductEqualPair(
+      const SymbolicDimProduct& lhs, const SymbolicDimProduct& rhs);
+
+ private:
+  // The module this SymbolicDimMgr runs on.
+  ModuleOp m_;
+
+  // A unique id to generate unique name.
+  int64_t nextSymbolicOpIdx_ = 0;
+
+  // Set of name of SymbolicDim.
+  // It stores all the seen names and is used to give a unique name
+  // to new created symbolic dim ops.
+  std::unordered_set<std::string> symbolNameSet_;
+
+  // map a symbolic dim -> its root SymbolicDim
+  // Here root symbolic dim means the representative member in the
+  // symbolic-equal symbolic dim set that this symbolic dim belongs to.
+  DenseMap<SymbolicDimOp, SymbolicDimOp> symbolDimUnionSet_;
+
+  // map a concret constant value to a symbolic dim instance that represents the
+  // constant.
+  // Note that here we do not use DenseMap since it does not support using large
+  // int64 (e.g. 9223372036854775807 generated by torch frontend) as key.
+  std::unordered_map<int64_t, SymbolicDimOp> constantSymbolicDimMap_;
+
+  // m[p1][p2] = true iff product(p1) == product(p2)
+  using SymbolicDimProductMap =
+      DenseMap<SymbolicDimProduct, DenseMap<SymbolicDimProduct, bool>>;
+  SymbolicDimProductMap productEqualityMap_;
+};
+
+}  // namespace disc_ral
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_DISC_TRANSFORMS_DISC_SHAPE_OPTIMIZATION_UTILS_H_
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/transforms/disc_strip_shape_constraint_ops.cc b/tao_compiler/mlir/disc/transforms/disc_strip_shape_constraint_ops.cc
new file mode 100644
index 00000000000..a2064958905
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/disc_strip_shape_constraint_ops.cc
@@ -0,0 +1,56 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic to strip disc shape constraint ops.
+#include "llvm/Support/Debug.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Support/LogicalResult.h"
+#include "tensorflow/compiler/mlir/disc/IR/disc_shape_ops.h"
+#include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
+#include "tensorflow/compiler/mlir/disc/transforms/disc_shape_optimization_utils.h"
+
+namespace mlir {
+namespace disc_ral {
+
+namespace {
+
+struct DiscStripShapeConstraintOpsPass
+    : public DiscStripShapeConstraintOpsPassBase<
+          DiscStripShapeConstraintOpsPass> {
+  void runOnOperation() override {
+    SmallVector<disc_shape::SymbolicDimOp> ops;
+    getOperation().walk(
+        [&](disc_shape::SymbolicDimOp op) { ops.push_back(op); });
+    for (disc_shape::SymbolicDimOp op : ops) {
+      op->erase();
+    }
+    StringRef funcName = SymbolicDimMgr::getShapeConstraintGraphFunctionName();
+    // first try to remove the old shape constraint graph
+    if (auto func = getOperation().lookupSymbol<func::FuncOp>(funcName))
+      func->erase();
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createStripShapeConstraintOpsPass() {
+  return std::make_unique<DiscStripShapeConstraintOpsPass>();
+}
+
+}  // namespace disc_ral
+}  // namespace mlir
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc b/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc
index 5d9955421dd..5fc3e646ac4 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc
@@ -238,11 +238,20 @@ RankedTensorType getTransposeOutputType(
     Value value, const SmallVectorImpl<int64_t>& permutation, OpBuilder& b) {
   // Compute the resulting shape.
   llvm::SmallVector<int64_t, 4> transposed_shape;
-  ShapedType input_type = value.getType().cast<ShapedType>();
+  auto input_type = value.getType().cast<RankedTensorType>();
   auto input_shape = input_type.getShape();
   for (int64_t val : permutation) {
     transposed_shape.push_back(input_shape[val]);
   }
+  if (auto attrs = input_type.getEncoding().dyn_cast_or_null<ArrayAttr>()) {
+    SmallVector<Attribute> newAttrs;
+    for (int64_t val : permutation) {
+      newAttrs.push_back(attrs[val]);
+    }
+    auto symbolicShapeAttr = ArrayAttr::get(value.getContext(), newAttrs);
+    return RankedTensorType::get(transposed_shape, input_type.getElementType(),
+                                 symbolicShapeAttr);
+  }
   return RankedTensorType::get(transposed_shape, input_type.getElementType());
 }
 
diff --git a/tao_compiler/mlir/disc/transforms/dot_rewriter.cc b/tao_compiler/mlir/disc/transforms/dot_rewriter.cc
index 700841a6a11..8269c4d5797 100644
--- a/tao_compiler/mlir/disc/transforms/dot_rewriter.cc
+++ b/tao_compiler/mlir/disc/transforms/dot_rewriter.cc
@@ -293,6 +293,7 @@ struct EinsumToDotGeneralPattern : public OpRewritePattern<mhlo::EinsumOp> {
         batching_tokens, lhs_non_contracting_tokens, rhs_non_contracting_tokens,
         lhs_original_tokens, rhs_original_tokens, result_original_tokens,
         result_reordered_tokens);
+    result.setType(op.getResult().getType());
 
     rewriter.replaceOp(op, result);
     return success();
diff --git a/tao_compiler/mlir/disc/transforms/element_type_converter.cc b/tao_compiler/mlir/disc/transforms/element_type_converter.cc
index ff6f5b4a8c7..a3feefc50b9 100644
--- a/tao_compiler/mlir/disc/transforms/element_type_converter.cc
+++ b/tao_compiler/mlir/disc/transforms/element_type_converter.cc
@@ -147,9 +147,10 @@ struct ConvertReduceOpWithSmallWidthIntType
     assert(new_op && "convert element type of reduce op failed");
 
     SmallVector<Value, 4> converted_results;
-    for (auto value : new_op.getResults()) {
-      converted_results.push_back(
-          rewriter.create<mhlo::ConvertOp>(loc, value, ty.getElementType()));
+    for (const auto& z : llvm::zip(new_op.getResults(), op.getResults())) {
+      converted_results.push_back(rewriter.create<mhlo::ConvertOp>(
+          loc, std::get<0>(z), ty.getElementType()));
+      converted_results.back().setType(std::get<1>(z).getType());
     }
 
     rewriter.replaceOp(op, converted_results);
@@ -188,6 +189,7 @@ struct ConvertConvOp : public OpRewritePattern<OpTy> {
     Value conv =
         rewriter.create<OpTy>(loc, f16_tensor_ty, newOperands, op->getAttrs());
     Value fp32_conv = rewriter.create<mhlo::ConvertOp>(loc, conv, f32_ty);
+    fp32_conv.setType(op.getResult().getType());
     rewriter.replaceOp(op, fp32_conv);
     return success();
   }
@@ -221,6 +223,7 @@ struct ConvertDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
         loc, f16_tensor_ty, lhs_f16, rhs_f16, op.dot_dimension_numbers(),
         nullptr);
     Value fp32_dot = rewriter.create<mhlo::ConvertOp>(loc, dot, f32_ty);
+    fp32_dot.setType(op.getResult().getType());
     rewriter.replaceOp(op, fp32_dot);
     return success();
   }
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
index b2b7eaacb1d..b8759511621 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "mlir/Pass/Pass.h"               // TF:local_config_mlir
 #include "mlir/Transforms/RegionUtils.h"  // TF:llvm-project
 #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"
 #include "tensorflow/compiler/mlir/disc/transforms/placement_utils.h"
 #include "tensorflow/compiler/mlir/disc/transforms/shape_utils.h"
@@ -501,6 +502,10 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
   void runOnOperation() override {
     func::FuncOp func = getOperation();
 
+    // skip shape constraint graph
+    if (func.getName() == SymbolicDimMgr::getShapeConstraintGraphFunctionName())
+      return;
+
     // collect all blocks inside the function.
     SmallVector<Block*, 4> blocks;
     CollectBlocksInsideFunction(func, blocks);
@@ -534,7 +539,9 @@ struct DiscFusionPass : public DiscFusionPassBase<DiscFusionPass> {
                                     &disc_expected_kernels_in_ut);
     if ((disc_expected_kernels_in_ut >= 0) &&
         (disc_expected_kernels_in_ut != fusion_pattern_number)) {
-      emitError(func.getLoc(), "fusion pattern number is not as expected.");
+      func->emitError() << "fusion pattern number is not as expected ("
+                        << disc_expected_kernels_in_ut << " vs "
+                        << fusion_pattern_number << ")\n";
       signalPassFailure();
       return;
     }
diff --git a/tao_compiler/mlir/disc/transforms/passes.h b/tao_compiler/mlir/disc/transforms/passes.h
index cf49bc063c9..08071b0ac00 100644
--- a/tao_compiler/mlir/disc/transforms/passes.h
+++ b/tao_compiler/mlir/disc/transforms/passes.h
@@ -208,6 +208,13 @@ std::unique_ptr<OperationPass<FuncOp>> createDiscDotMergePass();
 // Apply some basic algebra simplification optimizations.
 std::unique_ptr<OperationPass<FuncOp>> createDiscAlgebraSimplifierPass();
 
+// apply some shape-related optimization.
+std::unique_ptr<OperationPass<ModuleOp>> createDiscShapeOptimizationPass(
+    const std::string& entry_func_name = "main", bool keep_tie_shape = false);
+
+// strip disc shape constraint related ops.
+std::unique_ptr<OperationPass<ModuleOp>> createStripShapeConstraintOpsPass();
+
 }  // namespace disc_ral
 }  // namespace mlir
 


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/mlir_feature_test.cc b/tao_compiler/mlir/disc/tests/mlir_feature_test.cc
index e496a5e09e1..53f7ac98a6f 100644
--- a/tao_compiler/mlir/disc/tests/mlir_feature_test.cc
+++ b/tao_compiler/mlir/disc/tests/mlir_feature_test.cc
@@ -223,6 +223,58 @@ bool feature_test_read_input_from_file(const std::string& mlir_file_path,
                            input_vals, profiling);
 }
 
+// key -> (value, pre-exist?)
+using EnvSetting =
+    std::unordered_map<std::string, std::pair<std::string, bool>>;
+using EnvSettings = std::vector<EnvSetting>;
+
+void addBoolFlags(EnvSettings& envSettings, const std::string& key) {
+  char* value = getenv(key.c_str());
+  if (value) {
+    for (auto& setting : envSettings) {
+      setting[key].first = value;
+      setting[key].second = true;
+    }
+  } else {
+    size_t original_size = envSettings.size();
+    for (int i = 0; i < original_size; ++i) {
+      envSettings[i][key].first = "false";
+      envSettings.push_back(envSettings[i]);
+      envSettings[i][key].first = "true";
+    }
+  }
+}
+
+EnvSettings getEnvironmentSettings() {
+  EnvSettings envSettings{{}};
+  addBoolFlags(envSettings, "DISC_ENABLE_STITCH");
+  addBoolFlags(envSettings, "DISC_ENABLE_SHAPE_CONSTRAINT_IR");
+  return envSettings;
+}
+
+struct EnvSettingContext {
+  explicit EnvSettingContext(const EnvSetting& setting) : setting(setting) {
+    VLOG(0) << "Apply env setting:";
+    for (const auto& kv : setting) {
+      VLOG(0) << "\t" << kv.first << " = " << kv.second.first;
+      setenv(kv.first.c_str(), kv.second.first.c_str(), 1);
+    }
+  }
+
+  ~EnvSettingContext() {
+    VLOG(0) << "Unset env setting:";
+    for (const auto& kv : setting) {
+      // not a pre-exist flag, unset it.
+      if (!kv.second.second) {
+        VLOG(0) << "\t" << kv.first << " = " << kv.second.first;
+        unsetenv(kv.first.c_str());
+      }
+    }
+  }
+
+  EnvSetting setting;
+};
+
 bool feature_test_main(const std::string& mlir_file_path,
                        const std::vector<BackendType>& backend_types,
                        int num_inputs, int num_outputs,
@@ -232,17 +284,10 @@ bool feature_test_main(const std::string& mlir_file_path,
                        bool profiling, bool multi_cc_mode,
                        bool multi_cc_mode_dbg_ptx_only) {
   bool pass = true;
-  const char* stitch_name = "DISC_ENABLE_STITCH";
-  std::vector<const char*> stitch_fusion_flags;
-  char* stitch_flag = getenv(stitch_name);
-  if (stitch_flag) {
-    stitch_fusion_flags.push_back(stitch_flag);
-  } else {
-    stitch_fusion_flags.push_back("true");
-    stitch_fusion_flags.push_back("false");
-  }
-  for (const auto flag : stitch_fusion_flags) {
-    setenv(stitch_name, flag, 1);
+  auto envSettings = getEnvironmentSettings();
+  for (const auto& setting : envSettings) {
+    EnvSettingContext ctx(setting);
+
     for (auto backend_type : backend_types) {
       if (backend_type == BackendType::kCuda) {
 #if (GOOGLE_CUDA) || (TENSORFLOW_USE_ROCM)
diff --git a/tao_compiler/mlir/disc/tests/regression/algebra_simplifier.cc b/tao_compiler/mlir/disc/tests/regression/algebra_simplifier.cc
new file mode 100644
index 00000000000..6e1692c9b05
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/regression/algebra_simplifier.cc
@@ -0,0 +1,35 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/disc/tests/mlir_feature_test.h"
+#include "tensorflow/compiler/mlir/disc/tests/mlir_test.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir_test {
+
+const std::string c_ft_path =
+    "tensorflow/compiler/mlir/disc/tests/regression/data/";
+
+TEST(AlgebraSimplifier, PowF32DynamicTest) {
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "algebra_simplifier_pow_d_f32.mlir",
+      /*backend_types*/ kSupportedBackendList,
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"11x12xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
+}  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tests/regression/data/algebra_simplifier_pow_d_f32.mlir b/tao_compiler/mlir/disc/tests/regression/data/algebra_simplifier_pow_d_f32.mlir
new file mode 100644
index 00000000000..5ec1c9bd1bc
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/regression/data/algebra_simplifier_pow_d_f32.mlir
@@ -0,0 +1,10 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func @main(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.Const"() {value = dense<4.000000e+00> : tensor<f32>} : () -> tensor<f32>
+      %1:2 = tf_executor.island wraps "tf.Pow"(%arg0, %0) : (tensor<?x?xf32>, tensor<f32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %1 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc b/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc
index 284f7fe5065..bc607687760 100644
--- a/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc
+++ b/tao_compiler/mlir/disc/tests/regression/kstitch_fusion.cc
@@ -12,6 +12,7 @@
 #include "tensorflow/compiler/mlir/disc/tests/mlir_feature_test.h"
 #include "tensorflow/compiler/mlir/disc/tests/mlir_test.h"
 #include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/env_var.h"
 
 namespace mlir_test {
 
@@ -51,17 +52,26 @@ TEST(KStitchFusionGPUTest, KStitchSimpleLargeColumnF32) {
 // There is a root op that is not a skeleton op in the kStitch fusion.
 TEST(KStitchFusionGPUTest, KStitchNonSkeletonOutputF32) {
   setenv("DISC_ENABLE_STITCH", "true", 1);
-  // TODO: even though it means kStitch works when there are 2 kernels formed,
-  // we can further reduce the kernel number to 1 after optimizing shape ops
-  // like compute_reshape_shape successfully.
-  setenv("DISC_EXPECTED_KERNELS_IN_UT", "2", 1);
-  EXPECT_TRUE(feature_test_main(
-      /*mlir_file_path*/ c_ft_path + "kstitch_fusion_non_skl_output.mlir",
-      /*backend_types*/ {BackendType::kCuda},
-      /*num_inputs*/ 3,
-      /*num_outputs*/ 2,
-      /*input_descriptors*/ {"128x768xf32_X", "1x128x768xf32_X", "768xf32_X"},
-      /*output_descriptors*/ {"f32_X", "f32_X"}));
+  for (int i = 0; i < 2; ++i) {
+    bool enable_shape_constraint_ir = (i == 0);
+    setenv("DISC_ENABLE_SHAPE_CONSTRAINT_IR",
+           enable_shape_constraint_ir ? "true" : "false", 1);
+    if (enable_shape_constraint_ir) {
+      setenv("DISC_EXPECTED_KERNELS_IN_UT", "1", 1);
+    } else {
+      // TODO: even though it means kStitch works when there are 2 kernels
+      // formed, we can further reduce the kernel number to 1 after optimizing
+      // shape ops like compute_reshape_shape successfully.
+      setenv("DISC_EXPECTED_KERNELS_IN_UT", "2", 1);
+    }
+    EXPECT_TRUE(feature_test_main(
+        /*mlir_file_path*/ c_ft_path + "kstitch_fusion_non_skl_output.mlir",
+        /*backend_types*/ {BackendType::kCuda},
+        /*num_inputs*/ 3,
+        /*num_outputs*/ 2,
+        /*input_descriptors*/ {"128x768xf32_X", "1x128x768xf32_X", "768xf32_X"},
+        /*output_descriptors*/ {"f32_X", "f32_X"}));
+  }
   unsetenv("DISC_ENABLE_STITCH");
   unsetenv("DISC_EXPECTED_KERNELS_IN_UT");
 }
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_d_f32.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_d_f32.mlir
new file mode 100644
index 00000000000..0d8e28136c1
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_d_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func @main(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.Einsum"(%arg0, %arg1) {equation = "ijk,ikm->ijm"} : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>)
+      tf_executor.fetch %0 : tensor<?x?x?xf32>
+    }
+    return %graph : tensor<?x?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_p_f32.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_p_f32.mlir
new file mode 100644
index 00000000000..a9d8c8a81f5
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_p_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func @main(%arg0: tensor<?x4x5xf32>, %arg1: tensor<?x5x6xf32>) -> (tensor<?x4x6xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.Einsum"(%arg0, %arg1) {equation = "ijk,ikm->ijm"} : (tensor<?x4x5xf32>, tensor<?x5x6xf32>) -> (tensor<?x4x6xf32>)
+      tf_executor.fetch %0 : tensor<?x4x6xf32>
+    }
+    return %graph : tensor<?x4x6xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_s_f32.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_s_f32.mlir
new file mode 100644
index 00000000000..62bcb51991a
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/einsum_s_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func @main(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32>) -> (tensor<3x4x6xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.Einsum"(%arg0, %arg1) {equation = "ijk,ikm->ijm"} : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> (tensor<3x4x6xf32>)
+      tf_executor.fetch %0 : tensor<3x4x6xf32>
+    }
+    return %graph : tensor<3x4x6xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/einsum.cc b/tao_compiler/mlir/disc/tests/tensorflow_ops/einsum.cc
new file mode 100644
index 00000000000..a32bf58e937
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/einsum.cc
@@ -0,0 +1,61 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/disc/tests/mlir_feature_test.h"
+#include "tensorflow/compiler/mlir/disc/tests/mlir_test.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir_test {
+
+const std::string c_ft_path =
+    "tensorflow/compiler/mlir/disc/tests/tensorflow_ops/data/";
+
+// static shape test case
+TEST(EinsumTest, StaticShapeF32) {
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "einsum_s_f32.mlir",
+      /*backend_types*/
+      kSupportedBackendList,
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"3x4x5xf32_X", "3x5x6xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
+// partial dynamic shape test case
+TEST(EinsumTest, PartialDynamicShapeF32) {
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "einsum_p_f32.mlir",
+      /*backend_types*/
+      kSupportedBackendList,
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"3x4x5xf32_X", "3x5x6xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
+// fully dynamic shape test case
+TEST(EinsumTest, FullyDynamicShapeF32) {
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "einsum_d_f32.mlir",
+      /*backend_types*/
+      kSupportedBackendList,
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"3x4x5xf32_X", "3x5x6xf32_X"},
+      /*output_descriptors*/ {"f32_X"}));
+}
+
+}  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-algebra-simplifier.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-algebra-simplifier.mlir
index c88d4c7adf3..d681a478de8 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-algebra-simplifier.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-algebra-simplifier.mlir
@@ -55,4 +55,98 @@ func @bcast_const_not_integer(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
   return %3 : tensor<?x?xf32>
 }
 
+// -----
+
+// CHECK-LABEL: @broadcast
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<10x11xf32>) -> tensor<10x11xf32>
+func @broadcast(%arg0 : tensor<10x11xf32>) -> tensor<10x11xf32> {
+  // CHECK-NOT: mhlo.broadcast
+  // CHECK: return %[[ARG0]] : tensor<10x11xf32>
+  %0 = "mhlo.broadcast"(%arg0) {broadcast_sizes = dense<[]> : tensor<0xi64>} : (tensor<10x11xf32>) -> tensor<10x11xf32>
+  return %0 : tensor<10x11xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @broadcast_in_dim
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<10x11xf32>) -> tensor<10x11xf32>
+func @broadcast_in_dim(%arg0 : tensor<10x11xf32>) -> tensor<10x11xf32> {
+  // CHECK-NOT: mhlo.broadcast_in_dim
+  // CHECK: return %[[ARG0]] : tensor<10x11xf32>
+  %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<10x11xf32>) -> tensor<10x11xf32>
+  return %0 : tensor<10x11xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dynamic_broadcast_in_dim
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32, [@S0, @S1]>, %[[ARG1:.*]]: tensor<2xindex>) -> tensor<?x?xf32, [@S0, @S1]>
+func @dynamic_broadcast_in_dim(%arg0 : tensor<?x?xf32, [@S0, @S1]>, %arg1: tensor<2xindex>) -> tensor<?x?xf32, [@S0, @S1]> {
+  // CHECK-NOT: mhlo.dynamic_broadcast_in_dim
+  %0 = "mhlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<?x?xf32, [@S0, @S1]>, tensor<2xindex>) -> tensor<?x?xf32, [@S0, @S1]>
+  // CHECK return %[[ARG0]] : tensor<?x?xf32, [@S0, @S1]>
+  return %0 : tensor<?x?xf32, [@S0, @S1]>
+}
+
+// CHECK-DAG: "disc_shape.SymbolicDim"()
+// CHECK-SAME: sym_name = "S0"
+// CHECK-DAG: "disc_shape.SymbolicDim"()
+// CHECK-SAME: sym_name = "S1"
+"disc_shape.SymbolicDim"() {sym_name = "S0", value = -1 : i64} : () -> ()
+"disc_shape.SymbolicDim"() {sym_name = "S1", value = -1 : i64} : () -> ()
+
+// -----
+
+// CHECK-LABEL: broadcast_in_dim_of_reshape
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<10x11xf32>, %[[ARG1:.*]]: tensor<10x11x12xf32>) -> tensor<10x11x12xf32>
+func @broadcast_in_dim_of_reshape(%arg0: tensor<10x11xf32>, %arg1: tensor<10x11x12xf32>) -> tensor<10x11x12xf32> {
+  // CHECK-NOT: mhlo.reshape
+  // CHECK: %[[T0:.*]] = "mhlo.broadcast_in_dim"(%[[ARG0]])
+  // CHECK-SAME: broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
+  // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[ARG1]]
+  // CHECK: return %[[T1]]
+  %0 = "mhlo.reshape"(%arg0) : (tensor<10x11xf32>) -> tensor<10x11x1xf32>
+  %1 = "mhlo.broadcast_in_dim"(%0) {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>} : (tensor<10x11x1xf32>) -> tensor<10x11x12xf32>
+  %2 = mhlo.add %1, %arg1 : tensor<10x11x12xf32>
+  return %2 : tensor<10x11x12xf32>
+}
+
+// -----
+
+// CHECK-LABEL: dynamic_broadcast_in_dim_of_reshape
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32, [@S0, @S1]>, %[[ARG1:.*]]: tensor<?x?x?xf32, [@S0, @S1, @S2]>) -> tensor<?x?x?xf32, [@S0, @S1, @S2]>
+func @dynamic_broadcast_in_dim_of_reshape(%arg0: tensor<?x?xf32, [@S0, @S1]>, %arg1: tensor<?x?x?xf32, [@S0, @S1, @S2]>) -> tensor<?x?x?xf32, [@S0, @S1, @S2]> {
+  // CHECK-NOT: mhlo.dynamic_reshape
+  // CHECK: %[[T0:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]]
+  // CHECK-SAME: broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
+  // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[ARG1]]
+  // CHECK: return %[[T1]]
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32, [@S0, @S1]>
+  %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32, [@S0, @S1]>
+  %2 = tensor.from_elements %0, %1, %c1 : tensor<3xindex>
+  %3 = "mhlo.dynamic_reshape"(%arg0, %2) : (tensor<?x?xf32, [@S0, @S1]>, tensor<3xindex>) -> tensor<?x?x1xf32, [@S0, @S1, @C1]>
+  %4 = tensor.dim %arg1, %c0 : tensor<?x?x?xf32, [@S0, @S1, @S2]>
+  %5 = tensor.dim %arg1, %c1 : tensor<?x?x?xf32, [@S0, @S1, @S2]>
+  %6 = tensor.dim %arg1, %c2 : tensor<?x?x?xf32, [@S0, @S1, @S2]>
+  %7 = tensor.from_elements %4, %5, %6 : tensor<3xindex>
+  %8 = "mhlo.dynamic_broadcast_in_dim"(%3, %7) {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>} : (tensor<?x?x1xf32, [@S0, @S1, @C1]>, tensor<3xindex>) -> tensor<?x?x?xf32, [@S0, @S1, @S2]>
+  %9 = mhlo.add %8, %arg1 : tensor<?x?x?xf32, [@S0, @S1, @S2]>
+  return %9 : tensor<?x?x?xf32, [@S0, @S1, @S2]>
+}
+
+// CHECK-DAG: "disc_shape.SymbolicDim"()
+// CHECK-SAME: sym_name = "S0"
+// CHECK-DAG: "disc_shape.SymbolicDim"()
+// CHECK-SAME: sym_name = "S1"
+// CHECK-DAG: "disc_shape.SymbolicDim"()
+// CHECK-SAME: sym_name = "S2"
+// CHECK-DAG: "disc_shape.SymbolicDim"()
+// CHECK-SAME: sym_name = "C1"
+"disc_shape.SymbolicDim"() {sym_name = "S0", value = -1 : i64} : () -> ()
+"disc_shape.SymbolicDim"() {sym_name = "S1", value = -1 : i64} : () -> ()
+"disc_shape.SymbolicDim"() {sym_name = "S2", value = -1 : i64} : () -> ()
+"disc_shape.SymbolicDim"() {sym_name = "C1", value = 1 : i64} : () -> ()
 
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
new file mode 100644
index 00000000000..3eccb927c15
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
@@ -0,0 +1,226 @@
+// RUN: disc-opt --disc-shape-optimization -split-input-file %s | FileCheck %s
+
+// Test tensor.cast op
+
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<10xf32>) -> tensor<10xf32>
+func @main(%arg0 : tensor<?xf32>) -> tensor<10xf32> {
+  // CHECK: return %[[ARG0]] : tensor<10xf32>
+  %0 = tensor.cast %arg0 : tensor<?xf32> to tensor<10xf32>
+  return %0 : tensor<10xf32>
+}
+
+// -----
+
+// Test mhlo.add op
+
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<10xf32>, %[[ARG1:.*]]: tensor<10xf32>) -> tensor<10xf32>
+func @main(%arg0 : tensor<?xf32>, %arg1 : tensor<10xf32>) -> tensor<?xf32> {
+  %0 = tensor.cast %arg1 : tensor<10xf32> to tensor<?xf32>
+  // CHECK: %[[T1:.*]] = mhlo.add %[[ARG0]], %[[ARG1]] : tensor<10xf32>
+  %1 = "mhlo.add"(%arg0, %0) : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+  // CHECK: return %[[T1]] : tensor<10xf32>
+  return %1 : tensor<?xf32>
+}
+
+// -----
+
+// Test mhlo.concat op
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x10xf32, [@[[S0:.*]], @[[S1:.*]]]>, %[[ARG1:.*]]: tensor<?x10xf32, [@[[S0]], @[[S1]]]>) -> tensor<?x20xf32, [@[[S0]], @[[S2:.*]]]>
+func @main(%arg0 : tensor<?x10xf32>, %arg1 : tensor<?x10xf32>) -> tensor<?x?xf32> {
+  // CHECK: %[[T0:.*]] = "mhlo.concatenate"(%[[ARG0]], %[[ARG1]])
+  // CHECK: return %[[T0]] : tensor<?x20xf32, [@[[S0]], @[[S2]]]>
+  %0 = "mhlo.concatenate"(%arg0, %arg1) { dimension = 1 : i64 } : (tensor<?x10xf32>, tensor<?x10xf32>) -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+// Test mhlo.dot_general
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?x?xf32, [@[[S0:.*]], @[[S1:.*]], @[[S2:.*]], @[[S3:.*]]]>, %[[ARG1:.*]]: tensor<?x?x?x?xf32, [@[[S0]], @[[S1]], @[[S3]], @[[S4:.*]]]>) -> tensor<?x?x?x?xf32, [@[[S0]], @[[S1]], @[[S2]], @[[S4]]]>
+func @main(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+  // CHECK: %[[T0:.*]] = "mhlo.dot_general"(%[[ARG0]], %[[ARG1]])
+  // CHECK: return %[[T0]] : tensor<?x?x?x?xf32, [@[[S0]], @[[S1]], @[[S2]], @[[S4]]]>
+  %0 = "mhlo.dot_general"(%arg0, %arg1) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0, 1], rhs_batching_dimensions = [0, 1], lhs_contracting_dimensions = [3], rhs_contracting_dimensions = [2]>} : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
+  return %0 : tensor<?x?x?x?xf32>
+}
+
+// -----
+
+// Test mhlo.dot
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32, [@[[S0:.*]], @[[S1:.*]]]>, %[[ARG1:.*]]: tensor<?x?xf32, [@[[S1]], @[[S2:.*]]]>) -> tensor<?x?xf32, [@[[S0]], @[[S2]]]>
+func @main(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  // CHECK: %[[T0:.*]] = "mhlo.dot"(%[[ARG0]], %[[ARG1]])
+  // CHECK: return %[[T0]] : tensor<?x?xf32, [@[[S0]], @[[S2]]]>
+  %1 = "mhlo.dot"(%arg0, %arg1) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %1: tensor<?x?xf32>
+}
+
+// -----
+
+// Test mhlo.clamp: zero rank min/max
+
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32, [@[[S0:.*]]]>, %[[ARG1:.*]]: tensor<f32>, %[[ARG2:.*]]: tensor<f32>) -> tensor<?xf32, [@[[S0]]]>
+func.func @main(%arg0 : tensor<?xf32>, %arg1 : tensor<f32>, %arg2 : tensor<f32>) -> tensor<?xf32> {
+  // CHECK: %[[T0:.*]] = "mhlo.clamp"(%[[ARG1]], %[[ARG0]], %[[ARG2]])
+  // CHECK: return %[[T0]] : tensor<?xf32, [@[[S0]]]>
+  %0 = "mhlo.clamp"(%arg1, %arg0, %arg2) : (tensor<f32>, tensor<?xf32>, tensor<f32>) -> tensor<?xf32>
+  func.return %0 : tensor<?xf32>
+}
+
+// -----
+
+// Test mhlo.clamp: same shape
+
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32, [@[[S0:.*]]]>, %[[ARG1:.*]]: tensor<?xf32, [@[[S0]]]>, %[[ARG2:.*]]: tensor<?xf32, [@[[S0]]]>) -> tensor<?xf32, [@[[S0]]]>
+func.func @main(%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK: %[[T0:.*]] = "mhlo.clamp"(%[[ARG1]], %[[ARG0]], %[[ARG2]])
+  // CHECK: return %[[T0]] : tensor<?xf32, [@[[S0]]]>
+  %0 = "mhlo.clamp"(%arg1, %arg0, %arg2) : (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+  func.return %0 : tensor<?xf32>
+}
+
+// -----
+
+// Test mhlo.select : zero rank
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<i1>, %[[ARG1:.*]]: tensor<?xf32, [@[[S0:.*]]]>, %[[ARG2:.*]]: tensor<?xf32, [@[[S0]]]>) -> tensor<?xf32, [@[[S0]]]>
+func.func @main(%arg0: tensor<i1>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK: %[[T0:.*]] = "mhlo.select"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
+  // CHECK: return %[[T0]] : tensor<?xf32, [@[[S0]]]>
+  %0 = "mhlo.select"(%arg0, %arg1, %arg2)  : (tensor<i1>, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+  func.return %0 : tensor<?xf32>
+}
+
+// -----
+
+// Test mhlo.select : same shape
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xi1, [@[[S0:.*]]]>, %[[ARG1:.*]]: tensor<?xf32, [@[[S0]]]>, %[[ARG2:.*]]: tensor<?xf32, [@[[S0]]]>) -> tensor<?xf32, [@[[S0]]]>
+func.func @main(%arg0: tensor<?xi1>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK: %[[T0:.*]] = "mhlo.select"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
+  // CHECK: return %[[T0]] : tensor<?xf32, [@[[S0]]]>
+  %0 = "mhlo.select"(%arg0, %arg1, %arg2)  : (tensor<?xi1>, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+  func.return %0 : tensor<?xf32>
+}
+
+// -----
+
+// Test mhlo.einsum
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?xf32, [@[[S0:.*]], @[[S1:.*]], @[[S2:.*]]]>, %[[ARG1:.*]]: tensor<?x?x?xf32, [@[[S0]], @[[S2]], @[[S3:.*]]]>) -> tensor<?x?x?xf32, [@[[S0]], @[[S1]], @[[S3]]]>
+func @main(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> attributes {tf.entry_function = {input_placements = "cpu,cpu", inputs = "input0,input1", output_placements = "cpu", outputs = "output0"}} {
+  // CHECK: %[[T0:.*]] = "mhlo.einsum"(%[[ARG0]], %[[ARG1]])
+  // CHECK: return %[[T0]] : tensor<?x?x?xf32, [@[[S0]], @[[S1]], @[[S3]]]>
+  %0 = "mhlo.einsum"(%arg0, %arg1) {einsum_config = "ijk,ikm->ijm"} : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+  return %0 : tensor<?x?x?xf32>
+}
+
+// -----
+
+// Test arith.cmpi
+// CHECK-LABEL: main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?xf32, [@[[S0:.*]], @[[S1:.*]], @[[S2:.*]]]>, %[[ARG1:.*]]: tensor<?x?xf32, [@[[S3:.*]], @[[S4:.*]]]>) -> tensor<?x?xf32, [@[[S3]], @[[S4]]]>
+func @main(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c-1 = arith.constant -1 : index
+  %c2 = arith.constant 2 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[T0:.*]] = tensor.dim %[[ARG1]], %c0
+  // CHECK: %[[T1:.*]] = tensor.dim %[[ARG1]], %c1
+  // CHECK: %[[T2:.*]] = tensor.from_elements %[[T0]], %[[T1]] : tensor<2xindex>
+  // CHECK: %[[T3:.*]] = "mhlo.dynamic_reshape"(%[[ARG0]], %[[T2]])
+  // CHECK: return %[[T3]] : tensor<?x?xf32, [@[[S3]], @[[S4]]]>
+  %0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+  %2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
+  %3 = arith.muli %1, %0 : index
+  %4 = arith.muli %2, %3 : index
+  %5 = tensor.dim %arg1, %c0 : tensor<?x?xf32>
+  %6 = arith.muli %5, %c-1 : index
+  %7 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %8 = arith.muli %7, %6 : index
+  %9 = arith.cmpi eq, %4, %c0 : index
+  %10 = arith.cmpi slt, %8, %c0 : index
+  %11 = arith.ori %9, %10 : i1
+  %12 = arith.select %11, %c1, %8 : index
+  %13 = tensor.dim %arg1, %c0 : tensor<?x?xf32>
+  %14 = arith.cmpi eq, %13, %c-1 : index
+  %15 = arith.divui %4, %12 : index
+  %16 = arith.select %14, %15, %13 : index
+  %17 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %18 = arith.cmpi eq, %17, %c-1 : index
+  %19 = arith.divui %4, %12 : index
+  %20 = arith.select %18, %19, %17 : index
+  %21 = tensor.from_elements %16, %20 : tensor<2xindex>
+  %22 = "mhlo.dynamic_reshape"(%arg0, %21) : (tensor<?x?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
+  return %22 : tensor<?x?xf32>
+}
+
+// -----
+
+// Test disc_shape.tie_product_equal op case 1
+// CHECK-LABEL: @main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32, [@[[S0:.*]], @[[S1:.*]]]>, %[[ARG1:.*]]: tensor<1xindex>, %[[ARG2:.*]]: tensor<3xindex>) -> tensor<?x?x?xf32, [@[[S3:.*]], @[[S4:.*]], @[[S5:.*]]]>
+func @main(%arg0 : tensor<?x?xf32>, %arg1 : tensor<1xindex>, %arg2 : tensor<3xindex>) -> tensor<?x?x?xf32> {
+   // CHECK-NEXT: %[[T0:.*]] = "mhlo.dynamic_reshape"(%[[ARG0]], %[[ARG1]])
+   // CHECK-SAME: tensor<?xf32, [@[[S2:.*]]]>
+   %0 = "mhlo.dynamic_reshape"(%arg0, %arg1) : (tensor<?x?xf32>, tensor<1xindex>) -> tensor<?xf32>
+   // CHECK-NEXT: %[[T1:.*]] = mhlo.abs %[[T0]] : tensor<?xf32, [@[[S2]]]>
+   %1 = "mhlo.abs"(%0) : (tensor<?xf32>) -> tensor<?xf32>
+   // CHECK-NEXT: %[[T2:.*]] = "mhlo.dynamic_reshape"(%[[T1]], %[[ARG2]])
+   // CHECK-SAME: tensor<?x?x?xf32, [@[[S3]], @[[S4]], @[[S5]]]>
+   %2 = "mhlo.dynamic_reshape"(%1, %arg2) : (tensor<?xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+   // CHECK-NEXT: return %[[T2]] : tensor<?x?x?xf32, [@[[S3]], @[[S4]], @[[S5]]]>
+   return %2 : tensor<?x?x?xf32>
+}
+
+// CHECK-LABEL: @shape_constraint_graph
+// CHECK-DAG: %[[TT0:.*]] = "disc_shape.dim"() {name = @[[S0]]} : () -> index
+// CHECK-DAG: %[[TT1:.*]] = "disc_shape.dim"() {name = @[[S1]]} : () -> index
+// CHECK-DAG: %[[TT2:.*]] = "disc_shape.dim"() {name = @[[S2]]} : () -> index
+// CHECK-DAG: %[[TT3:.*]] = "disc_shape.dim"() {name = @[[S3]]} : () -> index
+// CHECK-DAG: %[[TT4:.*]] = "disc_shape.dim"() {name = @[[S4]]} : () -> index
+// CHECK-DAG: %[[TT5:.*]] = "disc_shape.dim"() {name = @[[S5]]} : () -> index
+// CHECK-DAG: "disc_shape.tie_product_equal"(%[[TT2]], %[[TT3]], %[[TT4]], %[[TT5]])
+// CHECK-DAG: "disc_shape.tie_product_equal"(%[[TT2]], %[[TT0]], %[[TT1]])
+// CHECK-DAG: "disc_shape.tie_product_equal"(%[[TT0]], %[[TT1]], %[[TT3]], %[[TT4]], %[[TT5]])
+
+// -----
+
+// Test disc_shape.tie_product_equal op case 2
+
+// CHECK-LABEL: @main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?xf32, [@[[S0:.*]], @[[S1:.*]], @[[S2:.*]]]>) -> tensor<?x?xf32, [@[[S1]], @[[S2]]]>
+func @main(%arg0: tensor<?x?x?xf32>) -> tensor<?x?xf32> {
+  %c2 = arith.constant 2 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %0 = mhlo.constant dense<-0.000000e+00> : tensor<f32>
+  %1 = mhlo.abs %arg0 : tensor<?x?x?xf32>
+  %2 = tensor.dim %1, %c0 : tensor<?x?x?xf32>
+  %3 = tensor.dim %1, %c1 : tensor<?x?x?xf32>
+  %4 = tensor.dim %1, %c2 : tensor<?x?x?xf32>
+  %5 = arith.muli %3, %4 : index
+  %6 = tensor.from_elements %2, %5 : tensor<2xindex>
+  %7 = "mhlo.dynamic_reshape"(%1, %6) : (tensor<?x?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
+  %8 = mhlo.reduce(%7 init: %0) applies mhlo.add across dimensions = [0] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?xf32>
+  %9 = tensor.dim %1, %c1 : tensor<?x?x?xf32>
+  %10 = tensor.dim %1, %c2 : tensor<?x?x?xf32>
+  %11 = tensor.from_elements %9, %10 : tensor<2xindex>
+  %12 = "mhlo.dynamic_reshape"(%8, %11) : (tensor<?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
+  return %12 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: @shape_constraint_graph
+// CHECK-DAG: %[[TT1:.*]] = "disc_shape.dim"() {name = @[[S1]]} : () -> index
+// CHECK-DAG: %[[TT2:.*]] = "disc_shape.dim"() {name = @[[S2]]} : () -> index
+// CHECK-DAG: %[[TT3:.*]] = "disc_shape.dim"() {name = @[[S3]]} : () -> index
+// CHECK-DAG: "disc_shape.tie_product_equal"(%[[TT3]], %[[TT1]], %[[TT2]])
+