[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)

399

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

23 ↗

(On Diff #463365)

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

23 ↗

(On Diff #463365)

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/IR/ShapeOps.td

def Shape_ShapeOfOp : Shape_Op<"shape_of",
 
   let extraClassDeclaration = [{
     // Returns when two result types are compatible for this op; method used by
     // InferTypeOpInterface
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
+def Shape_ValueOfOp : Shape_Op<"value_of", [NoSideEffect]> {
+  let summary = "Returns value of a !shape.value_shape operand";
+
+   let description = [{
+    The operation takes !shape.value_shape as an argument and it returns a
+    tensor, which represents the value of (value, shape) tuple of the argument.
+  }];
+
+  let arguments = (ins Shape_ValueShapeType:$arg);
+  let results = (outs AnyShaped:$result);
+
+  let assemblyFormat = "$arg attr-dict `:` type($result)";
+}
+
 def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [
     DeclareOpInterfaceMethods<CastOpInterface>, NoSideEffect
   ]> {
   let summary = "Casts between index types of the shape and standard dialect";
   let description = [{
     Converts a `shape.size` to a standard index. This operation and its
     inverse, `index_to_size`, facilitate index conversion between the standard
     and the shape dialect. The behavior is undefined for unknown and invalid
let description = [{
     This op need not be a refinement of the shape. In non-error cases the input
     ValueShape's value and shape are conformant and so too for the output, but
     the result may be less specified than `operand`'s shape as `shape` is
     merely used to construct the new ValueShape. If join behavior is desired
     then a join op should be used.
   }];
 
   let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$operand,
-                       Shape_ShapeType:$shape);
+                       Shape_ShapeOrExtentTensorType:$shape);
   let results = (outs Shape_ValueShapeType:$result);
 
   let assemblyFormat = "operands attr-dict `:` type($operand) `,` type($shape)";
 }
 
 def Shape_YieldOp : Shape_Op<"yield",
     [HasParent<"ReduceOp, FunctionLibraryOp">,
      NoSideEffect,
let description = [{
     function.
   }];
 
   let arguments = (ins SymbolNameAttr:$sym_name,
                        TypeAttrOf<FunctionType>:$function_type,
                        OptionalAttr<StrAttr>:$sym_visibility);
   let regions = (region AnyRegion:$body);
 
+  let builders = [OpBuilder<(ins
+    "StringRef":$name, "FunctionType":$type,
+    CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs,
+    CArg<"ArrayRef<DictionaryAttr>", "{}">:$argAttrs)
+  >];
+
   let extraClassDeclaration = [{
+    static FuncOp create(Location location, StringRef name, FunctionType type,
+                         ArrayRef<NamedAttribute> attrs = {});
+    static FuncOp create(Location location, StringRef name, FunctionType type,
+                         Operation::dialect_attr_range attrs);
+    static FuncOp create(Location location, StringRef name, FunctionType type,
+                         ArrayRef<NamedAttribute> attrs,
+                         ArrayRef<DictionaryAttr> argAttrs);
     //===------------------------------------------------------------------===//
     // CallableOpInterface
     //===------------------------------------------------------------------===//
 
     /// Returns the region on the current operation that is callable. This may
     /// return null in the case of an external callable object, e.g. an external
     /// function.
     ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); }

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 {
 /// Creates an instance of the ShapeToShapeLowering pass that legalizes Shape

 
 // 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 corresponding
+// symbols mapping to it.·
+std::unique_ptr<OperationPass<ModuleOp>>
+createOutlineShapeComputationPass(const std::string &entryFunc = "");
+
 //===----------------------------------------------------------------------===//
 // 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(denoted by shape.with_shape)
+    by adding shape.func and corresponding symbols mapping to it.
+
+    The symbol(s) is placed in the `encoding` attribute of RankedTensorType. If
+    the shape of the tensor is or equals to a dynamic shape source, it is a
+    single FlatSymbolRefAttr. Else it is a ArrayAttr, where the first element
+    represents corresponding shape function name, the others represents its 
+    arguments.
+
+    // TODO: support all functions after lenient call is supported.
+    Currently this pass is only limited to entry function. And internal
+    func.call is also not supported.
+    
+    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 = "tosa.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 = "tosa.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, [@shape_cal_0, @arg_0]> {
+      %0 = "tosa.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32, @arg_0>
+      %1 = "tosa.concat"(%0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32,
+            @arg_0>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32, [@shape_cal_0,
+            @arg_0]>
+      return %1 : tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+    }
+    shape.func private @shape_cal_0(%arg0: tensor<3xindex>) -> !shape.shape {
+      %c2 = arith.constant 2 : index
+      %c0 = arith.constant 0 : index
+      %c4 = arith.constant 4 : index
+      %0 = get_extent %arg0, %c2 : tensor<3xindex>, index -> index
+      %1 = get_extent %arg0, %c0 : tensor<3xindex>, index -> index
+      %2 = arith.addi %1, %c2 : index
+      %3 = from_extents %2, %c4, %0 : index, index, index
+      return %3 : !shape.shape
+    }
+    ```
+
+    For the above example, the shape computation is inlined in the input IR,
+    which is used for two values'(tosa.abs and tosa.concat) shape. And the shape
+    compuatation part is outlined in the output IR. The tosa.abs' output shape
+    equals to arg0's , so the encoding is a simple symbol @arg_0. The
+    tosa.concat's output shape is bound to a shape.func, the first symbol
+    @shape_cal_0 denotes the function name, and the second symbol @arg_0 denotes
+    its argument.
+  }];
+  let constructor = "mlir::createOutlineShapeComputationPass()";
+  let options = [
+    Option<"entryFunc", "entry-func", "std::string",
+            /*default=*/"\"main\"",
+            "Used to specify the entry-function.">,
+  ];
+  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 Arithmetic";
   let constructor = "mlir::createShapeToShapeLowering()";

mlir/lib/Dialect/Shape/IR/Shape.cpp

void FunctionLibraryOp::print(OpAsmPrinter &p) {
   p << " mapping ";
   p.printAttributeWithoutType(getMappingAttr());
 }
 
 //===----------------------------------------------------------------------===//
 // FuncOp
 //===----------------------------------------------------------------------===//
 
+FuncOp FuncOp::create(Location location, StringRef name, FunctionType type,
+                      ArrayRef<NamedAttribute> attrs) {
+  OpBuilder builder(location->getContext());
+  OperationState state(location, getOperationName());
+  FuncOp::build(builder, state, name, type, attrs);
+  return cast<FuncOp>(Operation::create(state));
+}
+FuncOp FuncOp::create(Location location, StringRef name, FunctionType type,
+                      Operation::dialect_attr_range attrs) {
+  SmallVector<NamedAttribute, 8> attrRef(attrs);
+  return create(location, name, type, llvm::makeArrayRef(attrRef));
+}
+FuncOp FuncOp::create(Location location, StringRef name, FunctionType type,
+                      ArrayRef<NamedAttribute> attrs,
+                      ArrayRef<DictionaryAttr> argAttrs) {
+  FuncOp func = create(location, name, type, attrs);
+  func.setAllArgAttrs(argAttrs);
+  return func;
+}
+
+void FuncOp::build(OpBuilder &builder, OperationState &state, StringRef name,
+                   FunctionType type, ArrayRef<NamedAttribute> attrs,
+                   ArrayRef<DictionaryAttr> argAttrs) {
+  state.addAttribute(FuncOp::getSymNameAttrName(state.name),
+                     builder.getStringAttr(name));
+  state.addAttribute(FuncOp::getFunctionTypeAttrName(state.name),
+                     TypeAttr::get(type));
+  state.attributes.append(attrs.begin(), attrs.end());
+  state.addRegion();
+
+  if (argAttrs.empty())
+    return;
+  assert(type.getNumInputs() == argAttrs.size());
+  function_interface_impl::addArgAndResultAttrs(builder, state, argAttrs,
+                                                /*resultAttrs=*/llvm::None);
+}
+
 ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) {
   auto buildFuncType =
       [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
          function_interface_impl::VariadicFlag,
          std::string &) { return builder.getFunctionType(argTypes, results); };
 
   return function_interface_impl::parseFunctionOp(
       parser, result, /*allowVariadic=*/false, buildFuncType);

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 "PassDetail.h"
+#include "mlir/Dialect/Func/IR/FuncOps.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>
+
+#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) {
+  assert(!cluster.empty() &&
+         "There must be at least one element in the cluster!");
+
+  ValueRange inputs = getInputsOfCluster(cluster);
+  auto fnType = b.getFunctionType(inputs.getTypes(), shape.getType());
+  b.setInsertionPointAfter(cluster[0]->getParentOp());
+  shape::FuncOp fnOp = b.create<shape::FuncOp>(loc, fnName, fnType);
+  Block *block = fnOp.addEntryBlock();
+  b.setInsertionPoint(block, block->end());
+  BlockAndValueMapping bvm;
+  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;
+}
+
+// Increment `idx` until find the next available symbol name
+std::string
+getNextAvailableSymbolName(const std::string &prefix, int &idx,
+                           std::unordered_set<std::string> &usedSymbolNames) {
+  std::string name = prefix + std::to_string(idx++);
+
+  while (usedSymbolNames.count(name)) {
+    name = prefix + std::to_string(idx++);
+  }
+  usedSymbolNames.insert(name);
+  return name;
+}
+
+// return argument index if `shape` is the output of a
+// shape.shape_of(func_arg), else return -1.
+llvm::Optional<int> getShapeOfFuncArgIdx(Value shape, func::FuncOp funcOp) {
+  shape::ShapeOfOp shapeOfOp = shape.getDefiningOp<shape::ShapeOfOp>();
+  if (shapeOfOp == nullptr)
+    return llvm::None;
+  Value inp = shapeOfOp.getArg();
+  for (int i = 0; i < int(funcOp.getNumArguments()); ++i) {
+    if (funcOp.getArgument(i) == inp)
+      return i;
+  }
+
+  return llvm::None;
+}
+
+void constructShapeFunc(
+    const std::vector<shape::WithOp> &allWithOps, MLIRContext *context,
+    DenseMap<Value, SmallVector<Operation *, 8>> &clusters,
+    SymbolTable &symbolTable, func::FuncOp funcOp,
+    std::unordered_set<std::string> &usedSymbolNames,
+    DenseMap<Value, FlatSymbolRefAttr> &dynShapeSrc2Symbol,
+    DenseMap<Value, SmallVector<Attribute>> &dynShape2ShapeFunc) {
+  std::string dynamicSourceNamePrefix = "s";
+  int dynamicSourceNameIdx = 0;
+  std::string shapeCalculationNamePrefix = "shape_cal_";
+  int shapeCalculationNameIdx = 0;
+  OpBuilder builder(context);
+
+  auto getOrConstructSymbolFromShape = [&](Value shape) {
+    auto symbolIt = dynShapeSrc2Symbol.find(shape);
+    if (symbolIt == dynShapeSrc2Symbol.end()) {
+      std::string name;
+      llvm::Optional<int> index = getShapeOfFuncArgIdx(shape, funcOp);
+      if (index.has_value())
+        name = "arg_" + std::to_string(index.value());
+      else
+        name = getNextAvailableSymbolName(
+            dynamicSourceNamePrefix, dynamicSourceNameIdx, usedSymbolNames);
+
+      auto symbol = FlatSymbolRefAttr::get(context, name);
+      dynShapeSrc2Symbol[shape] = symbol;
+      return symbol;
+    } else {
+      return symbolIt->second;
+    }
+  };
+
+  // Construct a shape function or a symbol for each cluster
+  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];
+    // The cluster is empty when the shape is equal to a dynamic shape source
+    if (cluster.empty()) {
+      FlatSymbolRefAttr symbol = getOrConstructSymbolFromShape(shape);
+      value.setType(RankedTensorType::get(rankedType.getShape(),
+                                          rankedType.getElementType(), symbol));
+      LLVM_DEBUG(llvm::dbgs()
+                 << "Symbol for " << shape << ": " << symbol << "\n");
+    } else {
+      auto it = dynShape2ShapeFunc.find(shape);
+      SmallVector<Attribute> symbols;
+      if (it == dynShape2ShapeFunc.end()) {
+        std::string name = getNextAvailableSymbolName(
+            shapeCalculationNamePrefix, shapeCalculationNameIdx,
+            usedSymbolNames);
+        Location loc = value.getLoc();
+        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);
+        symbols.push_back(symbol);
+        for (Value inp : inputs) {
+          FlatSymbolRefAttr argSymbol = getOrConstructSymbolFromShape(inp);
+          symbols.push_back(argSymbol);
+        }
+        dynShape2ShapeFunc[shape] = symbols;
+      } else {
+        symbols = it->second;
+      }
+      auto arrayAttr = ArrayAttr::get(context, symbols);
+      LLVM_DEBUG(llvm::dbgs()
+                 << "Symbol for " << shape << ": " << arrayAttr << "\n");
+      value.setType(RankedTensorType::get(
+          rankedType.getShape(), rankedType.getElementType(), arrayAttr));
+    }
+  }
+}
+
+struct OutlineShapeComputationPass
+    : public OutlineShapeComputationBase<OutlineShapeComputationPass> {
+
+  OutlineShapeComputationPass(const std::string &entryFunc)
+      : OutlineShapeComputationBase() {
+    this->entryFunc = entryFunc;
+  }
+
+  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);
+  // `usedSymbolNames` makes sure the created symbol is unique at Moudle level
+  std::unordered_set<std::string> usedSymbolNames;
+  DenseMap<Value, FlatSymbolRefAttr> dynShapeSrc2Symbol;
+  DenseMap<Value, SmallVector<Attribute>> dynShape2ShapeFunc;
+
+  moduleOp.walk([&](func::FuncOp funcOp) {
+    if (funcOp.getName() != entryFunc)
+      return;
+
+    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, funcOp,
+                       usedSymbolNames, dynShapeSrc2Symbol, dynShape2ShapeFunc);
+
+    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();
+
+    funcOp.setType(
+        FunctionType::get(context, funcOp.front().getArgumentTypes(),
+                          funcOp.front().getTerminator()->getOperandTypes()));
+  });
+}
+
+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 either the result
+// of shape.shape_of or function argument.
+void OutlineShapeComputationPass::getClusterFromValue(
+    Value shape, DenseMap<Value, DenseSet<Operation *>> &clusters) {
+  DenseSet<Operation *> cluster;
+
+  Operation *defOp = shape.getDefiningOp();
+  // defOp == nullptr means shape is the argument of the func op
+  if (nullptr == defOp) {
+    return;
+  }
+
+  DenseSet<Operation *> visited;
+  std::queue<Operation *> queue;
+  visited.insert(defOp);
+  queue.push(defOp);
+  while (!queue.empty()) {
+    Operation *op = queue.front();
+    queue.pop();
+    if (op->getNumOperands() == 0) {
+      cluster.insert(op);
+    } else if (llvm::isa<shape::ShapeOfOp>(op) &&
+               onlyUsedByWithShapes_.count(op)) {
+      // Stop when the op is type of shape.shape_of and only used by
+      // shape.with_shape ops
+      continue;
+    } else {
+      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(const std::string &entryFunc) {
+  return std::make_unique<OutlineShapeComputationPass>(entryFunc);
+}

mlir/lib/Dialect/Shape/Transforms/PassDetail.h

 //===- PassDetail.h - Shape Pass class details ------------------*- 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 DIALECT_SHAPE_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_SHAPE_TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 namespace bufferization {
 class BufferizationDialect;
 } // namespace bufferization
 
 namespace memref {
 class MemRefDialect;
 } // namespace memref
 
+namespace shape {
+class ShapeDialect;
+} // namespace shape
+
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/Shape/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // DIALECT_SHAPE_TRANSFORMS_PASSDETAIL_H_

mlir/test/Dialect/Shape/ops.mlir

func.func @test_parse_const_shape() {
   return
 }
 
 func.func @test_shape_of(%arg0: tensor<?xf32>) -> tensor<?xindex> {
   %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
 
+func.func @test_value_of(%arg0: !shape.value_shape) -> tensor<?xf32> {
+  %0 = shape.value_of %arg0 : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
 func.func @test_constraints() {
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
   %true = arith.constant true
   %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   %w1 = shape.cstr_eq %0, %1 : !shape.shape, !shape.shape
   %w2 = shape.const_witness true
   %w3 = shape.const_witness false

 // Testing invoking shape function from another. shape_equal_shapes is merely
 // a trivial helper function to invoke elsewhere.
 func.func @shape_equal_shapes(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
   %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
   %1 = shape.shape_of %b : !shape.value_shape -> !shape.shape
   %2 = "shape.meet"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   return %2 : !shape.shape
 }
+
 func.func @shape_with_shape(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
   %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
   %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
   %2 = call @shape_equal_shapes(%a, %1) : (!shape.value_shape, !shape.value_shape) -> !shape.shape
   return %2 : !shape.shape
 }
 
+func.func @shape_with_shape_extent_tensor_type(%a : tensor<?x?x?xf32>, %b : !shape.value_shape) -> !shape.value_shape {
+  %0 = shape.shape_of %a : tensor<?x?x?xf32> -> tensor<3xindex>
+  %1 = shape.with_shape %b, %0 : !shape.value_shape, tensor<3xindex>
+  return %1 : !shape.value_shape
+}
+
 func.func @any_on_shape(%a : !shape.shape, %b : !shape.shape, %c : !shape.shape)
     -> !shape.shape {
   %result = shape.any %a, %b, %c
       : !shape.shape, !shape.shape, !shape.shape -> !shape.shape
   return %result : !shape.shape
 }
 
 func.func @any_on_mixed(%a : tensor<?xindex>,

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

+// RUN: mlir-opt -outline-shape-computation="entry-func=main" -split-input-file -allow-unregistered-dialect %s | FileCheck %s
+
+// two dynamic shapes: a trivial symbol and a 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.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-SAME:   (%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) -> tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+// CHECK-NEXT:     %0 = "test.abs"(%arg0) : (tensor<?x4x?xf32>) -> tensor<?x4x?xf32, @arg_0>
+// CHECK-NEXT:     %1 = "test.concat"(%0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32, @arg_0>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+// CHECK-NEXT:     return %1 : tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+
+// CHECK-LABEL:   shape.func private @shape_cal_0
+// CHECK-SAME:     (%arg0: tensor<3xindex>) -> !shape.shape
+// CHECK-DAG:      %[[V0:.*]] = get_extent %arg0, %c2 : tensor<3xindex>, index -> index
+// CHECK-DAG:      %[[V1:.*]] = get_extent %arg0, %c0 : tensor<3xindex>, index -> index
+// CHECK-DAG:      %[[V2:.*]] = arith.addi %[[V1]], %c2 : index
+// CHECK-DAG:      %[[V3:.*]] = from_extents %[[V2]], %c4, %[[V0]] : index, index, index
+// CHECK-DAG:      return %[[V3]] : !shape.shape
+
+// -----
+
+// 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-SAME:  (%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) -> tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+// CHECK-NEXT:     %0 = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+// CHECK-NEXT:     %1 = "test.abs"(%0) : (tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>) -> tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+// CHECK-NEXT:     return %1 : tensor<?x4x?xf32, [@shape_cal_0, @arg_0]>
+
+// CHECK-LABEL:   shape.func private @shape_cal_0
+// CHECK-SAME:    (%arg0: tensor<3xindex>) -> !shape.shape
+// CHECK-DAG:       %[[V0:.*]] = get_extent %arg0, %c2 : tensor<3xindex>, index -> index
+// CHECK-DAG:       %[[V1:.*]] = get_extent %arg0, %c0 : tensor<3xindex>, index -> index
+// CHECK-DAG:       %[[V2:.*]] = arith.addi %[[V1]], %c2 : index
+// CHECK-DAG:       %[[V3:.*]] = from_extents %[[V2]], %c4, %[[V0]] : index, index, index
+// CHECK-DAG:       return %[[V3]] : !shape.shape
+
+
+// -----
+
+// 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-SAME:  (%arg0: tensor<?x4xf32>) -> tensor<?xi32, @s0>
+// CHECK-NEXT:     %0 = "test.nonzero"(%arg0) : (tensor<?x4xf32>) -> tensor<?xi32, @s0>
+// CHECK-NEXT:     %1 = "test.abs"(%0) : (tensor<?xi32, @s0>) -> tensor<?xi32, @s0>
+// CHECK-NEXT:     %2 = "test.negate"(%1) : (tensor<?xi32, @s0>) -> tensor<?xi32, @s0>
+// CHECK-NEXT:     return %2 : tensor<?xi32, @s0>