[mlir][shape] add outline-shape-computation pass

[yaochengji]

Summary

• add outline-shape-computation pass

[jpienaar]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

15

No punctuation at the end (https://mlir.llvm.org/docs/OpDefinitions/#operation-documentation)

Could you also add a description of the pass?

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

52

Style followed here is for single line conditionals the {} is omitted (https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements)

60

\p notatiion isn't generally used here. shape is more common.

64

This looks like an assert? (assert(!cluster.empty() && ...)

70

If you use ValueRange , you can do getTypes() and can also skip creating vector here.

76

Is there a better location that could be given? Perhaps NameLoc related to function.

83

Same re single instruction {}

94

typo: inconvenient

95

Could you expand here how this orders instructions ?

124

Does symbol names need to be unique at Module level? Can symbol table behavior of Module be used to handle this automatically for you?

136

Lets use some constant rather than -1 (-1 is also used in ShapedType::kDynamic and if the goal is return the equivalent of a dynamic dim, then that constant could be used here)

320

Please add newline at end

mlir/test/Dialect/Shape/outline-shape-computation.mlir

4

Could you add a description of what is being tested?

[yaochengji]

misc fixes

[yaochengji]

Note that this diff should be finalized after https://reviews.llvm.org/D131977 is merged.

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

52

It is fixed as well as other corresponding places.

76

Here I use the loc of value' of shape.with_shape, do you think it's OK?

95

Add some explanation, hope it could help.

124

1. It does if we the arguments of the shape computation function is out of current original function.
2. Symbol table could not be since it could only handle operations with symbol trait, while it is symbol attribute here.

136

use llvm::Optional instead.

mlir/test/Dialect/Shape/outline-shape-computation.mlir

4

Added a concise description here, as well as more tests.

Detailed description could be fetched from the Passes.td.

[yaochengji]

rebase on new main branch

[jpienaar]

Nice, thanks

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

17

A bit more context would be good. E.g., do we start with a function with reified shape computations? Where is this expected kind of level. And then post that the why one would do that (this is more detail than the other passes here but I think the others are simpler/more obvious:-))

17

Nit: space before ( (here and below)

20

What happens if the RankedTensorType already has an encoding? I'm guessing no expectations here if another pass post this also changes encoding

22

Is there scoping for the symbol? E.g., is this a global value, per function definition/call site?

24

And the agurments are symbols again?

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

76

Yes, that will help in error reporting. If this turns out to not be sufficient we can look at FusedLoc (but that'll be very large in some cases)

89

Could you document return and expected usage? (DenseMap has non deterministic iteration order I believe, so shouldn't be iterated upon unless result of iteration doesn't effect produced IR)

244

And this effects main function or only inside outlined shape function (s)?

369

What does cal mean here? (Perhaps val/value?)

[yaochengji]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

17

Done.

17

Added more details, including

1. usually there's shape reification before this pass
2. which level
3. why this pass is needed?

I'm not a good doc writer, feel free to modify it directly :)

20

It is not handled by now. And I supplemented a TODO here.

22

It is global. Considering symbol might be referred across function.

24

Yes. Symbols are added in Type.

It is because Value and Operation are volatile, while Type is relatively stable.

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

244

This pattern is applied on main function before the shape-outlining logic.

369

cal is short for calculate :)

[jpienaar]

Nice and sorry for delays, as mentioned traveling a bit but luckily not blocker

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

20

The pass should not overwrite the type in that case/exit when these encountered.

That means this pass can't be used with sparsity today and with MHLO & StableHLO shortly (i think today already it can't be, but only partially rolled out). That's undesirable. Instead we can get more use if this pass is split in two. Have the outline functions & populating an Analysis in this pass (for Analysis a key of Value and index should be sufficient, extending that to subtypes with nested shapes would be a natural extension later) and the pass (potentially downstream) that follows it would serialize the Analysis into the type and perform type propagation/joins across call sites (the follow up pass could also mark the Analysis as retained).

This would allow this outlining and analysis to be used in more general contexts (incl MHLO ones, torch-mlir with it's custom tensor types and I believe MemRef/Vector type cases) which would be great.

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

369

Ha :-)

mlir/test/Dialect/Shape/outline-shape-computation.mlir

2

Could you add a test case where the shape computation shares values with a value being computed? (E.g., shape computation interleaved with and feeding into other computations/returned)

Could you add a test where you have multiple reused shape computations? Say you have two dims, concat, associate , concat dims while repeating last 1x, associate, concat dims while repeating last dim 2x, associate etc and return all associated values - as if you had unrolled a loop with a fill inside. This one would be interesting in both unsimplified and CSE'd form.

63

You may need a CHECK-NOT to verify only one is generated.

[yaochengji]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

20

1. To confirm, I'm using two dynamic shapes: a trivial symbol and a shape.func test as an example.

After the pass, the IR will be like this? The difference with the current one is that all the types remain unchanged. And there will be a retained Analysis used to mapping the Values to their shape computation function. In this case, the table will be {%0 : %arg0, %1 : [ @shape_cal_0, %arg0]} ?

func.func @main(%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
    %0 = "test.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
    %1 = "test.concat"(%0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
     return %1 : tensor<?x4x?xf32>

shape.func private @shape_cal_0
    ...

If so (specifically the mapping table uses Value as a key), it will conflict our use cases. The original purpose of the pass is to preserve shape computation info as well as not let the computation part affect lowering or optimization passes badly (Ex. shape computation part will add more users to original Value). Based on these two preconditions, it is quite possible that %0 and %1 will be replaced by some new Values after some passes, which invalidates the mapping table of the above Analysis.

2. I'm also interested about why type cannot be overwrite with sparsity/mhlo/stablehlo, could you give me more details, a link to doc or code? We also use the encoding attr of RankedTensorType to represent bounded shape info. Maybe I need to persuade my colleagues to think of a better way.

[jpienaar]

Sorry fell off review queue while I was out.

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

20

1. Yes, so that's why in this pass you populate an Analysis and the pass just after this one you encode it in the type in some way. Encoding it here in the type would limit utility, but if it is useful for your flow then the pass just after (and both have to be module passes here whether or not 1 pass as you are mutating types across callsites) would capture it in type. That way you do get the behavior supporting your original use case but also have this as a general utility pass that can be used where folks have different assumptions. E.g., you _still_ encode in Type, just in the pass that runs directly after this one and that is more specialized for a deployment route (see 2 below).

2. Sure, so if you and sparse compiler folks both use encoding attribute, what attribute is stored there? Sparse uses mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td while MHLO & StableHLO uses (or uses some variant of) https://github.com/tensorflow/mlir-hlo/blob/8a3ec6e5187e94a6c4f91d67a69545b9a1075d3e/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_attrs.td#L171. Now what you could do here is: create a new Attribute that contains both of these or implements both sets of classes (and can be cast to both) and then you can still encode in Type. BUT you can't do that here as you don't have access to the downstream Attribute type here. Which is why if you make this 2 passes, you could have this pass here that does 99% of the work and then another pass that does know about the MHLO/StableHLO attributes and can produce that new Attribute. For sparse you could do the same, but it gets a bit messy to just accumulate all different encodings into a new mega encoding, so I'd consider: where do you need what? Potentially where you do the outlining you could just nest sparsity inside your attribute and then post where you utilize it in codegen, flip to the nested sparse attribute and let the sparse codegen do its thing (or vice versa). Dropping info once you've utilized it is fine, doing it so needlessly would be bad.

An alternative would be to encode it via an op. This is where some of the shape assuming ops do (and this would be good candidate to add one that has a function reference to make it even simpler). Also what IREE does flow dialect side. The main downside to ops is that it obscures use-def chains. So doing it such that it gets out of the way of what you do want to retain or optimize is good (of course you could also have passes that optimize across such metadata ops knowingly). I think the op route will be better long term. It'll be cheaper and a bit more flexible/robust. But I don't think an ergonomic version of it exists at the moment.

[yaochengji]

populate shape mapping information into Analysis instead and add two required tests.

[jpienaar]

I think this is a good place to iterate from! It can be refined. Thanks for making this into 2. Couple of changes, but nothing major. I'll help landing post.

[seems I didn't send my last reply of "Sorry fell off review queue while I was out.", I get confused with all the review systems when which one sends what comments]

mlir/include/mlir/Dialect/Shape/Analysis/ShapeMappingAnalysis.h

24

Could you document these?

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

179

The naming convention here doesn't follow suffixed with _ for class members.

241

This would be wrong here I believe as you are mutating the module structure and the folding above could change behavior too. You don't need to do this one here if the goal is to avoid invalidation before use, that you can control in your pipeline for this case. Unless there is another analysis that existed before here that you wish to retain.

mlir/test/Dialect/Shape/outline-shape-computation.mlir

8

Unfortunately I think DAG will allow this block arg matching to be interchanged. E.g., even if you had the result from shape_cal_1 swapped with here it would not complain. Unfortunately we don't have a DAG-NEXT thing.

Is DAG here for blocks of function? (e.g., could line 5-8 be switched with 10-13 conceptually). In your pass I think the ordering is fixed, it is not load bearing the order that the functions are emitted but it is deterministically fixed. I just want to avoid the case where the result is arbitrarily mangled. Another alternative is to sort the shape functions in the print function - given its only for debugging the extra work/cost is OK and then here you don't need DAG.

[yaochengji]

[seems I didn't send my last reply of "Sorry fell off review queue while I was out.", I get confused with all the review systems when which one sends what comments]

---

I saw this and it's fine. I didn't actively ping you because I was also quite busy previously.

[yaochengji]

mlir/include/mlir/Dialect/Shape/Analysis/ShapeMappingAnalysis.h

24

Done.

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

179

Done.

241

Deleted. I cannot find an example about controlling analysis in pass pipeline. Could you point out it for me if possible? And as currently the analysis is not preserved, the analysis output is not tested in lit test for now.

mlir/test/Dialect/Shape/outline-shape-computation.mlir

8

Line 5-8 could be switched with 10-13, because it's a map and the order of keys is undetermined when iterating.

BTW, the order of created shape.funcs is determined.

[yaochengji]

format code

[jpienaar]

Did a scan and submitted with small changes.

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

61

Note: this returns a reference to on stack object. Fixed it.

241

I couldn't find good reference, that needs to be fixed (Analysis themselves need a bit of TLC). And you were correct one has to mark the Analysis being retained here as this is not a pure Analysis. Added back the check and marked only this Analysis as retained.

I think this can be improved still, but don't want to block on that now as its not directly related to this change.

mlir/test/Dialect/Shape/outline-shape-computation.mlir

8

Yes, but CHECK-DAG will allow interleaving at line level. So

Shape function symbol: @shape_cal_1
Shape function symbol: @shape_cal_0
Shape function arguments:
...

would pass the test, and so in particular here it would be unable to detect if you got the analysis wrong and switched results of 1 with those of 0. Here the block arguments being the same would result in not finding this.

Updated to make this more concise and have DAG matching work without allowing the above accidental passing case. Can revise later.

[csigg]

FYI, I pushed 5faebb5 fix for Windows for the test added in this change.

[yaochengji]

Thanks @jpienaar and @csigg

mlir/include/mlir/Dialect/Shape/Analysis/ShapeMappingAnalysis.h

+//===- ShapeMappingAnalysis.h - Preserve shape Info  ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SHAPE_ANALYSIS_SHAPEMAPPINGANALYSIS_H_
+#define MLIR_DIALECT_SHAPE_ANALYSIS_SHAPEMAPPINGANALYSIS_H_
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+
+namespace shape {
+
+/// ShapeMappingValue works as the value of ShapeMappingAnalysis table, where
+/// `funcSymbol` is the symbol of mapping function, and `inputs` are the actual
+/// parameters for the function.
+struct ShapeMappingValue {
+  ShapeMappingValue() = default;
+  ShapeMappingValue(FlatSymbolRefAttr symbol, llvm::SmallVector<Value> &&inps)
+      : funcSymbol(symbol), inputs(inps) {}
+
+  FlatSymbolRefAttr funcSymbol;
+  llvm::SmallVector<Value> inputs;
+};
+
+/// ShapeMappingAnalysis is used together with OutlineShapeComputationPass to
+/// preserve Value and corresponding shape function / arguments mapping
+/// infomation
+struct ShapeMappingAnalysis {
+
+  ShapeMappingAnalysis(Operation *op) : operation(op) {}
+
+  /// Dumps the shape mapping information to the given stream.
+  void print(raw_ostream &os) const {
+    os << "// ---- Shape Mapping Infomation -----\n";
+    for (auto it : shapeMapping) {
+      const ShapeMappingValue &mappingValue = it.second;
+      os << "// - Shape for: " << it.first << "\n";
+      os << "// - Shape function symbol: " << mappingValue.funcSymbol << "\n";
+      os << "// - Shape function arguments: \n";
+      for (Value v : mappingValue.inputs) {
+        os << "// - " << v << "\n";
+      }
+      os << "// - End of one shape mapping item -----\n\n";
+    }
+  }
+
+  llvm::DenseMap<Value, ShapeMappingValue> shapeMapping;
+
+private:
+  Operation *operation;
+};
+
+} // namespace shape
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SHAPE_ANALYSIS_SHAPEMAPPINGANALYSIS_H_

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

 
 #ifndef MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 class ConversionTarget;
+class ModuleOp;
 class TypeConverter;
 namespace func {
 class FuncOp;
 } // namespace func
 } // namespace mlir
 
 namespace mlir {
 

 
 // Bufferizes shape dialect ops.
 //
 // Note that most shape dialect ops must be converted to std before
 // bufferization happens, as they are intended to be bufferized at the std
 // level.
 std::unique_ptr<OperationPass<func::FuncOp>> createShapeBufferizePass();
 
+/// Outline the shape computation part by adding shape.func and populate
+/// conrresponding mapping infomation into ShapeMappingAnalysis.
+std::unique_ptr<OperationPass<ModuleOp>> createOutlineShapeComputationPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Shape/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/Shape/Transforms/Passes.td

 //===-- Passes.td - ShapeOps pass definition file ----------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES
 #define MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 
+def OutlineShapeComputation : Pass<"outline-shape-computation", "ModuleOp"> {
+  let summary = "Using shape.func to preserve shape computation";
+  let description = [{
+    This pass outlines the shape computation part in high level IR by adding
+    shape.func and and populate conrresponding mapping infomation into
+    ShapeMappingAnalysis. The shape computation part is usually introduced by
+    shape reification, and each single dynamic shape is denoted by shape.with_shape.
+    
+    There're two main reasons this shape-outline pass is needed:
+    1. Many passes don't take shape reification part into consideration.
+       Therefore we need to "remove" the shape reification part temporarily for
+       these passes.
+    2. Sometimes we cannot redo shape reification after converting from dialect
+       A to dialect B. Because op-level shape reification is only implemented
+       on A.
+    
+    Input:
+    ```mlir
+    func.func @main(%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) ->
+      tensor<?x4x?xf32> {
+      %c2 = arith.constant 2 : index
+      %c0 = arith.constant 0 : index
+      %c4 = arith.constant 4 : index
+      %0 = shape.shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+      %1 = shape.get_extent %0, %c2 : tensor<3xindex>, index -> index
+      %2 = "test.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
+      %3 = shape.with_shape %2, %0 : tensor<?x4x?xf32>, tensor<3xindex>
+      %4 = shape.value_of %3 : tensor<?x4x?xf32>
+      %5 = "test.concat"(%4, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>,
+            tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+      %6 = shape.get_extent %0, %c0 : tensor<3xindex>, index -> index
+      %7 = arith.addi %6, %c2 : index
+      %8 = shape.from_extents %7, %c4, %1 : index, index, index
+      %9 = shape.with_shape %5, %8 : tensor<?x4x?xf32>, !shape.shape
+      %10 = shape.value_of %9 : tensor<?x4x?xf32>
+      return %10 : tensor<?x4x?xf32>
+    }
+    ```
+
+    Output
+    ```mlir
+    func.func @main(%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) ->
+      tensor<?x4x?xf32> {
+      %0 = "test.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
+      %1 = "test.concat"(%0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, 
+            tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+      return %1 : tensor<?x4x?xf32>
+    }
+    shape.func private @shape_cal_1(%arg0: tensor<?x4x?xf32>) -> !shape.shape {
+      %c2 = arith.constant 2 : index
+      %c0 = arith.constant 0 : index
+      %c4 = arith.constant 4 : index
+      %0 = shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+      %1 = get_extent %0, %c2 : tensor<3xindex>, index -> index
+      %2 = get_extent %0, %c0 : tensor<3xindex>, index -> index
+      %3 = arith.addi %2, %c2 : index
+      %4 = from_extents %3, %c4, %1 : index, index, index
+      return %4 : !shape.shape
+    }
+    shape.func private @shape_cal_0(%arg0: tensor<?x4x?xf32>) -> tensor<3xindex> {
+      %0 = shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+      return %0 : tensor<3xindex>
+    }
+    ```
+
+    For the above example, the shape computation is inlined in the input IR,
+    which is used for two values' (test.abs and test.concat) shape. And the shape
+    compuatation part is outlined in the output IR.
+
+    And the shape mapping infomation will be:
+    // ---- Shape Mapping Infomation -----
+    // - Shape for: %0 = "test.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
+    // - Shape function symbol: @shape_cal_0
+    // - Shape function arguments:
+    // - <block argument> of type 'tensor<?x4x?xf32>' at index: 0
+    // - End of one shape mapping item -----
+
+    // - Shape for: %1 = "test.concat"(%0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+    // - Shape function symbol: @shape_cal_1
+    // - Shape function arguments:
+    // - <block argument> of type 'tensor<?x4x?xf32>' at index: 0
+    // - End of one shape mapping item -----
+
+  }];
+  let constructor = "mlir::createOutlineShapeComputationPass()";
+  let dependentDialects = ["shape::ShapeDialect"];
+}
+
 def RemoveShapeConstraints : Pass<"remove-shape-constraints", "func::FuncOp"> {
   let summary = "Replace all cstr_ ops with a true witness";
   let constructor = "mlir::createRemoveShapeConstraintsPass()";
 }
 
 def ShapeToShapeLowering : Pass<"shape-to-shape-lowering", "func::FuncOp"> {
   let summary = "Legalize Shape dialect to be convertible to Arith";
   let constructor = "mlir::createShapeToShapeLowering()";

mlir/lib/Dialect/Shape/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRShapeOpsTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
+  OutlineShapeComputation.cpp
   RemoveShapeConstraints.cpp
   ShapeToShapeLowering.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ShapeOps/Transforms
 
   DEPENDS
   MLIRShapeTransformsIncGen

mlir/lib/Dialect/Shape/Transforms/OutlineShapeComputation.cpp

+//====----- OutlineShapeComputation.cpp -------------------------*- C++-*--===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Shape/Analysis/ShapeMappingAnalysis.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Shape/Transforms/Passes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/Support/Debug.h"
+#include <queue>
+#include <unordered_set>
+#include <vector>
+
+namespace mlir {
+#define GEN_PASS_DEF_OUTLINESHAPECOMPUTATION
+#include "mlir/Dialect/Shape/Transforms/Passes.h.inc"
+} // namespace mlir
+
+#define DEBUG_TYPE "outline-shape-computation"
+
+using namespace mlir;
+
+namespace {
+
+// A Value is an input of the cluster if it is an operand of an operation in the
+// cluster and its defining operation is not in the cluster.
+ValueRange
+getInputsOfCluster(const llvm::SmallVector<Operation *, 8> &cluster) {
+  SmallVector<Value, 4> inputs;
+  llvm::SmallDenseSet<Value> inputSet;
+  llvm::SmallDenseSet<Operation *> opSet;
+  for (Operation *op : cluster) {
+    bool inserted = opSet.insert(op).second;
+    (void)inserted;
+    assert(inserted && "cluster contains duplicate operations");
+  }
+
+  for (Operation *op : cluster) {
+    for (Value operand : op->getOperands()) {
+      Operation *operandOp = operand.getDefiningOp();
+      if (opSet.find(operandOp) != opSet.end()) {
+        // skip if defining op is in the cluster
+        continue;
+      }
+      if (inputSet.insert(operand).second)
+        inputs.push_back(operand);
+    }
+  }
+  return inputs;
+}
+
+// Create a shape.func representing the shape computation for `shape`.
+std::pair<shape::FuncOp, SmallVector<Value>>
+createFuncFromCluster(OpBuilder &b, const SmallVector<Operation *, 8> &cluster,
+                      Value shape, StringRef fnName, Location loc) {
+  ValueRange inputs = getInputsOfCluster(cluster);
+  auto fnType = cluster.empty()
+                    ? b.getFunctionType(shape.getType(), shape.getType())
+                    : b.getFunctionType(inputs.getTypes(), shape.getType());
+  shape::FuncOp fnOp = b.create<shape::FuncOp>(loc, fnName, fnType);
+  Block *block = fnOp.addEntryBlock();
+  b.setInsertionPoint(block, block->end());
+  BlockAndValueMapping bvm;
+  if (cluster.empty())
+    bvm.map(shape, fnOp.getArgument(0));
+  else
+    for (auto inputAndArg : llvm::zip(inputs, fnOp.getArguments())) {
+      bvm.map(std::get<0>(inputAndArg), std::get<1>(inputAndArg));
+    }
+
+  for (Operation *op : cluster) {
+    b.clone(*op, bvm);
+  }
+  llvm::SmallVector<Value, 4> fnReturns;
+  fnReturns.push_back(bvm.lookupOrDefault(shape));
+
+  b.create<shape::ReturnOp>(UnknownLoc::get(b.getContext()), fnReturns);
+  fnOp.setPrivate();
+  return std::make_pair(fnOp, inputs);
+}
+
+// The operations in the cluster might be unsorted, which could be inconvenient
+// when creating shape.func op.
+DenseMap<Value, SmallVector<Operation *, 8>>
+getOrderedClusters(const DenseMap<Value, DenseSet<Operation *>> &clusters,
+                   func::FuncOp funcOp) {
+  // Compute all clusters that each operation is in
+  DenseMap<Operation *, SmallVector<Value>> op2Shapes;
+  for (auto it : clusters) {
+    Value shape = it.first;
+    const DenseSet<Operation *> &cluster = it.second;
+    for (Operation *cOp : cluster) {
+      op2Shapes[cOp].push_back(shape);
+    }
+  }
+
+  // Iterate through all operations in order. Get all the clusters `cOp` belongs
+  // to and construct the new ordered cluster as it traverses.
+  DenseMap<Value, SmallVector<Operation *, 8>> orderedClusters;
+  funcOp.walk([&](Operation *op) {
+    auto it = op2Shapes.find(op);
+    if (it != op2Shapes.end()) {
+      Operation *cOp = it->first;
+      for (Value shape : it->second) {
+        orderedClusters[shape].push_back(cOp);
+      }
+    }
+  });
+
+  return orderedClusters;
+}
+
+void constructShapeFunc(
+    const std::vector<shape::WithOp> &allWithOps, MLIRContext *context,
+    DenseMap<Value, SmallVector<Operation *, 8>> &clusters,
+    SymbolTable &symbolTable,
+    DenseMap<Value, shape::ShapeMappingValue> &dynShape2ShapeFunc, func::FuncOp funcOp,
+    shape::ShapeMappingAnalysis &shapeMappingAnalysis) {
+  std::string shapeCalculationNamePrefix = "shape_cal_";
+  int shapeCalculationNameIdx = 0;
+  OpBuilder builder(context);
+
+  // Construct a shape function
+  for (shape::WithOp withOp : allWithOps) {
+    Value value = withOp.getOperand();
+    Value shape = withOp.getShape();
+    RankedTensorType rankedType = value.getType().dyn_cast<RankedTensorType>();
+    if (rankedType == nullptr)
+      continue;
+
+    const SmallVector<Operation *, 8> &cluster = clusters[shape];
+    shape::ShapeMappingValue shapeMappingValue;
+    auto it = dynShape2ShapeFunc.find(shape);
+    if (it == dynShape2ShapeFunc.end()) {
+      std::string name = shapeCalculationNamePrefix +
+                         std::to_string(shapeCalculationNameIdx++);
+      Location loc = value.getLoc();
+      builder.setInsertionPointAfter(funcOp);
+      auto pair = createFuncFromCluster(builder, cluster, shape, name, loc);
+      const SmallVector<Value> &inputs = pair.second;
+      shape::FuncOp shapeFuncOp = pair.first;
+      StringAttr insertedName = symbolTable.insert(shapeFuncOp);
+      auto symbol = FlatSymbolRefAttr::get(context, insertedName);
+
+      shapeMappingValue.funcSymbol = symbol;
+      shapeMappingValue.inputs = inputs;
+    } else {
+      shapeMappingValue = it->second;
+    }
+    dynShape2ShapeFunc[shape] = shapeMappingValue;
+    shapeMappingAnalysis.shapeMapping.insert(
+        std::make_pair(value, shapeMappingValue));
+  }
+}
+
+struct OutlineShapeComputationPass
+    : public impl::OutlineShapeComputationBase<OutlineShapeComputationPass> {
+
+  void runOnOperation() override;
+
+private:
+  bool calOnlyUsedByWithShapesRecursively(Operation *op, Value prevOutput);
+  void getClusterFromValue(Value shape,
+                           DenseMap<Value, DenseSet<Operation *>> &clusters);
+  DenseMap<Value, SmallVector<Operation *, 8>>
+  constructClustersForEachShape(const std::vector<shape::WithOp> &allWithOps,
+                                func::FuncOp funcOp);
+  DenseMap<Operation *, bool> onlyUsedByWithShapes;
+};
+
+class TensorDimOpRewriter : public OpRewritePattern<tensor::DimOp> {
+  using OpRewritePattern<tensor::DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::DimOp op,
+                                PatternRewriter &rewriter) const override {
+    auto shapeOf =
+        rewriter.create<shape::ShapeOfOp>(op.getLoc(), op.getSource());
+    rewriter.replaceOpWithNewOp<shape::GetExtentOp>(op, op.getType(), shapeOf,
+                                                    op.getIndex());
+    return success();
+  }
+};
+
+void OutlineShapeComputationPass::runOnOperation() {
+  ModuleOp moduleOp = getOperation();
+  SymbolTable symbolTable(moduleOp);
+  DenseMap<Value, shape::ShapeMappingValue> dynShape2ShapeFunc;
+  auto &shapeMappingAnalysis = getAnalysis<shape::ShapeMappingAnalysis>();
+
+  moduleOp.walk([&](func::FuncOp funcOp) {
+    MLIRContext *context = funcOp.getContext();
+    RewritePatternSet prevPatterns(context);
+    prevPatterns.insert<TensorDimOpRewriter>(context);
+    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(prevPatterns))))
+      return signalPassFailure();
+
+    // initialize class member `onlyUsedByWithShapes`
+    onlyUsedByWithShapes.clear();
+    funcOp.walk([&](Operation *op) {
+      calOnlyUsedByWithShapesRecursively(op, /*prevOutput=*/nullptr);
+    });
+    // clang-format off
+    LLVM_DEBUG(llvm::dbgs() << "onlyUsedByWithShapes table: \n";
+               for (auto it : onlyUsedByWithShapes) 
+                 if (it.second)
+                   llvm::dbgs() << *(it.first) << ": " << it.second << "\n";);
+    // clang-format on
+
+    // collect all the shape.with_shape ops.
+    std::vector<shape::WithOp> allWithOps;
+    funcOp.walk([&](shape::WithOp withOp) { allWithOps.push_back(withOp); });
+
+    DenseMap<Value, SmallVector<Operation *, 8>> clusters =
+        constructClustersForEachShape(allWithOps, funcOp);
+    constructShapeFunc(allWithOps, context, clusters, symbolTable,
+                       dynShape2ShapeFunc, funcOp, shapeMappingAnalysis);
+
+    for (shape::WithOp withOp : allWithOps) {
+      Value value = withOp.getOperand();
+      for (Operation *user : withOp.getResult().getUsers()) {
+        if (Value valueOf = llvm::dyn_cast<shape::ValueOfOp>(user))
+          valueOf.replaceAllUsesExcept(value, withOp);
+      }
+    }
+
+    // dce
+    if (failed(applyPatternsAndFoldGreedily(funcOp, {})))
+      return signalPassFailure();
+  });
+}
+
+DenseMap<Value, SmallVector<Operation *, 8>>
+OutlineShapeComputationPass::constructClustersForEachShape(
+    const std::vector<shape::WithOp> &allWithOps, func::FuncOp funcOp) {
+  DenseMap<Value, DenseSet<Operation *>> clusters;
+  for (shape::WithOp withOp : allWithOps) {
+    Value shape = withOp.getShape();
+    if (clusters.count(shape) == 0)
+      getClusterFromValue(shape, clusters);
+  }
+  return getOrderedClusters(clusters, funcOp);
+}
+
+// The output of a cluster is the `shape`, and the inputs are the outputs of
+// operations who are not in `onlyUsedByWithShapes`
+void OutlineShapeComputationPass::getClusterFromValue(
+    Value shape, DenseMap<Value, DenseSet<Operation *>> &clusters) {
+  DenseSet<Operation *> cluster;
+  Operation *defOp = shape.getDefiningOp();
+
+  DenseSet<Operation *> visited;
+  std::queue<Operation *> queue;
+  // defOp == nullptr means shape is the argument of the func op
+  if (defOp != nullptr) {
+    visited.insert(defOp);
+    queue.push(defOp);
+  }
+  while (!queue.empty()) {
+    Operation *op = queue.front();
+    queue.pop();
+    if (onlyUsedByWithShapes.lookup(op)) {
+      cluster.insert(op);
+      for (Value inp : op->getOperands()) {
+        Operation *inpDefOp = inp.getDefiningOp();
+        if (nullptr != inpDefOp && !visited.contains(inpDefOp)) {
+          visited.insert(inpDefOp);
+          queue.push(inpDefOp);
+        }
+      }
+    }
+  }
+
+  clusters[shape] = std::move(cluster);
+}
+
+// return true if `op` is a shape.with_shape, or all the users' of `op`
+// eventually point to the shape operand of shape.with_shape ops
+bool OutlineShapeComputationPass::calOnlyUsedByWithShapesRecursively(
+    Operation *op, Value prevOutput) {
+  auto it = onlyUsedByWithShapes.find(op);
+  if (it != onlyUsedByWithShapes.end())
+    return it->second;
+
+  if (auto withOp = llvm::dyn_cast<shape::WithOp>(op)) {
+    bool retVal = withOp.getShape() == prevOutput;
+    return retVal;
+  }
+
+  if (op->use_empty()) {
+    onlyUsedByWithShapes[op] = false;
+    return false;
+  }
+
+  bool allUsers = true;
+  for (Value oup : op->getResults())
+    for (Operation *user : oup.getUsers())
+      allUsers &= calOnlyUsedByWithShapesRecursively(user, oup);
+
+  onlyUsedByWithShapes[op] = allUsers;
+  return allUsers;
+}
+
+} // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createOutlineShapeComputationPass() {
+  return std::make_unique<OutlineShapeComputationPass>();
+}

mlir/test/Dialect/Shape/outline-shape-computation.mlir

+// RUN: mlir-opt -outline-shape-computation -split-input-file -allow-unregistered-dialect %s 2>&1 | FileCheck %s
+
+// two dynamic shapes: one of direct shape.shape_of(arg) and the other 
+func.func @main(%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) -> tensor<?x4x?xf32> {
+  %c2 = arith.constant 2 : index
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %0 = shape.shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+  %1 = shape.get_extent %0, %c2 : tensor<3xindex>, index -> index
+  %2 = "test.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
+  %3 = shape.with_shape %2, %0 : tensor<?x4x?xf32>, tensor<3xindex>
+  %4 = shape.value_of %3 : tensor<?x4x?xf32>
+  %5 = "test.concat"(%4, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+  %6 = shape.get_extent %0, %c0 : tensor<3xindex>, index -> index
+  %7 = arith.addi %6, %c2 : index
+  %8 = shape.from_extents %7, %c4, %1 : index, index, index
+  %9 = shape.with_shape %5, %8 : tensor<?x4x?xf32>, !shape.shape
+  %10 = shape.value_of %9 : tensor<?x4x?xf32>
+  return %10 : tensor<?x4x?xf32>
+}
+
+// CHECK-LABEL:  func.func @main
+// CHECK-NEXT:     %0 = "test.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
+// CHECK-NEXT:     %1 = "test.concat"(%0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+// CHECK-NEXT:     return %1 : tensor<?x4x?xf32>
+
+// CHECK: shape.func private @shape_cal_1(%arg0: tensor<?x4x?xf32>) -> !shape.shape {
+// CHECK-DAG:      %[[V0:.*]] = shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+// CHECK-DAG:      %[[V1:.*]] = get_extent %[[V0]], %c2 : tensor<3xindex>, index -> index
+// CHECK-DAG:      %[[V2:.*]] = get_extent %[[V0]], %c0 : tensor<3xindex>, index -> index
+// CHECK-DAG:      %[[V3:.*]] = arith.addi %[[V2]], %c2 : index
+// CHECK-DAG:      %[[V4:.*]] = from_extents %[[V3]], %c4, %[[V1]] : index, index, index
+// CHECK-DAG:      return %[[V4]] : !shape.shape
+
+// CHECK: shape.func private @shape_cal_0(%arg0: tensor<?x4x?xf32>) -> tensor<3xindex> {
+// CHECK-DAG:   %0 = shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+// CHECK-DAG:   return %0 : tensor<3xindex>
+
+// -----
+
+// two dynamic shapes and they share the same shape.func
+func.func @main(%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) -> tensor<?x4x?xf32> {
+  %c2 = arith.constant 2 : index
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %0 = shape.shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+  %1 = shape.get_extent %0, %c2 : tensor<3xindex>, index -> index
+  %2 = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+  %3 = shape.get_extent %0, %c0 : tensor<3xindex>, index -> index
+  %4 = arith.addi %3, %c2 : index
+  %5 = shape.from_extents %4, %c4, %1 : index, index, index
+  %6 = shape.with_shape %2, %5 : tensor<?x4x?xf32>, !shape.shape
+  %7 = shape.value_of %6 : tensor<?x4x?xf32>
+  %8 = "test.abs"(%7) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
+  %9 = shape.with_shape %8, %5 : tensor<?x4x?xf32>, !shape.shape
+  %10 = shape.value_of %9 : tensor<?x4x?xf32>
+  return %10 : tensor<?x4x?xf32>
+}
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT:     %0 = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+// CHECK-NEXT:     %1 = "test.abs"(%0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32>
+// CHECK-NEXT:     return %1 : tensor<?x4x?xf32>
+
+// CHECK:       shape.func private @shape_cal_0(%arg0: tensor<?x4x?xf32>) -> !shape.shape {
+// CHECK-DAG:     %[[V0:.*]] = shape_of %arg0 : tensor<?x4x?xf32> -> tensor<3xindex>
+// CHECK-DAG:     %[[V1:.*]] = get_extent %[[V0]], %c2 : tensor<3xindex>, index -> index
+// CHECK-DAG:     %[[V2:.*]] = get_extent %[[V0]], %c0 : tensor<3xindex>, index -> index
+// CHECK-DAG:     %[[V3:.*]] = arith.addi %[[V2]], %c2 : index
+// CHECK-DAG:     %[[V4:.*]] = from_extents %[[V3]], %c4, %[[V1]] : index, index, index
+// CHECK-DAG:     return %4 : !shape.shape
+// CHECK-NOT: shape_cal_1
+
+// -----
+
+// there's an internal dynamic shape source, and two other dynamic shapes shares it
+func.func @main(%arg0: tensor<?x4xf32>) -> tensor<?xi32> {
+  %0 = "test.nonzero"(%arg0) : (tensor<?x4xf32>) -> tensor<?xi32>
+  %1 = shape.shape_of %0 : tensor<?xi32> -> tensor<1xindex>
+  %2 = shape.with_shape %0, %1 : tensor<?xi32>, tensor<1xindex>
+  %3 = shape.value_of %2 : tensor<?xi32>
+  %4 = "test.abs"(%3) : (tensor<?xi32>) -> tensor<?xi32>
+  %5 = shape.with_shape %4, %1 : tensor<?xi32>, tensor<1xindex>
+  %6 = shape.value_of %5 : tensor<?xi32>
+  %7 = "test.negate"(%6) : (tensor<?xi32>) -> tensor<?xi32>
+  %8 = shape.with_shape %7, %1 : tensor<?xi32>, tensor<1xindex>
+  %9 = shape.value_of %8 : tensor<?xi32>
+  return %9 : tensor<?xi32>
+}
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT:     %0 = "test.nonzero"(%arg0) : (tensor<?x4xf32>) -> tensor<?xi32>
+// CHECK-NEXT:     %1 = "test.abs"(%0) : (tensor<?xi32>) -> tensor<?xi32>
+// CHECK-NEXT:     %2 = "test.negate"(%1) : (tensor<?xi32>) -> tensor<?xi32>
+// CHECK-NEXT:     return %2 : tensor<?xi32>
+
+// CHECK:      shape.func private @shape_cal_0(%arg0: tensor<?xi32>) -> tensor<1xindex> {
+// CHECK-NEXT:   %0 = shape_of %arg0 : tensor<?xi32> -> tensor<1xindex>
+// CHECK-NEXT:   return %0 : tensor<1xindex>
+// CHECK-NOT: shape_cal_1
+
+// -----
+
+// there's only a return op in the constructed shape.func
+func.func @main(%arg0: tensor<?x4xf32>, %arg1: tensor<1xindex>) -> tensor<?xi32> {
+  %0 = "test.nonzero"(%arg0) : (tensor<?x4xf32>) -> tensor<?xi32>
+  %1 = shape.with_shape %0, %arg1 : tensor<?xi32>, tensor<1xindex>
+  %2 = shape.value_of %1 : tensor<?xi32>
+  return %2 : tensor<?xi32>
+}
+// CHECK-LABEL: func.func @main(%arg0: tensor<?x4xf32>, %arg1: tensor<1xindex>) -> tensor<?xi32> {
+// CHECK-NEXT:   %0 = "test.nonzero"(%arg0) : (tensor<?x4xf32>) -> tensor<?xi32>
+// CHECK-NEXT:   return %0 : tensor<?xi32>
+
+// CHECK:      shape.func private @shape_cal_0(%arg0: tensor<1xindex>) -> tensor<1xindex> {
+// CHECK-NEXT:   return %arg0 : tensor<1xindex>
+
+// -----
+
+// shape computation part interleaves with general computation
+func.func @main(%arg0: tensor<?x4x5xf32>, %arg1: tensor<?x4x5xf32>, %arg2: tensor<?x4x5xf32>) -> (tensor<?x4x5xf32>, index) {
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %c5 = arith.constant 5 : index
+  %0 = shape.shape_of %arg0 : tensor<?x4x5xf32> -> tensor<3xindex>
+  %1 = shape.shape_of %arg1 : tensor<?x4x5xf32> -> tensor<3xindex>
+  %2 = shape.shape_of %arg2 : tensor<?x4x5xf32> -> tensor<3xindex>
+  %3 = "test.concat"(%arg0, %arg1, %arg2) {axis = 0 : i64} : (tensor<?x4x5xf32>, tensor<?x4x5xf32>, tensor<?x4x5xf32>) -> tensor<?x4x5xf32>
+  %4 = shape.get_extent %0, %c0 : tensor<3xindex>, index -> index
+  %5 = shape.get_extent %1, %c0 : tensor<3xindex>, index -> index
+  %6 = shape.get_extent %2, %c0 : tensor<3xindex>, index -> index
+  %7 = arith.addi %4, %5 : index
+  %8 = arith.addi %7, %6 : index
+  %9 = shape.from_extents %8, %c4, %c5 : index, index, index
+  %10 = shape.with_shape %3, %9 : tensor<?x4x5xf32>, !shape.shape
+  %11 = shape.value_of %10 : tensor<?x4x5xf32>
+  return %11, %7 : tensor<?x4x5xf32>, index
+}
+// CHECK-LABEL: func.func @main
+// CHECK-DAG:   %[[V0:.*]] = shape.shape_of %arg0 : tensor<?x4x5xf32> -> tensor<3xindex>
+// CHECK-DAG:   %[[V1:.*]] = shape.shape_of %arg1 : tensor<?x4x5xf32> -> tensor<3xindex>
+// CHECK-DAG:   %[[V2:.*]] = "test.concat"(%arg0, %arg1, %arg2) {axis = 0 : i64} : (tensor<?x4x5xf32>, tensor<?x4x5xf32>, tensor<?x4x5xf32>) -> tensor<?x4x5xf32>
+// CHECK-DAG:   %[[V3:.*]] = shape.get_extent %[[V0]], %c0 : tensor<3xindex>, index -> index
+// CHECK-DAG:   %[[V4:.*]] = shape.get_extent %[[V1]], %c0 : tensor<3xindex>, index -> index
+// CHECK-DAG:   %[[V5:.*]] = arith.addi %[[V3]], %[[V4]] : index
+// CHECK-DAG:   return %[[V2]], %[[V5]] : tensor<?x4x5xf32>, index
+
+// CHECK:     shape.func private @shape_cal_0(%arg0: tensor<?x4x5xf32>, %arg1: index, %arg2: index) -> !shape.shape {
+// CHECK-DAG:   %[[V0:.*]] = shape_of %arg0 : tensor<?x4x5xf32> -> tensor<3xindex>
+// CHECK-DAG:   %[[V1:.*]] = get_extent %[[V0]], %arg1 : tensor<3xindex>, index -> index
+// CHECK-DAG:   %[[V2:.*]] = arith.addi %arg2, %[[V1]] : index
+// CHECK-DAG:   %[[V3:.*]] = from_extents %[[V2]], %c4, %c5 : index, index, index
+// CHECK-DAG:   return %[[V3]] : !shape.shape
+
+// -----
+
+// there're multiple reused shape computations
+func.func @main(%arg0: tensor<?x4xf32>, %arg1: tensor<?x4xf32>) -> (tensor<?x4xf32>, tensor<?x4xf32>, tensor<?x4xf32>, tensor<?x4xf32>) {
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %0 = shape.shape_of %arg0 : tensor<?x4xf32> -> tensor<2xindex>
+  %1 = shape.shape_of %arg1 : tensor<?x4xf32> -> tensor<2xindex>
+  %2 = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+  %3 = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+  %4 = shape.get_extent %0, %c0 : tensor<2xindex>, index -> index
+  %5 = shape.get_extent %1, %c0 : tensor<2xindex>, index -> index
+  %6 = arith.addi %4, %5 : index
+  %7 = shape.from_extents %6, %c4 : index, index
+  %8 = shape.with_shape %2, %7 : tensor<?x4xf32>, !shape.shape
+  %9 = shape.with_shape %3, %7 : tensor<?x4xf32>, !shape.shape
+  %10 = shape.value_of %8 : tensor<?x4xf32>
+  %11 = shape.value_of %9 : tensor<?x4xf32>
+  %12 = "test.concat"(%arg0, %2) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+  %13 = "test.concat"(%arg0, %3) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+  %14 = arith.addi %6, %4 : index
+  %15 = shape.from_extents %14, %c4 : index, index
+  %16 = shape.with_shape %12, %15 : tensor<?x4xf32>, !shape.shape
+  %17 = shape.with_shape %13, %15 : tensor<?x4xf32>, !shape.shape
+  %18 = shape.value_of %16 : tensor<?x4xf32>
+  %19 = shape.value_of %17 : tensor<?x4xf32>
+  return %10, %11, %18, %19 : tensor<?x4xf32>, tensor<?x4xf32>, tensor<?x4xf32>, tensor<?x4xf32>
+}
+// CHECK-LABEL: func.func @main
+// CHECK-DAG:     %[[V0:.*]] = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+// CHECK-DAG:     %[[V1:.*]] = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+// CHECK-DAG:     %[[V2:.*]] = "test.concat"(%arg0, %[[V0]]) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+// CHECK-DAG:     %[[V3:.*]] = "test.concat"(%arg0, %[[V1]]) {axis = 0 : i64} : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32>
+// CHECK-DAG:     return %[[V0]], %[[V1]], %[[V2]], %[[V3]] : tensor<?x4xf32>, tensor<?x4xf32>, tensor<?x4xf32>, tensor<?x4xf32>
+  
+// CHECK:      shape.func private @shape_cal_1(%arg0: tensor<?x4xf32>, %arg1: tensor<?x4xf32>) -> !shape.shape {
+// CHECK-DAG:    %[[V0:.*]] = shape_of %arg0 : tensor<?x4xf32> -> tensor<2xindex>
+// CHECK-DAG:    %[[V1:.*]] = shape_of %arg1 : tensor<?x4xf32> -> tensor<2xindex>
+// CHECK-DAG:    %[[V2:.*]] = get_extent %[[V0]], %c0 : tensor<2xindex>, index -> index
+// CHECK-DAG:    %[[V3:.*]] = get_extent %[[V1]], %c0 : tensor<2xindex>, index -> index
+// CHECK-DAG:    %[[V4:.*]] = arith.addi %[[V2]], %[[V3]] : index
+// CHECK-DAG:    %[[V5:.*]] = arith.addi %[[V4]], %[[V2]] : index
+// CHECK-DAG:    %[[V6:.*]] = from_extents %[[V5]], %c4 : index, index
+// CHECK-DAG:    return %[[V6]] : !shape.shape
+
+// CHECK:     shape.func private @shape_cal_0(%arg0: tensor<?x4xf32>, %arg1: tensor<?x4xf32>) -> !shape.shape {
+// CHECK-DAG:   %[[V0:.*]] = shape_of %arg0 : tensor<?x4xf32> -> tensor<2xindex>
+// CHECK-DAG:   %[[V1:.*]] = shape_of %arg1 : tensor<?x4xf32> -> tensor<2xindex>
+// CHECK-DAG:   %[[V2:.*]] = get_extent %[[V0]], %c0 : tensor<2xindex>, index -> index
+// CHECK-DAG:   %[[V3:.*]] = get_extent %[[V1]], %c0 : tensor<2xindex>, index -> index
+// CHECK-DAG:   %[[V4:.*]] = arith.addi %[[V2]], %[[V3]] : index
+// CHECK-DAG:   %[[V5:.*]] = from_extents %[[V4]], %c4 : index, index
+// CHECK-DAG:   return %[[V5]] : !shape.shape

mlir/test/lib/Dialect/Shape/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRShapeTestPasses
   TestShapeFunctions.cpp
+  TestShapeMappingAnalysis.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
+  MLIRShapeOpsTransforms
   MLIRShapeDialect
   MLIRSupport
   )

mlir/test/lib/Dialect/Shape/TestShapeMappingAnalysis.cpp

+//===- TestShapeMappingInfo.cpp -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Shape/Analysis/ShapeMappingAnalysis.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+struct TestShapeMappingPass
+    : public PassWrapper<TestShapeMappingPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestShapeMappingPass)
+
+  StringRef getArgument() const final { return "test-print-shape-mapping"; }
+  StringRef getDescription() const final {
+    return "Print the contents of a constructed shape mapping information.";
+  }
+  void runOnOperation() override {
+    llvm::Optional<std::reference_wrapper<shape::ShapeMappingAnalysis>>
+        maybeAnalysis = getCachedAnalysis<shape::ShapeMappingAnalysis>();
+    if (maybeAnalysis.has_value())
+      maybeAnalysis.value().get().print(llvm::errs());
+    else
+      llvm::errs() << "No cached ShapeMappingAnalysis existed.";
+  }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestShapeMappingPass() {
+  PassRegistration<TestShapeMappingPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestPadFusion();
 void registerTestPDLByteCodePass();
 void registerTestPDLLPasses();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
+void registerTestShapeMappingPass();
 void registerTestSliceAnalysisPass();
 void registerTestTensorTransforms();
 void registerTestTilingInterface();
 void registerTestTopologicalSortAnalysisPass();
 void registerTestTransformDialectInterpreterPass();
 void registerTestVectorLowerings();
 void registerTestNvgpuLowerings();
 } // namespace test
#endif
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestPDLLPasses();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
+  mlir::test::registerTestShapeMappingPass();
   mlir::test::registerTestSliceAnalysisPass();
   mlir::test::registerTestTensorTransforms();
   mlir::test::registerTestTilingInterface();
   mlir::test::registerTestTopologicalSortAnalysisPass();
   mlir::test::registerTestTransformDialectInterpreterPass();
   mlir::test::registerTestVectorLowerings();
   mlir::test::registerTestNvgpuLowerings();
 }