[MLIR] Introduce the type inference pass.

[arames]

Summary

Type Inference uses the existing forward dataflow analysis mechanisms, with a
lattice for the types, and the join relationship.

It is designed to be safe to run on existing IR: it will not modify types
wherever uninstructed.
Interfaces AllowsInputTypesRefinementInterface and
AllowsOuptputTypesRefinementInterface allow fine-grained control over what
input/output types can be specialized. The type inference pass will insert
type_specialization and type_relaxation ops to satisfy the specified
constraints.

Notes:

• Due to potentially arbitrary constraints on block argument types coming from parent blocks' terminators, block argument types are never specialized. Instead, type_specialization ops are inserted when appropriate.
• Support for input types refinement is fully controlled by AllowsInputTypesRefinementInterface.
• Support for ouptut types refinement is explicit if AllowsOuptputTypesRefinementInterface is implemented. If not, it is implicit if InferTypeOpInterface is implemented.

[arames]

mlir/lib/Transforms/TypeInference.cpp

187

This function will need to be updated following the decision on D110764.

[stephenneuendorffer]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2148 ↗

(On Diff #376020)

more specialized?

[stephenneuendorffer]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2107 ↗

(On Diff #376020)

I'd prefer an active void here. "relax_types"

2144–2145 ↗

(On Diff #376020)

specialize_types?

2152–2153 ↗

(On Diff #376020)

This seems a little awkward to me. 'join_types' and 'type_specialization' seem to work together to implement a 'cast' with a dummy %join argument to pass the type information? is it useful to have them separate? (perhaps to force multiple casts to be performed the same way?)

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

34–40

Should probably be documented that this needs to be a monotonic function on the type lattice?

181–182

"specialization" is used elsewhere. It would be good to consistently use one term or the other.

181–218

These interfaces seem to be able to take into account the current type of the input or output. However, they probably should be monotonic functions representing type constraints on the type lattice. Would it be better to have a less stateful interface that passed the current type and was required to be a pure function?

mlir/include/mlir/Transforms/Passes.td

706 ↗

(On Diff #376020)

Output?

707 ↗

(On Diff #376020)

What is implicit? That input and output types can both be refined?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2820 ↗

(On Diff #376020)

isLessSpecializedOrSame should probably be allowed and then canonicalized out?

mlir/lib/Transforms/TypeInference.cpp

2

InferTypes.cpp?

15

Can you add a picture of the type lattice?

107

Why is this a special case and not just another op that implements the right interfaces?

108

specialize

166

Perhaps this should just imply a default implementation of inferTypeOpInterface, to avoid a special case here?

182

Perhaps this should just imply a default implementation of inferTypeOpInterface, to avoid a special case here?

mlir/test/Transforms/type-inference.mlir

33–34

This seems weird because the type information is being propagated through join_types without its output type changing?

47–50

This also seems weird. why doesn't join_types output the join of its input types?

[rriddle]

Will try to take a look in the next few days.

mlir/lib/Transforms/TypeInference.cpp

20

I'd suggest avoiding any kind of dialect dependency, because it prevents this from being a "general" transformation. I'd drive the creation of cast operations through a dialect interface specific to this transformation, and allow for dialects to decide how to do things.

[stephenneuendorffer]

Have you tried this on any graph regions? I guess maybe the same question goes for ForwardDataflowAnalysis?

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2107 ↗

(On Diff #376020)

I'd prefer an active void here. "relax_types"

active voice :)

mlir/lib/Transforms/TypeInference.cpp

10–15

It seems that there is a critical assumption here: that static type information represents an abstract interpretation of the behavior of an op. Meaning that given the universe of values a particular op will always map input in the set to an output in the set regardless of the type. In this case specializing a type only represents the presence of more information about the op (i.e. a smaller set of possible values) and cannot change the behavior of an op. Certain kinds of ops are dangerous/would be affected by this transformation, like an op that outputs an integer given the number of input dimensions with unknown size.

[arames]

I went through the first round of comments.
I addressed the minor ones, and am working on the other open:

• Consider a different mechanism than specialize_type and relax_type ops.
• No special case for relax_type (if it is kept).
• Maybe SameOperandsAndResultType (and similar traits) imply a default implementation of InferTypeOpInterface.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2152–2153 ↗

(On Diff #376020)

I tried to use a short example taken from the tests, but it must be confusing and so not useful.
I both the "specialize" and "relax" examples.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

181–182

I updated the patch to use "specialize" and "relax" everywhere.

181–218

I was only aware that some exotic may need some control over this, but not to what extent. So I set it up this way to allow maximum control on the op side.
I am not sure looking at the current type would be enough here. For example, I can imagine the decision may depend on the input/output index, or operation attributes.

I think @silvas may want to weigh on this. I personally don't have a use for this, so would be fine restricting it.

mlir/include/mlir/Transforms/Passes.td

707 ↗

(On Diff #376020)

"support for output types refinement" is implicit.

I rewrote to clarify.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2820 ↗

(On Diff #376020)

Sounds fair.
It should not happen with the current usage, but I can see it being useful in the future. (For example if we somehow support updating block argument types separately.)

mlir/lib/Transforms/TypeInference.cpp

2

Looking at other filenames in that directory, there seem to be a mix of active and passive voice names, with a majority of passive.
For example:

• BufferDeallocation.cpp
• Canonicalizer.cpp
• LoopInvariantCodeMotion.cpp

But:

• Bufferize.cpp
• NormalizeMemRefs.cpp

Has a preference been established ? I'm happy to update to InferTypes.cpp and InferTypesPass if that's the established preference.

10–15

I believe the introduced the AllowsInputTypesSpecializationInterface and AllowsOutputTypesSpecializationInterface interfaces specifically address this.
By default, ops do NOT allow input or output type specialization (see details in the description for AllowsOutputTypesSpecializationInterface).

Could you check that addresses your concerns ?

20

Sounds interesting. I'll look into it.

107

I missed this as I was thinking of it as special, but it probably can simply use InferReturnTypes. Looking into it.

166

I'll look into it.
I suppose it would require the presence of the trait to imply implementation of the interface.
It should be relatively easy to do via tblgen, but is that something we already do ? or are fine to do ?

mlir/test/Transforms/type-inference.mlir

33–34

Yes. The behavior is intentional.

This follows the mechanisms to control type specialization for input types, that might be required for exotic ops.

For example I think the relax_type op itself (if kept) should be updated to use this instead of the special case:

• infer the result type to be the same as the input type
• but never specialize the output type

I expect most ops will simply rely on implicit support for output types specialization via InferTypeOpInterface, or maybe conditional via AllowsOutputTypesSpecializationInterface for specific outputs.

47–50

test.ti.join_types is a test op that "processes" its inputs, and somehow returns a result with a type that is the join of the input types.
Here we simply check that type-inference correctly propagates type information through type_relaxation (relax_type).
I agree the test as-is may be confusing. I reworked it to be clearer.

[arames]

Address minor comments. Still some work to do.

[silvas]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

181–218

These interfaces only determine whether an output or input is allowed to be refined. They don't actually mutate anything or calculate a type to refine to, so I don't think monotonicity concerns intersect with these two specific interfaces.

mlir/lib/Transforms/TypeInference.cpp

36

In torch-mlir, we found that DataFlowAnalysis.h miscompiles because it does not "meet" at control flow merge points: https://bugs.llvm.org/show_bug.cgi?id=51636

I suspect that this patch will be affected as well -- I don't think the analysis we want here can be implemented correctly without a "meet" callback on the lattice value as well.

CC @rriddle too

[silvas]

mlir/lib/Transforms/TypeInference.cpp

2

For me personally (take it with a grain of salt), I have been recently always using active voice (inferTypes.cpp) because it reads nicer on the mlir-opt command line (like a sequence of actions). I don't mind either way though.

[arames]

mlir/lib/Transforms/TypeInference.cpp

2

I was already using the term infer-types for the mlir-opt command.
I'll rename the file, pass, etc. to match this in the next update of the diff.

[arames]

mlir/lib/Transforms/TypeInference.cpp

166

I see this is already being done in tblgen. (For example code added in 31f40f603d0c0.).
So I'll extend that mechanism in a separate patch.

[arames]

mlir/lib/Transforms/TypeInference.cpp

36

Looking at the bug report, I am not understanding what is stated.
For type inference, it seems to me the dataflow analysis should always join. And the join of 3x? and ?x4 should be ?x?, not 3x4.
The refinement part becomes allowed because value types start with uninitialized, not their original type.
I test this for example in test_regionbranchopinterface in this patch.
I'll look at the bug. But I'd also be interested in chatting more.

[arames]

Address first round of comments.

[arames]

mlir/lib/Transforms/TypeInference.cpp

2

Looking at this again, it seems to me filenames and classes use nouns more often than "active voice" names.
I do use -infer-types for the pass argument though.
I'm open to change if there is more support for the other approach.

36

Here is an example. Imagine test.ti.join_types are data-processing ops that happen to have the same return type as the input.

func @test_branch_join(%cond: i1, %x : tensor<1x2xi32>, %y : tensor<1x9xi32>) {
  cond_br %cond, ^bb1, ^bb2
^bb1:
  %val_true = "test.ti.join_types"(%x) : (tensor<1x2xi32>) -> tensor<*xi32>
  br ^bb3(%val_true : tensor<*xi32>)
^bb2:
  %val_false = "test.ti.join_types"(%y) : (tensor<1x9xi32>) -> tensor<*xi32>
  br ^bb3(%val_false : tensor<*xi32>)
^bb3(%val: tensor<*xi32>):
  "test.ti.allow_input_type_specialization"(%val) : (tensor<*xi32>) -> ()
  "test.ti.disallow_input_type_specialization"(%val) : (tensor<*xi32>) -> ()
  return
}

% /bin/mlir-opt -infer-types ~/work/tmp/test.mlir
ninja: no work to do.
module  {
  func @test_branch_join(%arg0: i1, %arg1: tensor<1x2xi32>, %arg2: tensor<1x9xi32>) {
    cond_br %arg0, ^bb1, ^bb2
  ^bb1:  // pred: ^bb0
    %0 = "test.ti.join_types"(%arg1) : (tensor<1x2xi32>) -> tensor<1x2xi32>
    %1 = relax_type %0 : tensor<1x2xi32> to tensor<*xi32>
    br ^bb3(%1 : tensor<*xi32>)
  ^bb2:  // pred: ^bb0
    %2 = "test.ti.join_types"(%arg2) : (tensor<1x9xi32>) -> tensor<1x9xi32>
    %3 = relax_type %2 : tensor<1x9xi32> to tensor<*xi32>
    br ^bb3(%3 : tensor<*xi32>)
  ^bb3(%4: tensor<*xi32>):  // 2 preds: ^bb1, ^bb2
    %5 = specialize_type %4 : tensor<*xi32> to tensor<1x?xi32>
    "test.ti.allow_input_type_specialization"(%5) : (tensor<1x?xi32>) -> ()
    "test.ti.disallow_input_type_specialization"(%4) : (tensor<*xi32>) -> ()
    return
  }
}

We correctly see the join of tensor<1x2xi32> and tensor<1x9xi32> to tensor<1x?xi32>.
The operation required is a join, not a meet.

[silvas]

mlir/lib/Transforms/TypeInference.cpp

36

Yes, for control flow confluence what you want is to join (which can only derefine).

What I'm confused about is how your updateLatticeElement function end up refining. Using your example, I'd like to understand better how %val_true is treated.

Eventually it runs:

// latticeElement has initial value of tensor<*xi32>
latticeElement = getLatticeElement(%val_true);
// then you run:
latticeElement.join(tensor<1x2xi32>)
// Now latticeElement should still be tensor<*xi32> per join semantics?
// How does latticeElement get refined to the final value of tensor<1x2xi32>?

[arames]

mlir/lib/Transforms/TypeInference.cpp

36

I think I see what you missed.

// latticeElement has initial value of tensor<*xi32>
latticeElement = getLatticeElement(%val_true);

No it does not !
The current type of a value is only used (in the analysis) for ops that do not support InferTypeOpInterface. This is the point that allows refining types.

During the analysis:
All lattice elements start in the bottom state.
Then, there are two cases:

• an op does not support InferTypeOpInterface. In that case, we do join with the current type: join(uninitialized, current_type) -> current_type
• an op does support InferTypesOpInterface. Then we do not use the current type. Instead join with the result of the inferReturnTypes function.

When using the results of the analysis, we verify that each inferred type is more specialized than the current type.

[arames]

mlir/lib/Transforms/TypeInference.cpp

36

To go in the implementation details, after

// latticeElement has initial value of tensor<*xi32>
latticeElement = getLatticeElement(%val_true);

,
latticeElement.isUninitialized() is true. So in

latticeElement.join(value with type tensor<1x2xi32>)

LatticeElement::join will select the tensor type for optimisticValue.

I'll read the code again.
Here is a dump with the debug info. Updates to lattices elements for the control-flow are not visible because they happen in the dataflow analysis code.

% ./bin/mlir-opt -debug-only=infer-types -infer-types ~/work/tmp/test.mlir
update lattice element for :    %1 = "test.ti.join_types"(%arg2) : (tensor<1x9xi32>) -> tensor<*xi32>
        original type:  uninitialized
        joining with:   tensor<1x9xi32>
        yielding:       tensor<1x9xi32>
update lattice element for :    %0 = "test.ti.join_types"(%arg1) : (tensor<1x2xi32>) -> tensor<*xi32>
        original type:  uninitialized
        joining with:   tensor<1x2xi32>
        yielding:       tensor<1x2xi32>
updated type from tensor<*xi32> to tensor<1x2xi32> (with 0/1 specialized uses) for %0 = "test.ti.join_types"(%arg1) : (tensor<1x2xi32>) -> tensor<1x2xi32>
updated type from tensor<*xi32> to tensor<1x9xi32> (with 0/1 specialized uses) for %2 = "test.ti.join_types"(%arg2) : (tensor<1x9xi32>) -> tensor<1x9xi32>
updated type from tensor<*xi32> to tensor<1x?xi32> (with 1/1 specialized uses) for <block argument> of type 'tensor<*xi32>' at index: 0
module  {
  func @test_branch_join(%arg0: i1, %arg1: tensor<1x2xi32>, %arg2: tensor<1x9xi32>) {
    cond_br %arg0, ^bb1, ^bb2
  ^bb1:  // pred: ^bb0
    %0 = "test.ti.join_types"(%arg1) : (tensor<1x2xi32>) -> tensor<1x2xi32>
    %1 = relax_type %0 : tensor<1x2xi32> to tensor<*xi32>
    br ^bb3(%1 : tensor<*xi32>)
  ^bb2:  // pred: ^bb0
    %2 = "test.ti.join_types"(%arg2) : (tensor<1x9xi32>) -> tensor<1x9xi32>
    %3 = relax_type %2 : tensor<1x9xi32> to tensor<*xi32>
    br ^bb3(%3 : tensor<*xi32>)
  ^bb3(%4: tensor<*xi32>):  // 2 preds: ^bb1, ^bb2
    %5 = specialize_type %4 : tensor<*xi32> to tensor<1x?xi32>
    "test.ti.allow_input_type_specialization"(%5) : (tensor<1x?xi32>) -> ()
    return
  }
}

[arames]

Include changes to remove explicit support for Same* traits.
They will be handled in a separate patch by providing a default implementation of InferTypeOpInterface.

[arames]

NFC refactor.

[arames]

mlir/lib/Transforms/TypeInference.cpp

166

I removed the code to handle the different interfaces. It now only supports InferTypeOpInterface (or a default is not supported).
I'll add support for other interfaces in a separate patch, as you suggested, by supplying a default implementation of InferTypeOpInterface.

[silvas]

mlir/lib/Transforms/TypeInference.cpp

36

To go in the implementation details, after

// latticeElement has initial value of tensor<*xi32>
latticeElement = getLatticeElement(%val_true);
,
latticeElement.isUninitialized() is true. So in

What if the fixed-point is not reached in the first iteration? What if it takes multiple iterations to reach fixed-point? Won't isUninitialized return false and no further refinement will happen after the first iteration, so the true fixed-point won't be reached?

Or to put it another way, I view this as a pessimistic dataflow problem where you start with the current result type, and then work your way down the lattice to a more refined type.

From what you describe, the code currently starts at the bottom of the lattice, and moves up, which means that the first inferred type (which might not be the most refined possible) gets "stuck".

Oh, actually I may have answered my own question. The SCCP-like framework uses optimistic values for backedges, so any later visitations can only make the type less refined. I see.

Carry on.... sorry for the noise.

[arames]

mlir/lib/Transforms/TypeInference.cpp

36

From what you describe, the code currently starts at the bottom of the lattice, and moves up, which means that the first inferred type (which might not be the most refined possible) gets "stuck".

True.
I need to think more about cases that could hit this.

Or to put it another way, I view this as a pessimistic dataflow problem where you start with the current result type, and then work your way down the lattice to a more refined type.

Not sure how much of this is obvious already: I don't think you can walk down the lattice to a more refined type. Walking down the lattice means using the meet operation.
Let's imagine for a given value will see types t1, t2, and t3.
The inferred type we need is the most specialized type that is less specialized than either of them. It is by definition the join(t1, t2, t3).
We cannot refine the initial result type, because any meet operation could be incompatible with other observed types.

[silvas]

mlir/lib/Transforms/TypeInference.cpp

36

Not sure how much of this is obvious already: I don't think you can walk down the lattice to a more refined type. Walking down the lattice means using the meet operation.
Let's imagine for a given value will see types t1, t2, and t3.
The inferred type we need is the most specialized type that is less specialized than either of them. It is by definition the join(t1, t2, t3).
We cannot refine the initial result type, because any meet operation could be incompatible with other observed types.

Not sure I follow. It is totally possible for a result type to be statically written as tensor<*xf32> and then be refined to tensor<?x?xf32> and then be refined to tensor<2x2xf32> as we analyze the program further. These could all come from the same transfer function, but as the operand lattice values are refined, the corresponding result lattice values will get refined. This does not happen in the optimistic formulation we are using here though (since it only moves up the lattice, and uses its special handling of backedges optimistically to make it all work out).

True.
I need to think more about cases that could hit this.

Because of the optimistic backedge handling, the "stuck" value is guaranteed to be the most refined with the current framework. (assuming visitOperation only updates lattice values for results, and not for other unrelated values in the program)

[arames]

Rebase on top of tree.
Slightly factor out the creation of type specialization and relaxation ops in TI.

[arames]

mlir/lib/Transforms/TypeInference.cpp

176

Working on this right now.

[arames]

Introduce mechanism to customize explicit type specialization/relaxation.

[jpienaar]

(nice, just couple of fly by comments)

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1094 ↗

(On Diff #383421)

+1 (on my wishlist,, I told Alex I'd get to this too as it makes Python side gen needlessly complicated)

1103 ↗

(On Diff #383421)

Nit: Lets start with what the op does and then where introduced.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

192

And this is on the "flattened" form? E.g.., is variadics as in ODS considered

[arames]

Address nit for op descriptions.

[arames]

Reworked TI to be a utility class instead of a pass, removing dependency on standard ops.

[arames]

Upload the right patch.

mlir/include/mlir/Analysis/TypeInference.h

+//===- TypeInference.h - Type inference analysis for MLIR -------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// The type inference analysis uses the dataflow analysis infrastructure to
+// infer MLIR value types over the IR.
+//
+// The `join` and `meet` relationships over the type lattice are described more
+// in details in `JoinMeetTypeInterface.h`.
+// The type lattice looks like
+//
+//    join(t1, t2)
+//    tensor<?x?x3xf32>
+//    /              \
+//   /                \
+//  /                  \
+// t1                  t2
+// tensor<1x?x3xf32>   tensor<?x2x3xf32>
+//  \                  /
+//   \                /
+//    \              /
+//    meet(t1, t2)
+//    tensor<1x2x3xf32>
+// .
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_TYPEINFERENCE_H
+#define MLIR_ANALYSIS_TYPEINFERENCE_H
+
+#include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Interfaces/JoinMeetTypeInterface.h"
+#include "llvm/Support/Debug.h"
+
+namespace mlir {
+struct TypeLatticeValue {
+  TypeLatticeValue() = default;
+  TypeLatticeValue(Type type) : type(type) {}
+
+  static TypeLatticeValue getPessimisticValueState(MLIRContext *context) {
+    return TypeLatticeValue();
+  }
+
+  static TypeLatticeValue getPessimisticValueState(Value value) {
+    return value.getType();
+  }
+
+  static TypeLatticeValue join(const TypeLatticeValue &lhs,
+                               const TypeLatticeValue &rhs) {
+    return joinTypes(lhs.type, rhs.type);
+  }
+
+  bool operator==(const TypeLatticeValue &rhs) const {
+    return type == rhs.type;
+  }
+
+  Type type;
+};
+
+class TypeInferenceAnalysis : public ForwardDataFlowAnalysis<TypeLatticeValue> {
+public:
+  using ForwardDataFlowAnalysis<TypeLatticeValue>::ForwardDataFlowAnalysis;
+  TypeInferenceAnalysis(MLIRContext *context)
+      : ForwardDataFlowAnalysis(context), status(success()) {}
+  ~TypeInferenceAnalysis() override = default;
+
+  ChangeResult visitOperation(
+      Operation *op,
+      ArrayRef<LatticeElement<TypeLatticeValue> *> operands) override;
+
+  LogicalResult getStatus() const { return status; }
+
+private:
+  /// A wrapper allowing tracking updates for debugging purposes.
+  ChangeResult updateLatticeElement(Value value, Type ty);
+
+  LogicalResult status;
+};
+
+} // end namespace mlir
+
+#endif // MLIR_ANALYSIS_TYPEINFERENCE_H

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

let description = [{
     during op construction, verification or type inference.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     StaticInterfaceMethod<
       /*desc=*/[{Infer the return types that an op would generate.
 
       The method takes an optional location which, if set, will be used to
       report errors on. The operands and attributes correspond to those with
       which an Operation would be created (e.g., as used in Operation::create)
       and the regions of the op. Be aware that this method is supposed to be
       called with valid arguments, e.g., operands are verified, or it may result
       in an undefined behavior.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"inferReturnTypes",
       /*args=*/(ins "::mlir::MLIRContext *":$context,
                     "::llvm::Optional<::mlir::Location>":$location,
                     "::mlir::ValueRange":$operands,
                     "::mlir::DictionaryAttr":$attributes,
                     "::mlir::RegionRange":$regions,
                     "::llvm::SmallVectorImpl<::mlir::Type>&":$inferredReturnTypes)
InterfaceMethod<
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyResultShapes",
       /*args=*/(ins "::mlir::OpBuilder &":$builder,
         "::mlir::ReifiedRankedShapedTypeDims &":$reifiedReturnShapes)
     >
   ];
 }
 
+def AllowsInputTypesSpecializationInterface :
+    OpInterface<"AllowsInputTypesSpecializationInterface"> {
+  let description = [{
+    Interface to specify what inputs may be specialized.
+  }];
+  let cppNamespace = "::mlir";
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{Indicate whether the specified input type may be specialized}],
+      /*retTy=*/"bool",
+      /*methodName=*/"allowsInputTypeSpecialization",
+      /*args=*/(ins "unsigned":$operandIndex),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{return true;}]
+    >,
+  ];
+}
+
+def AllowsOutputTypesSpecializationInterface :
+    OpInterface<"AllowsOutputTypesSpecializationInterface"> {
+  let description = [{
+    Interface to specify what outputs may be specialized.
+
+    Support for output types specialization is by default conditional on support for
+    `InferTypeOpInterface`.
+    When implemented, this interface allows overriding this default to
+    explicitly specify what outputs may be specialized.
+  }];
+  let cppNamespace = "::mlir";
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{Indicate whether the specified output type can be specialized}],
+      /*retTy=*/"bool",
+      /*methodName=*/"allowsOutputTypeSpecialization",
+      /*args=*/(ins "unsigned":$outputIndex),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{return true;}]
+    >,
+  ];
+}
+
 #endif // MLIR_INFERTYPEOPINTERFACE

mlir/include/mlir/Transforms/TypeInference.h

+//===- TypeInference.h - Type inference pass --------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Declaration of the type inference utilities, using the results of the type
+// inference analysis to update value types through the IR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_TYPEINFERENCE_H
+#define MLIR_TRANSFORMS_TYPEINFERENCE_H
+
+#include "mlir/Analysis/TypeInference.h"
+#include "mlir/IR/Types.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir {
+using TypeInferenceTypeAdapter = std::function<Value(Value, Type)>;
+
+/// The `TypeInference` class runs the type inference analysis, and uses the
+/// results to update value types through the IR.
+///
+/// The operation interface `InferTypeOpInterface` allows expressing
+/// per-operation type inference logic.
+///
+/// Support for input types specialization is controlled by
+/// `AllowsInputTypesSpecializationInterface`.
+///
+/// Support for input types specialization is controlled by
+/// `AllowsInputTypesSpecializationInterface` and `InferTypeOpInterface`.
+/// If `AllowsOutputTypesSpecializationInterface` is implemented, it explicitly
+/// specifies what output types specialization is allowed.
+/// If `AllowsOutputTypesSpecializationInterface` is not implemented, output
+/// types specialization is allowed if `InferTypeOpInterface` is implemented.
+///
+/// To satisfy type constraints, the type inference pass may need to explicitly
+/// relax or specialize a type when it cannot be done in place.
+/// This requires the user to provide a "type adapter" fuction, to indicate how
+/// such type "reinterpretations" should be performed.
+/// An example custom adapter would be:
+/// ```
+/// Value customTypeAdapter(Value value, Type type) {
+///   OpBuilder builder(value.getContext());
+///   builder.setInsertionPointAfterValue(value);
+///   if (type.isa<TensorType>())
+///     return builder.createOrFold<tensor::CastOp>(value.getLoc(), type,
+///                                                 value);
+///   if (type.isa<CustomType>())
+///     return builder.createOrFold<CustomCastOp>(value.getLoc(), type, value);
+///   // We not expect other types in our use-case.
+///   return {};
+/// }
+/// ```
+///
+/// A concrete example is available via the `TestTypeInference` pass.
+/// For example, running it on the input
+///
+/// ```
+/// func @test_input_specialization(%x : tensor<1x?xi32>, %y : tensor<?x2xi32>)
+/// {
+///   %tmp = "test.ti.join_types"(%x, %y) : (tensor<1x?xi32>, tensor<?x2xi32>)
+///   -> tensor<*xi32> "test.ti.disallow_input_type_specialization"(%tmp) :
+///   (tensor<*xi32>) -> () "test.ti.allow_input_type_specialization"(%tmp) :
+///   (tensor<*xi32>) -> ()
+///   "test.ti.conditionally_allow_input_type_specialization"(%tmp, %tmp) :
+///   (tensor<*xi32>, tensor<*xi32>) -> () return
+/// }
+/// ```
+///
+/// will output
+///
+/// ```
+/// func @test_input_specialization(%arg0: tensor<1x?xi32>, %arg1:
+/// tensor<?x2xi32>) {
+///   %0 = "test.ti.join_types"(%arg0, %arg1) : (tensor<1x?xi32>,
+///   tensor<?x2xi32>) -> tensor<?x?xi32> %1 = test.ti.relax_type %0 :
+///   tensor<?x?xi32> to tensor<*xi32>
+///   "test.ti.disallow_input_type_specialization"(%1) : (tensor<*xi32>) -> ()
+///   "test.ti.allow_input_type_specialization"(%0) : (tensor<?x?xi32>) -> ()
+///   "test.ti.conditionally_allow_input_type_specialization"(%0, %1) :
+///   (tensor<?x?xi32>, tensor<*xi32>) -> () return
+/// }
+/// ```
+/// .
+///
+/// And running on the input
+///
+/// ```
+/// func @test_branch_join(%cond: i1, %x : tensor<1xi32>, %y : tensor<2xi32>) {
+///   cond_br %cond, ^bb1, ^bb2
+/// ^bb1:
+///   %val_true = "test.ti.join_types"(%x) : (tensor<1xi32>) -> tensor<*xi32>
+///   br ^bb3(%val_true : tensor<*xi32>)
+/// ^bb2:
+///   %val_false = "test.ti.join_types"(%y) : (tensor<2xi32>) -> tensor<*xi32>
+///   br ^bb3(%val_false : tensor<*xi32>)
+/// ^bb3(%val: tensor<*xi32>):
+///   "test.ti.allow_input_type_specialization"(%val) : (tensor<*xi32>) -> ()
+///   "test.ti.disallow_input_type_specialization"(%val) : (tensor<*xi32>) -> ()
+///   return
+/// }
+/// ```
+///
+/// will output
+///
+/// ```
+/// func @test_branch_join(%arg0: i1, %arg1: tensor<1xi32>, %arg2:
+/// tensor<2xi32>) {
+///   cond_br %arg0, ^bb1, ^bb2
+/// ^bb1:  // pred: ^bb0
+///   %0 = "test.ti.join_types"(%arg1) : (tensor<1xi32>) -> tensor<1xi32>
+///   %1 = test.ti.relax_type %0 : tensor<1xi32> to tensor<*xi32>
+///   br ^bb3(%1 : tensor<*xi32>)
+/// ^bb2:  // pred: ^bb0
+///   %2 = "test.ti.join_types"(%arg2) : (tensor<2xi32>) -> tensor<2xi32>
+///   %3 = test.ti.relax_type %2 : tensor<2xi32> to tensor<*xi32>
+///   br ^bb3(%3 : tensor<*xi32>)
+/// ^bb3(%4: tensor<*xi32>):  // 2 preds: ^bb1, ^bb2
+///   %5 = test.ti.specialize_type %4 : tensor<*xi32> to tensor<?xi32>
+///   "test.ti.allow_input_type_specialization"(%5) : (tensor<?xi32>) -> ()
+///   "test.ti.disallow_input_type_specialization"(%4) : (tensor<*xi32>) -> ()
+///   return
+/// ```
+/// .
+class TypeInference {
+public:
+  explicit TypeInference(Operation *op, TypeInferenceTypeAdapter typeAdapter)
+      : op(op), analysis(op->getContext()), typeAdapter(typeAdapter) {}
+
+  /// Analyzes and update types for the processed operation.
+  LogicalResult run() {
+    analysis.run(op);
+    if (failed(analysis.getStatus()))
+      return analysis.getStatus();
+    return updateOp(op);
+  }
+
+private:
+  LogicalResult updateOp(Operation *op);
+  LogicalResult updateBlock(Block &block);
+  LogicalResult updateOpResults(Operation *op);
+  LogicalResult updateType(Value result);
+
+  Operation *op;
+  TypeInferenceAnalysis analysis;
+  TypeInferenceTypeAdapter typeAdapter;
+};
+} // end namespace mlir
+
+#endif // MLIR_TRANSFORMS_TYPEINFERENCEUTILS_H

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   AffineAnalysis.cpp
   AffineStructures.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
   LinearTransform.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   NumberOfExecutions.cpp
   PresburgerSet.cpp
   SliceAnalysis.cpp
+  TypeInference.cpp
   Utils.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
   )
 
 add_mlir_library(MLIRAnalysis
   AliasAnalysis.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
   Liveness.cpp
   NumberOfExecutions.cpp
   SliceAnalysis.cpp
+  TypeInference.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 
   DEPENDS
   mlir-headers

mlir/lib/Analysis/TypeInference.cpp

+//===- TypeInference.cpp - Type inference analysis for MLIR -----*- 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/Analysis/TypeInference.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+
+#define DEBUG_TYPE "type-inference-analysis"
+
+using llvm::dbgs;
+
+mlir::ChangeResult mlir::TypeInferenceAnalysis::visitOperation(
+    Operation *op, ArrayRef<LatticeElement<TypeLatticeValue> *> operands) {
+  // Do not continue with the analysis if something went wrong.
+  // TODO: Having a way to interrupt the analysis on error would be
+  // convenient.
+  if (LLVM_UNLIKELY(failed(status)))
+    return ChangeResult::NoChange;
+
+  // FIXME: Have `SameOperandsAndResultType` and other similar traits provide a
+  // default implementation of `InferTypeOpInterface`, so that they are
+  // automatically handled here.
+  auto interface = dyn_cast<InferTypeOpInterface>(op);
+  if (!interface) {
+    // By default, simply use the current result types.
+    ChangeResult change = ChangeResult::NoChange;
+    for (Value result : op->getResults())
+      change |= updateLatticeElement(result, result.getType());
+    return change;
+  }
+
+  // Ephemerally override operand types with inferred types.
+  SmallVector<Type> originalOperandTys;
+  const unsigned numOperands = op->getNumOperands();
+  originalOperandTys.resize(numOperands);
+  for (unsigned i = 0; i != numOperands; ++i) {
+    auto operand = op->getOperand(i);
+    originalOperandTys[i] = operand.getType();
+    operand.setType(operands[i]->getValue().type);
+  }
+
+  // Infer types for the operation.
+  SmallVector<Type> inferredReturnTypes;
+  status = interface.inferReturnTypes(
+      op->getContext(), op->getLoc(), op->getOperands(),
+      op->getAttrDictionary(), op->getRegions(), inferredReturnTypes);
+
+  // Immediately restore original operand types.
+  for (unsigned i = 0; i != numOperands; ++i)
+    op->getOperand(i).setType(originalOperandTys[i]);
+
+  if (failed(status)) {
+    op->emitError("failed to infer types");
+    return ChangeResult::NoChange;
+  }
+
+  ChangeResult change = ChangeResult::NoChange;
+  for (unsigned int i = 0; i < op->getNumResults(); i++)
+    change |= updateLatticeElement(op->getResult(i), inferredReturnTypes[i]);
+  return change;
+}
+
+mlir::ChangeResult
+mlir::TypeInferenceAnalysis::updateLatticeElement(Value value, Type ty) {
+  LatticeElement<TypeLatticeValue> &latticeElement = getLatticeElement(value);
+  LLVM_DEBUG(dbgs() << "update lattice element for :\t"; value.print(dbgs());
+             dbgs() << "\n\toriginal type:\t";
+             if (latticeElement.isUninitialized()) dbgs() << "uninitialized";
+             else latticeElement.getValue().type.print(dbgs());
+             dbgs() << "\n\tjoining with:\t"; ty.print(dbgs()););
+  ChangeResult change = latticeElement.join(ty);
+  LLVM_DEBUG(dbgs() << "\n\tyielding:\t";
+             latticeElement.getValue().type.print(dbgs()); dbgs() << "\n");
+  return change;
+}

mlir/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTransforms
   LoopInvariantCodeMotion.cpp
   NormalizeMemRefs.cpp
   OpStats.cpp
   ParallelLoopCollapsing.cpp
   PipelineDataTransfer.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
+  TypeInference.cpp
   ViewOpGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRStandardOpsIncGen
   MLIRTransformsPassIncGen

mlir/lib/Transforms/TypeInference.cpp

+//===- TypeInference.cpp - Infer Types of MLIR operations -----------------===//
+//
+// 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/IR/TypeUtilities.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/Format.h"
+
+#include "mlir/Transforms/TypeInference.h"
+
+#define DEBUG_TYPE "type-inference"
+
+using namespace mlir;
+using llvm::dbgs;
+
+LogicalResult mlir::TypeInference::updateOp(Operation *op) {
+  for (auto &region : op->getRegions()) {
+    for (auto &block : region) {
+      if (failed(updateBlock(block)))
+        return failure();
+    }
+  }
+  return updateOpResults(op);
+}
+
+LogicalResult mlir::TypeInference::updateBlock(Block &block) {
+  for (auto blockArg : block.getArguments())
+    if (failed(updateType(blockArg)))
+      return failure();
+  for (auto &nestedOp : block)
+    if (failed(updateOp(&nestedOp)))
+      return failure();
+  return success();
+}
+
+LogicalResult mlir::TypeInference::updateOpResults(Operation *op) {
+  for (auto result : op->getResults())
+    if (failed(updateType(result)))
+      return failure();
+  return success();
+}
+
+LogicalResult mlir::TypeInference::updateType(Value value) {
+  LatticeElement<TypeLatticeValue> *latticeElement =
+      analysis.lookupLatticeElement(value);
+  if (!latticeElement)
+    return success();
+
+  Type originalTy = value.getType();
+  Type inferredTy = latticeElement->getValue().type;
+  if (inferredTy == originalTy)
+    return success();
+
+  Operation *definingOp = value.getDefiningOp();
+
+  if (!isMoreSpecialized(inferredTy, originalTy))
+    return emitError(value.getLoc(), "inferred type ")
+           << inferredTy << " is not more specialized than the original type"
+           << originalTy;
+
+  // Values and uses may or may not allow specializing input or output types.
+  // Keep track of the "value" with its original type, and after type
+  // specialization.
+  Value valueWithOriginalType = value;
+  Value valueWithInferredType = {};
+
+  // Attempt to specialize the output type.
+  auto interface =
+      dyn_cast_or_null<AllowsOutputTypesSpecializationInterface>(definingOp);
+  bool explicitlyAllowedOutputTypeSpecialization =
+      interface && interface.allowsOutputTypeSpecialization(
+                       value.cast<OpResult>().getResultNumber());
+  bool implicitlyAllowedOutputTypeSpecialization =
+      !interface && dyn_cast_or_null<InferTypeOpInterface>(definingOp);
+  if (explicitlyAllowedOutputTypeSpecialization ||
+      implicitlyAllowedOutputTypeSpecialization) {
+    value.setType(inferredTy);
+    valueWithOriginalType = {};
+    valueWithInferredType = value;
+  }
+
+  // Update uses with the appropriate "value", specialized or not.
+  unsigned nSpecializedUses = 0;
+  for (auto &use : llvm::make_early_inc_range(value.getUses())) {
+    auto interface =
+        dyn_cast<AllowsInputTypesSpecializationInterface>(use.getOwner());
+    bool allowsInputTypeSpecialization =
+        interface &&
+        interface.allowsInputTypeSpecialization(use.getOperandNumber());
+    if (allowsInputTypeSpecialization) {
+      if (!valueWithInferredType) {
+        valueWithInferredType = typeAdapter(value, inferredTy);
+        if (!valueWithInferredType)
+          return emitError(value.getLoc(),
+                           "type adapter failed to specialize the type of ")
+                 << value;
+      }
+      use.set(valueWithInferredType);
+      ++nSpecializedUses;
+    } else {
+      if (!valueWithOriginalType) {
+        valueWithOriginalType = typeAdapter(value, originalTy);
+        if (!valueWithOriginalType)
+          return emitError(value.getLoc(),
+                           "type adapter failed to relax the type of ")
+                 << value;
+      }
+      use.set(valueWithOriginalType);
+    }
+  }
+
+  LLVM_DEBUG(dbgs() << "updated type from "; originalTy.print(dbgs());
+             dbgs() << " to "; inferredTy.print(dbgs());
+             dbgs() << llvm::format(
+                 " (with %u/%u specialized uses) for ", nSpecializedUses,
+                 std::distance(value.getUses().begin(), value.getUses().end()));
+             value.print(dbgs()); dbgs() << "\n");
+
+  return success();
+}

mlir/test/Interfaces/ControlFlowInterfaces/operand-types.mlir

 ^bb0:
   %0 = "getValue"() : () -> tensor<1xi32>
   "test.br"(%0)[^bb1] : (tensor<1xi32>) -> ()
   // expected-error@-1 {{type mismatch for bb argument #0 of successor #0}}
 ^bb1(%arg1: tensor<?xi32>):
   return
 
 ^bb2:
-  "test.cfe.br"(%0)[^bb3] : (tensor<1xi32>) -> ()
+  "test.ti.br"(%0)[^bb3] : (tensor<1xi32>) -> ()
 ^bb3(%arg2: tensor<?xi32>):
   return
 }
 
 // -----
 
 func @succ_type_dynamic_to_fixed() {
 ^bb0:
   %0 = "getValue"() : () -> tensor<?xi32>
-  "test.cfe.br"(%0)[^bb1] : (tensor<?xi32>) -> ()
+  "test.ti.br"(%0)[^bb1] : (tensor<?xi32>) -> ()
   // expected-error@-1 {{type mismatch for bb argument #0 of successor #0}}
 ^bb1(%arg3: tensor<1xi32>):
   return
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
%tmp1 = test.region_if %0 : memref<1xi32> -> (memref<?xi32>) then {
     ^bb0(%arg1 : memref<?xi32>):
       %false_value = "getValue"(%arg1) : (memref<?xi32>) -> (memref<3xi32>)
       test.region_if_yield %false_value : memref<3xi32>
   } join {
     ^bb0(%arg1 : memref<?xi32>):
       test.region_if_yield %arg1 : memref<?xi32>
   }
 
-  %tmp2 = test.cfe.region_if %0 : memref<1xi32> -> (memref<?xi32>) then {
+  %tmp2 = test.ti.region_if %0 : memref<1xi32> -> (memref<?xi32>) then {
     ^bb0(%arg1 : memref<1xi32>):
       %true_value = "getValue"(%arg1) : (memref<1xi32>) -> (memref<2xi32>)
-      test.cfe.region_if_yield %true_value : memref<2xi32>
+      test.ti.region_if_yield %true_value : memref<2xi32>
   } else {
     ^bb0(%arg1 : memref<?xi32>):
       %false_value = "getValue"(%arg1) : (memref<?xi32>) -> (memref<3xi32>)
-      test.cfe.region_if_yield %false_value : memref<3xi32>
+      test.ti.region_if_yield %false_value : memref<3xi32>
   } join {
     ^bb0(%arg1 : memref<?xi32>):
-      test.cfe.region_if_yield %arg1 : memref<?xi32>
+      test.ti.region_if_yield %arg1 : memref<?xi32>
   }
 
   return
 }

mlir/test/Transforms/type-inference.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -test-type-inference %s -split-input-file -o %t1
+// RUN: FileCheck %s < %t1
+// Verify idempotence.
+// RUN: mlir-opt -allow-unregistered-dialect -test-type-inference %t1 -split-input-file -o %t2
+// RUN: FileCheck %s < %t2
+
+func @test_input_specialization(%x : tensor<1x?xi32>, %y : tensor<?x2xi32>) {
+  %tmp = "test.ti.join_types"(%x, %y) : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<*xi32>
+  "test.ti.disallow_input_type_specialization"(%tmp) : (tensor<*xi32>) -> ()
+  "test.ti.allow_input_type_specialization"(%tmp) : (tensor<*xi32>) -> ()
+  "test.ti.conditionally_allow_input_type_specialization"(%tmp, %tmp) : (tensor<*xi32>, tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK: [[TMP:%.+]] = "test.ti.join_types"{{.+}} : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<?x?xi32>
+// CHECK: [[RELAXED_TMP:%.+]] = test.ti.relax_type [[TMP]] : tensor<?x?xi32> to tensor<*xi32>
+// CHECK: "test.ti.disallow_input_type_specialization"([[RELAXED_TMP]]) : (tensor<*xi32>) -> ()
+// CHECK: "test.ti.allow_input_type_specialization"([[TMP]]) : (tensor<?x?xi32>) -> ()
+// CHECK: "test.ti.conditionally_allow_input_type_specialization"([[TMP]], [[RELAXED_TMP]]) : (tensor<?x?xi32>, tensor<*xi32>) -> ()
+
+// -----
+
+func @test_output_specialization(%x : tensor<1x?xi32>, %y : tensor<?x2xi32>) {
+  %join1 = "test.ti.join_types"(%x, %y) { allowOutputTypeSpecialization = false } : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<*xi32>
+  "test.ti.allow_input_type_specialization"(%join1) : (tensor<*xi32>) -> ()
+
+  %join2 = "test.ti.join_types"(%x, %y) { allowOutputTypeSpecialization = true } : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<*xi32>
+  "test.ti.allow_input_type_specialization"(%join2) : (tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK: [[JOIN1:%.+]] = "test.ti.join_types"({{.+}}) {allowOutputTypeSpecialization = false} : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<*xi32>
+// CHECK: [[SPECIALIZED_JOIN1:%.+]] = test.ti.specialize_type [[JOIN1]] : tensor<*xi32> to tensor<?x?xi32>
+// CHECK: "test.ti.allow_input_type_specialization"([[SPECIALIZED_JOIN1]]) : (tensor<?x?xi32>) -> ()
+// CHECK: [[JOIN2:%.+]] = "test.ti.join_types"({{.+}}) {allowOutputTypeSpecialization = true} : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<?x?xi32>
+// CHECK: "test.ti.allow_input_type_specialization"([[JOIN2]]) : (tensor<?x?xi32>) -> ()
+
+// -----
+
+func @test_relaxation_propagation(%value : tensor<1xi32>) {
+  %relaxed_value = test.ti.relax_type %value : tensor<1xi32> to tensor<*xi32>
+  "test.br"(%relaxed_value)[^bb1] : (tensor<*xi32>) -> ()
+^bb1(%arg: tensor<*xi32>):
+  "test.ti.allow_input_type_specialization"(%arg) : (tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK: func @test_relaxation_propagation([[VALUE:%.+]]: tensor<1xi32>) {
+// CHECK:   [[RELAXED_VALUE:%.+]] = test.ti.relax_type [[VALUE]] : tensor<1xi32> to tensor<*xi32>
+// CHECK:   "test.br"([[RELAXED_VALUE]])[^bb1] : (tensor<*xi32>) -> ()
+// CHECK: ^bb1([[ARG:%.+]]: tensor<*xi32>):  // pred: ^bb0
+// CHECK:   [[SPECIALIZED_ARG:%.+]] = test.ti.specialize_type [[ARG]] : tensor<*xi32> to tensor<1xi32>
+// CHECK:   "test.ti.allow_input_type_specialization"([[SPECIALIZED_ARG]]) : (tensor<1xi32>) -> ()
+// CHECK:   return
+// CHECK: }
+
+// -----
+
+func @test_branch_disallowing_input_specialization(%x : tensor<1x?xi32>, %y : tensor<?x2xi32>) {
+  %tmp = "test.ti.join_types"(%x, %y) : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<*xi32>
+  "test.br"(%tmp)[^bb1] : (tensor<*xi32>) -> ()
+^bb1(%arg: tensor<*xi32>):
+  "test.ti.disallow_input_type_specialization"(%arg) : (tensor<*xi32>) -> ()
+  "test.ti.allow_input_type_specialization"(%arg) : (tensor<*xi32>) -> ()
+  "test.ti.conditionally_allow_input_type_specialization"(%arg, %arg) : (tensor<*xi32>, tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK:   [[TMP:%.+]] = "test.ti.join_types"({{.+}}) : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<?x?xi32>
+// CHECK:   [[RELAXED_TMP:%.+]] = test.ti.relax_type [[TMP]] : tensor<?x?xi32> to tensor<*xi32>
+// CHECK:   "test.br"([[RELAXED_TMP]])[^bb1] : (tensor<*xi32>) -> ()
+// CHECK: ^bb1([[ARG:%.+]]: tensor<*xi32>):  // pred: ^bb0
+// CHECK:   [[SPECIALIZED_ARG:%.+]] = test.ti.specialize_type [[ARG]] : tensor<*xi32> to tensor<?x?xi32>
+// CHECK:   "test.ti.disallow_input_type_specialization"([[ARG]]) : (tensor<*xi32>) -> ()
+// CHECK:   "test.ti.allow_input_type_specialization"([[SPECIALIZED_ARG]]) : (tensor<?x?xi32>) -> ()
+// CHECK:   "test.ti.conditionally_allow_input_type_specialization"([[SPECIALIZED_ARG]], [[ARG]]) : (tensor<?x?xi32>, tensor<*xi32>) -> ()
+// CHECK:   return
+
+// -----
+
+func @test_branch_allowing_input_specialization(%x : tensor<1x?xi32>, %y : tensor<?x2xi32>) {
+  %tmp = "test.ti.join_types"(%x, %y) : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<*xi32>
+  "test.ti.br"(%tmp)[^bb1] : (tensor<*xi32>) -> ()
+^bb1(%arg: tensor<*xi32>):
+  "test.ti.disallow_input_type_specialization"(%arg) : (tensor<*xi32>) -> ()
+  "test.ti.allow_input_type_specialization"(%arg) : (tensor<*xi32>) -> ()
+  "test.ti.conditionally_allow_input_type_specialization"(%arg, %arg) : (tensor<*xi32>, tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK:   [[TMP:%.+]] = "test.ti.join_types"({{.+}}) : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<?x?xi32>
+// CHECK:   "test.ti.br"([[TMP]])[^bb1] : (tensor<?x?xi32>) -> ()
+// CHECK: ^bb1([[ARG:%.+]]: tensor<*xi32>):  // pred: ^bb0
+// CHECK:   [[SPECIALIZED_ARG:%.+]] = test.ti.specialize_type [[ARG]] : tensor<*xi32> to tensor<?x?xi32>
+// CHECK:   "test.ti.disallow_input_type_specialization"([[ARG]]) : (tensor<*xi32>) -> ()
+// CHECK:   "test.ti.allow_input_type_specialization"([[SPECIALIZED_ARG]]) : (tensor<?x?xi32>) -> ()
+// CHECK:   "test.ti.conditionally_allow_input_type_specialization"([[SPECIALIZED_ARG]], [[ARG]]) : (tensor<?x?xi32>, tensor<*xi32>) -> ()
+// CHECK:   return
+
+// -----
+
+func @test_branch_join(%cond: i1, %x : tensor<1xi32>, %y : tensor<2xi32>) {
+  cond_br %cond, ^bb1, ^bb2
+^bb1:
+  %val_true = "test.ti.join_types"(%x) : (tensor<1xi32>) -> tensor<*xi32>
+  br ^bb3(%val_true : tensor<*xi32>)
+^bb2:
+  %val_false = "test.ti.join_types"(%y) : (tensor<2xi32>) -> tensor<*xi32>
+  br ^bb3(%val_false : tensor<*xi32>)
+^bb3(%val: tensor<*xi32>):
+  "test.ti.allow_input_type_specialization"(%val) : (tensor<*xi32>) -> ()
+  "test.ti.disallow_input_type_specialization"(%val) : (tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK: func @test_branch_join([[COND:%.+]]: i1, [[X:%.+]]: tensor<1xi32>, [[Y:%.+]]: tensor<2xi32>) {
+// CHECK:   cond_br [[COND]], ^bb1, ^bb2
+// CHECK: ^bb1:
+// CHECK:   [[JOIN:%.+]] = "test.ti.join_types"([[X]]) : (tensor<1xi32>) -> tensor<1xi32>
+// CHECK:   [[RELAXED_JOIN:%.+]] = test.ti.relax_type [[JOIN]] : tensor<1xi32> to tensor<*xi32>
+// CHECK:   br ^bb3([[RELAXED_JOIN]] : tensor<*xi32>)
+// CHECK: ^bb2:
+// CHECK:   [[JOIN:%.+]] = "test.ti.join_types"([[Y]]) : (tensor<2xi32>) -> tensor<2xi32>
+// CHECK:   [[RELAXED_JOIN:%.+]] = test.ti.relax_type [[JOIN]] : tensor<2xi32> to tensor<*xi32>
+// CHECK:   br ^bb3([[RELAXED_JOIN]] : tensor<*xi32>)
+// CHECK: ^bb3([[VAL:%.+]]: tensor<*xi32>):
+// CHECK:   [[SPECIALIZED_VAL:%.+]] = test.ti.specialize_type [[VAL]] : tensor<*xi32> to tensor<?xi32>
+// CHECK:   "test.ti.allow_input_type_specialization"([[SPECIALIZED_VAL]]) : (tensor<?xi32>) -> ()
+// CHECK:   "test.ti.disallow_input_type_specialization"([[VAL]]) : (tensor<*xi32>) -> ()
+// CHECK:   return
+// CHECK: }
+// -----
+
+func @test_regionbranchopinterface() -> tensor<?x?xi32> {
+  %cond = "getValue"() : () -> (i1)
+
+  %res = test.region_if %cond : i1 -> (tensor<?x?xi32>) then {
+    ^bb0(%arg1 : i1):
+      %tmp1 = "getValue"() : () -> (tensor<1x2xi32>)
+      %then_value = "test.ti.join_types"(%tmp1) : (tensor<1x2xi32>) -> tensor<?x?xi32>
+      test.region_if_yield %then_value : tensor<?x?xi32>
+  } else {
+    ^bb0(%arg1 : i1):
+      %tmp1 = "getValue"() : () -> (tensor<1x9xi32>)
+      %else_value = "test.ti.join_types"(%tmp1) : (tensor<1x9xi32>) -> tensor<?x?xi32>
+      test.region_if_yield %else_value : tensor<?x?xi32>
+  } join {
+    ^bb0(%arg1 : tensor<?x?xi32>):
+      test.region_if_yield %arg1 : tensor<?x?xi32>
+  }
+  "test.ti.allow_input_type_specialization"(%res) : (tensor<?x?xi32>) -> ()
+  return %res : tensor<?x?xi32>
+}
+
+// CHECK: [[COND:%.+]] = "getValue"() : () -> i1
+// CHECK: [[RES:%.+]] = test.region_if [[COND]]: i1 -> tensor<?x?xi32> then {
+// CHECK: ^bb0({{.+}}: i1):  // no predecessors
+// CHECK:   [[TMP1:%.+]] = "getValue"() : () -> tensor<1x2xi32>
+// CHECK:   [[THEN_VALUE:%.+]] = "test.ti.join_types"([[TMP1]]) : (tensor<1x2xi32>) -> tensor<1x2xi32>
+// CHECK:   [[RELAXED_THEN_VALUE:%.+]] = test.ti.relax_type [[THEN_VALUE]] : tensor<1x2xi32> to tensor<?x?xi32>
+// CHECK:   test.region_if_yield [[RELAXED_THEN_VALUE]] : tensor<?x?xi32>
+// CHECK: } else {
+// CHECK: ^bb0({{.+}}: i1):  // no predecessors
+// CHECK:   [[TMP2:%.+]] = "getValue"() : () -> tensor<1x9xi32>
+// CHECK:   [[ELSE_VALUE:%.+]] = "test.ti.join_types"([[TMP2]]) : (tensor<1x9xi32>) -> tensor<1x9xi32>
+// CHECK:   [[RELAXED_ELSE_VALUE:%.+]] = test.ti.relax_type [[ELSE_VALUE]] : tensor<1x9xi32> to tensor<?x?xi32>
+// CHECK:   test.region_if_yield [[RELAXED_ELSE_VALUE]] : tensor<?x?xi32>
+// CHECK: } join {
+// CHECK: ^bb0([[ARG:%.+]]: tensor<?x?xi32>):  // no predecessors
+// CHECK:   test.region_if_yield [[ARG]] : tensor<?x?xi32>
+// CHECK: }
+// CHECK: [[SPECIALIZED_RES:%.+]] = test.ti.specialize_type [[RES]] : tensor<?x?xi32> to tensor<1x?xi32>
+// CHECK: "test.ti.allow_input_type_specialization"([[SPECIALIZED_RES]]) : (tensor<1x?xi32>) -> ()
+// CHECK: return [[RES]] : tensor<?x?xi32>
+
+// -----
+
+func @test_regionbranchopinterface() -> tensor<?x?xi32> {
+  %cond = "getValue"() : () -> (i1)
+
+  %res = test.ti.region_if %cond : i1 -> (tensor<?x?xi32>) then {
+    ^bb0(%arg1 : i1):
+      %tmp1 = "getValue"() : () -> (tensor<1x2xi32>)
+      %then_value = "test.ti.join_types"(%tmp1) : (tensor<1x2xi32>) -> tensor<?x?xi32>
+      test.ti.region_if_yield %then_value : tensor<?x?xi32>
+  } else {
+    ^bb0(%arg1 : i1):
+      %tmp1 = "getValue"() : () -> (tensor<1x9xi32>)
+      %else_value = "test.ti.join_types"(%tmp1) : (tensor<1x9xi32>) -> tensor<?x?xi32>
+      test.ti.region_if_yield %else_value : tensor<?x?xi32>
+  } join {
+    ^bb0(%arg1 : tensor<?x?xi32>):
+      test.ti.region_if_yield %arg1 : tensor<?x?xi32>
+  }
+  "test.ti.allow_input_type_specialization"(%res) : (tensor<?x?xi32>) -> ()
+  return %res : tensor<?x?xi32>
+}
+
+// CHECK: [[COND:%.+]] = "getValue"() : () -> i1
+// CHECK: [[RES:%.+]] = test.ti.region_if [[COND]]: i1 -> tensor<1x?xi32> then {
+// CHECK: ^bb0({{.+}}: i1):  // no predecessors
+// CHECK:   [[TMP1:%.+]] = "getValue"() : () -> tensor<1x2xi32>
+// CHECK:   [[THEN_VALUE:%.+]] = "test.ti.join_types"([[TMP1]]) : (tensor<1x2xi32>) -> tensor<1x2xi32>
+// CHECK:   test.ti.region_if_yield [[THEN_VALUE]] : tensor<1x2xi32>
+// CHECK: } else {
+// CHECK: ^bb0({{.+}}: i1):  // no predecessors
+// CHECK:   [[TMP2:%.+]] = "getValue"() : () -> tensor<1x9xi32>
+// CHECK:   [[ELSE_VALUE:%.+]] = "test.ti.join_types"([[TMP2]]) : (tensor<1x9xi32>) -> tensor<1x9xi32>
+// CHECK:   test.ti.region_if_yield [[ELSE_VALUE]] : tensor<1x9xi32>
+// CHECK: } join {
+// CHECK: ^bb0([[ARG:%.+]]: tensor<?x?xi32>):  // no predecessors
+// CHECK:   [[SPECIALIZED_ARG:%.+]] = test.ti.specialize_type [[ARG]] : tensor<?x?xi32> to tensor<1x?xi32>
+// CHECK:   test.ti.region_if_yield [[SPECIALIZED_ARG]] : tensor<1x?xi32>
+// CHECK: }
+// CHECK: [[RELAXED_RES:%.+]] = test.ti.relax_type [[RES]] : tensor<1x?xi32> to tensor<?x?xi32>
+// CHECK: "test.ti.allow_input_type_specialization"([[RES]]) : (tensor<1x?xi32>) -> ()
+// CHECK: return [[RELAXED_RES]] : tensor<?x?xi32>
+
+// -----
+
+func @test_dynamicize(%arg : tensor<1x2x3xi32>) {
+  %tmp1 = "test.ti.dynamicize"(%arg) : (tensor<1x2x3xi32>) -> tensor<*xi32>
+  %tmp2 = "test.ti.dynamicize"(%tmp1) : (tensor<*xi32>) -> tensor<*xi32>
+  %tmp3 = "test.ti.dynamicize"(%tmp2) : (tensor<*xi32>) -> tensor<*xi32>
+  %tmp4 = "test.ti.dynamicize"(%tmp3) : (tensor<*xi32>) -> tensor<*xi32>
+  return
+}
+
+// CHECK: func @test_dynamicize([[ARG:%.+]]: tensor<1x2x3xi32>) {
+// CHECK: [[TMP1:%.+]] = "test.ti.dynamicize"([[ARG]]) : (tensor<1x2x3xi32>) -> tensor<?x2x3xi32>
+// CHECK: [[TMP2:%.+]] = "test.ti.dynamicize"([[TMP1]]) : (tensor<?x2x3xi32>) -> tensor<?x?x3xi32>
+// CHECK: [[TMP3:%.+]] = "test.ti.dynamicize"([[TMP2]]) : (tensor<?x?x3xi32>) -> tensor<?x?x?xi32>
+// CHECK: [[TMP4:%.+]] = "test.ti.dynamicize"([[TMP3]]) : (tensor<?x?x?xi32>) -> tensor<?x?x?xi32>
+
+// -----
+
+func @test_simple_while(%func_arg : tensor<1x2x3xi32>) {
+  %relaxed_func_arg = test.ti.relax_type %func_arg : tensor<1x2x3xi32> to tensor<*xi32>
+  %while = scf.while (%before_arg = %relaxed_func_arg) : (tensor<*xi32>) -> tensor<*xi32> {
+    %cond = "getValue"() : () -> (i1)
+    scf.condition(%cond) %before_arg : tensor<*xi32>
+  } do {
+  ^bb0(%after_arg: tensor<*xi32>):
+    %join = "test.ti.join_types"(%after_arg) : (tensor<*xi32>) -> (tensor<*xi32>)
+    scf.yield %join : tensor<*xi32>
+  }
+  "test.ti.allow_input_type_specialization"(%while) : (tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK: func @test_simple_while([[FUNC_ARG:%.+]]: tensor<1x2x3xi32>) {
+// CHECK:   [[RELAXED_FUNC_ARG:%.+]] = test.ti.relax_type [[FUNC_ARG]] : tensor<1x2x3xi32> to tensor<*xi32>
+// CHECK:   [[WHILE:%.+]] = scf.while ([[BEFORE_ARG:%.+]] = [[RELAXED_FUNC_ARG]]) : (tensor<*xi32>) -> tensor<*xi32> {
+// CHECK:     [[COND:%.+]] = "getValue"() : () -> i1
+// CHECK:     scf.condition([[COND]]) [[BEFORE_ARG]] : tensor<*xi32>
+// CHECK:   } do {
+// CHECK:   ^bb0([[AFTER_ARG:%.+]]: tensor<*xi32>):  // no predecessors
+// CHECK:     [[JOIN:%.+]] = "test.ti.join_types"([[AFTER_ARG]]) : (tensor<*xi32>) -> tensor<1x2x3xi32>
+// CHECK:     [[RELAXED_JOIN:%.+]] = test.ti.relax_type [[JOIN]] : tensor<1x2x3xi32> to tensor<*xi32>
+// CHECK:     scf.yield [[RELAXED_JOIN]] : tensor<*xi32>
+// CHECK:   }
+// CHECK:   [[SPECIALIZED_WHILE:%.+]] = test.ti.specialize_type [[WHILE]] : tensor<*xi32> to tensor<1x2x3xi32>
+// CHECK:   "test.ti.allow_input_type_specialization"([[SPECIALIZED_WHILE]]) : (tensor<1x2x3xi32>) -> ()
+// CHECK:   return
+// CHECK: }
+
+// -----
+
+func @test_while(%func_arg : tensor<1x2x3xi32>) {
+  %relaxed_func_arg = test.ti.relax_type %func_arg : tensor<1x2x3xi32> to tensor<*xi32>
+  %while = scf.while (%before_arg = %relaxed_func_arg) : (tensor<*xi32>) -> tensor<*xi32> {
+    %cond = "getValue"() : () -> (i1)
+    scf.condition(%cond) %before_arg : tensor<*xi32>
+  } do {
+  ^bb0(%after_arg: tensor<*xi32>):
+    %dyn = "test.ti.dynamicize"(%after_arg) : (tensor<*xi32>) -> (tensor<*xi32>)
+    scf.yield %dyn : tensor<*xi32>
+  }
+  "test.ti.allow_input_type_specialization"(%while) : (tensor<*xi32>) -> ()
+  return
+}
+
+// CHECK: func @test_while([[FUNC_ARG:%.+]]: tensor<1x2x3xi32>) {
+// CHECK:   [[RELAXED_FUNC_ARG:%.+]] = test.ti.relax_type [[FUNC_ARG]] : tensor<1x2x3xi32> to tensor<*xi32>
+// CHECK:   [[WHILE:%.+]] = scf.while ([[BEFORE_ARG:%.+]] = [[RELAXED_FUNC_ARG]]) : (tensor<*xi32>) -> tensor<*xi32> {
+// CHECK:     [[COND:%.+]] = "getValue"() : () -> i1
+// CHECK:     scf.condition([[COND]]) [[BEFORE_ARG]] : tensor<*xi32>
+// CHECK:   } do {
+// CHECK:   ^bb0([[AFTER_ARG:%.+]]: tensor<*xi32>):  // no predecessors
+// CHECK:     [[SPECIALIZED_AFTER_ARG:%.+]] = test.ti.specialize_type [[AFTER_ARG]] : tensor<*xi32> to tensor<?x?x?xi32>
+// CHECK:     [[DYN:%.+]] = "test.ti.dynamicize"([[SPECIALIZED_AFTER_ARG]]) : (tensor<?x?x?xi32>) -> tensor<?x?x?xi32>
+// CHECK:     [[RELAXED_DYN:%.+]] = test.ti.relax_type [[DYN]] : tensor<?x?x?xi32> to tensor<*xi32>
+// CHECK:     scf.yield [[RELAXED_DYN]] : tensor<*xi32>
+// CHECK:   }
+// CHECK:   [[SPECIALIZED_WHILE:%.+]] = test.ti.specialize_type [[WHILE]] : tensor<*xi32> to tensor<?x?x?xi32>
+// CHECK:   "test.ti.allow_input_type_specialization"([[SPECIALIZED_WHILE]]) : (tensor<?x?x?xi32>) -> ()
+// CHECK:   return
+// CHECK: }

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/JoinMeetTypeInterface.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringSwitch.h"
 
 // Include this before the using namespace lines below to
 // test that we don't have namespace dependencies.
 #include "TestOpsDialect.cpp.inc"
static void print(SingleNoTerminatorCustomAsmOp op, OpAsmPrinter &printer) {
   printer.printRegion(
       op.getRegion(), /*printEntryBlockArgs=*/false,
       // This op has a single block without terminators. But explicitly mark
       // as not printing block terminators for testing.
       /*printBlockTerminators=*/false);
 }
 
 //===----------------------------------------------------------------------===//
-// Control-Flow Edge Operand Types Test Ops
+// Test Type Inferences
 //===----------------------------------------------------------------------===//
 
+OpFoldResult TIRelaxTypeOp::fold(ArrayRef<Attribute> operands) {
+  return impl::foldCastOp(*this);
+}
+
+static LogicalResult verify(TIRelaxTypeOp op) {
+  if (!isLessSpecializedOrSame(op.getType(), op.getIn().getType()))
+    return op.emitOpError(
+        "output type is not less specialized than the input type");
+  return success();
+}
+
+LogicalResult
+TIRelaxTypeOp::inferReturnTypes(MLIRContext *, Optional<Location> location,
+                                ValueRange operands, DictionaryAttr attributes,
+                                RegionRange regions,
+                                SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(operands[0].getType());
+  return success();
+}
+
+bool TIRelaxTypeOp::isCompatibleReturnTypes(TypeRange inferredReturnTypes,
+                                            TypeRange resultTypes) {
+  return llvm::all_of_zip(inferredReturnTypes, resultTypes,
+                          isMoreSpecializedOrSame);
+}
+
+OpFoldResult TISpecializeTypeOp::fold(ArrayRef<Attribute> operands) {
+  return impl::foldCastOp(*this);
+}
+
+static LogicalResult verify(TISpecializeTypeOp op) {
+  if (!isMoreSpecializedOrSame(op.getType(), op.getIn().getType()))
+    return op.emitOpError(
+        "output type is not more specialized than the input type");
+  return success();
+}
+
+LogicalResult TIJoinTypesOp::inferReturnTypes(
+    ::mlir::MLIRContext *context, ::llvm::Optional<::mlir::Location> location,
+    ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
+    ::mlir::RegionRange regions,
+    ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+  Type resultTy;
+  for (auto operand : operands)
+    resultTy =
+        !resultTy ? operand.getType() : joinTypes(resultTy, operand.getType());
+  inferredReturnTypes.push_back(resultTy);
+  return success();
+}
+
+bool TIJoinTypesOp::isCompatibleReturnTypes(::mlir::TypeRange lhs,
+                                            ::mlir::TypeRange rhs) {
+  return llvm::all_of_zip(lhs, rhs, isMoreSpecializedOrSame) ||
+         llvm::all_of_zip(lhs, rhs, isLessSpecializedOrSame);
+}
+
 Optional<MutableOperandRange>
-CFEBranchOp::getMutableSuccessorOperands(unsigned index) {
+TIBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return getTargetOperandsMutable();
 }
 
-bool CFEBranchOp::areCompatibleControlFlowEdgeOperandTypes(Type operandTy,
-                                                           Type blockArgTy) {
+bool TIBranchOp::areCompatibleControlFlowEdgeOperandTypes(Type operandTy,
   return isMoreSpecializedOrSame(operandTy, blockArgTy);
 }
 
-static void print(OpAsmPrinter &p, CFERegionIfOp op) { printRegionIfOp(p, op); }
+static void print(OpAsmPrinter &p, TIRegionIfOp op) { printRegionIfOp(p, op); }
 
-OperandRange CFERegionIfOp::getSuccessorEntryOperands(unsigned index) {
+OperandRange TIRegionIfOp::getSuccessorEntryOperands(unsigned index) {
   assert(index < 2 && "invalid region index");
   return getOperands();
 }
 
-void CFERegionIfOp::getSuccessorRegions(
+void TIRegionIfOp::getSuccessorRegions(
     Optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
   // We always branch to the join region.
   if (index.hasValue()) {
     if (index.getValue() < 2)
       regions.push_back(RegionSuccessor(&getJoinRegion(), getJoinArgs()));
     else
       regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   // The then and else regions are the entry regions of this op.
   regions.push_back(RegionSuccessor(&getThenRegion(), getThenArgs()));
   regions.push_back(RegionSuccessor(&getElseRegion(), getElseArgs()));
 }
 
-bool CFERegionIfOp::areCompatibleControlFlowEdgeOperandTypes(Type operandTy,
-                                                             Type blockArgTy) {
+bool TIRegionIfOp::areCompatibleControlFlowEdgeOperandTypes(Type operandTy,
   return isMoreSpecializedOrSame(operandTy, blockArgTy);
 }
 
+LogicalResult
+TIDynamicizeOp::inferReturnTypes(MLIRContext *, Optional<Location> location,
+                                 ValueRange operands, DictionaryAttr attributes,
+                                 RegionRange regions,
+                                 SmallVectorImpl<Type> &inferredReturnTypes) {
+  RankedTensorType inputTy = operands[0].getType().dyn_cast<RankedTensorType>();
+  if (!inputTy) {
+    inferredReturnTypes.push_back(operands[0].getType());
+    return success();
+  }
+
+  SmallVector<int64_t> shape(inputTy.getShape().begin(),
+                             inputTy.getShape().end());
+  for (auto &dim : shape) {
+    if (!ShapedType::isDynamic(dim)) {
+      dim = ShapedType::kDynamicSize;
+      break;
+    }
+  }
+
+  inferredReturnTypes.push_back(inputTy.clone(shape));
+
+  return success();
+}
+
+bool TIDynamicizeOp::isCompatibleReturnTypes(::mlir::TypeRange lhs,
+                                             ::mlir::TypeRange rhs) {
+  return llvm::all_of_zip(lhs, rhs, isMoreSpecializedOrSame) ||
+         llvm::all_of_zip(lhs, rhs, isLessSpecializedOrSame);
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestOpStructs.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/Dialect/Test/TestOps.td

let arguments = (ins
     Variadic<AnyTensor>:$input1,
     AnyTensor:$input2,
     Variadic<AnyTensor>:$input3
   );
 }
 def VariadicWithSameOperandsResult :
       TEST_Op<"variadic_with_same_operand_results",
               [SameOperandsAndResultType]> {
-  let arguments = (ins Variadic<AnySignlessInteger>);
-  let results = (outs AnySignlessInteger:$result);
+  let arguments = (ins Variadic<AnyTensor>:$operands);
+  let results = (outs AnyTensor:$result);
 }
 
 def SameOperandsResultType : TEST_Op<
     "same_operand_result_type", [SameOperandsAndResultType]> {
   let arguments = (ins AnyTensor:$operand);
   let results = (outs AnyTensor:$result);
 }
 
def TestDefaultStrAttrHasValueOp : TEST_Op<"has_str_value"> {
   let arguments = (ins DefaultValuedStrAttr<StrAttr, "">:$value);
   let assemblyFormat = "attr-dict";
 }
 
 def : Pat<(TestDefaultStrAttrNoValueOp $value),
           (TestDefaultStrAttrHasValueOp ConstantStrAttr<StrAttr, "foo">)>;
 
 //===----------------------------------------------------------------------===//
-// Test Control-Flow Edge Operand Types
+// Test Type Inferences
 //===----------------------------------------------------------------------===//
 
-def CFEBranchOp : TEST_Op<"cfe.br",
-    [DeclareOpInterfaceMethods<BranchOpInterface,
+// FIXME: Have `SameOperandsAndResultType` and other similar traits provide a
+// default implementation of `InferTypeOpInterface`, and use it here.
+def TIRelaxTypeOp
+    : TEST_Op<"ti.relax_type",
+              [NoSideEffect,
+               DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>,
+               DeclareOpInterfaceMethods<AllowsOutputTypesSpecializationInterface>,
+               DeclareOpInterfaceMethods<InferTypeOpInterface, ["inferReturnTypes"]>]> {
+  let summary = "relax the value type";
+  let description = [{
+    Relax the type of the input value. The "content" of the value itself is
+    unchanged. Only the type is modified.
+
+    It is typically introduced by the type inference pass to satisfy type
+    constraints, when a less specialized type is required.
+
+    Example:
+    ```
+    // The hypothetical `process` op implements `SameOperandsAndResultType` and
+    // allows output type specialization.
+    %val = "process"(%x, %y) : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<1x2xi32>
+    %relaxed_value = test.ti.relax_type %value : tensor<1x2xi32> to tensor<*xi32>
+    "op_disallowing_input_type_specialization"(%relaxed_value) : (tensor<*xi32>) -> ()
+    ```
+  }];
+
+  let arguments = (ins AnyType:$in);
+  let results = (outs AnyType:$out);
+
+  let builders = [
+    OpBuilder<(ins "::mlir::Value":$source, "::mlir::Type":$destType), [{
+       impl::buildCastOp($_builder, $_state, source, destType);
+    }]>
+  ];
+
+  let extraClassDeclaration = [{
+    static bool isCompatibleReturnTypes(::mlir::TypeRange lhs, ::mlir::TypeRange rhs);
+    bool allowsOutputTypeSpecialization(unsigned operandIndex) { return false; }
+  }];
+
+  let assemblyFormat = "$in attr-dict `:` type($in) `to` type($out)";
+
+  let hasFolder = 1;
+}
+
+def TISpecializeTypeOp : TEST_Op<"ti.specialize_type", [NoSideEffect]> {
+  let summary = "specialize the value type";
+  let description = [{
+    Specialize the type of the input value. The "content" of the value itself is
+    unchanged. Only the type is modified.
+
+    It is typically introduced by the type inference pass to satisfy type
+    constraints, when a more specialized type is required.
+
+    Example:
+    ```
+    // The hypothetical `process` op implements `SameOperandsAndResultType` but
+    // disallows output type specialization.
+    %val = "process"(%x, %y) : (tensor<1x?xi32>, tensor<?x2xi32>) -> tensor<*xi32>
+    %specialized_val = specialize_type %val : tensor<*xi32> to tensor<1x2xi32>
+    "op_allowing_input_type_specialization"(%specialized_val) : (tensor<1x2xi32>) -> ()
+    ```
+  }];
+
+  let arguments = (ins AnyType:$in);
+  let results = (outs AnyType:$out);
+
+  let builders = [
+    OpBuilder<(ins "::mlir::Value":$source, "::mlir::Type":$destType), [{
+       impl::buildCastOp($_builder, $_state, source, destType);
+    }]>
+  ];
+
+  let assemblyFormat = "$in attr-dict `:` type($in) `to` type($out)";
+
+  let hasFolder = 1;
+}
+
+def TIDisallowInputTypeSpecializationOp
+    : TEST_Op<"ti.disallow_input_type_specialization"> {
+  let description = [{Disallow type specialization for all operands.}];
+  let arguments = (ins Variadic<AnyType>:$inputs);
+}
+
+def TIAllowInputTypeSpecializationOp
+    : TEST_Op<"ti.allow_input_type_specialization",
+              [DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>]> {
+  let description = [{Allow type specialization for all operands.}];
+  let arguments = (ins Variadic<AnyType>:$inputs);
+}
+
+def TIConditionallyAllowInputTypeSpecializationOp
+    : TEST_Op<"ti.conditionally_allow_input_type_specialization",
+              [DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>]> {
+  let description = [{Allow type specialization for evenly-indexed operands.}];
+  let arguments = (ins Variadic<AnyType>:$inputs);
+  let extraClassDeclaration = [{
+    bool allowsInputTypeSpecialization(unsigned operandIndex) {
+      return operandIndex % 2 == 0;
+    }
+  }];
+}
+
+def TIJoinTypesOp
+    : TEST_Op<"ti.join_types",
+              [DeclareOpInterfaceMethods<InferTypeOpInterface>,
+               DeclareOpInterfaceMethods<AllowsOutputTypesSpecializationInterface>]> {
+  let description = [{
+    Type inference joins the operand types.
+    Output type specialization is configurable via a boolean attribute.
+  }];
+  let arguments = (ins Variadic<AnyType>:$inputs,
+                       DefaultValuedAttr<BoolAttr, "true">:$allowOutputTypeSpecialization);
+  let results = (outs AnyType:$result);
+  let extraClassDeclaration = [{
+    static bool isCompatibleReturnTypes(::mlir::TypeRange lhs, ::mlir::TypeRange rhs);
+    bool allowsOutputTypeSpecialization(unsigned outputIndex) {
+      return this->getAllowOutputTypeSpecialization();
+    }
+  }];
+}
+
+def TIBranchOp : TEST_Op<"ti.br",
+    [DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>,
+     DeclareOpInterfaceMethods<BranchOpInterface,
                                ["areCompatibleControlFlowEdgeOperandTypes"]>,
      Terminator]> {
+  let description = [{Branch allowing type specialization for all operands.}];
   let arguments = (ins Variadic<AnyType>:$targetOperands);
   let successors = (successor AnySuccessor:$target);
 }
 
-def CFERegionIfYieldOp : TEST_Op<"cfe.region_if_yield",
-      [NoSideEffect, ReturnLike, Terminator]> {
+def TIRegionIfYieldOp : TEST_Op<"ti.region_if_yield",
+      [DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>,
+       NoSideEffect, ReturnLike, Terminator]> {
   let arguments = (ins Variadic<AnyType>:$results);
   let assemblyFormat = [{
     $results `:` type($results) attr-dict
   }];
 }
 
-def CFERegionIfOp
-    : TEST_Op<"cfe.region_if",
-      [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+def TIRegionIfOp
+    : TEST_Op<"ti.region_if",
+      [DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>,
+       DeclareOpInterfaceMethods<AllowsOutputTypesSpecializationInterface>,
+       DeclareOpInterfaceMethods<RegionBranchOpInterface,
                                  ["areCompatibleControlFlowEdgeOperandTypes"]>,
-       SingleBlockImplicitTerminator<"CFERegionIfYieldOp">,
+       SingleBlockImplicitTerminator<"TIRegionIfYieldOp">,
        RecursiveSideEffects]> {
   let description =[{
     Represents an abstract if-then-else-join pattern. In this context, the then
     and else regions jump to the join region, which finally returns to its
     parent op.
     }];
 
   let printer = [{ return ::print(p, *this); }];
::mlir::Block::BlockArgListType getJoinArgs() {
       return getBody(2)->getArguments();
     }
     ::mlir::OperandRange getSuccessorEntryOperands(unsigned index);
   }];
 
   let verifier = [{ return RegionBranchOpInterface::verifyTypes(*this); }];
 }
 
+def TIDynamicizeOp
+    : TEST_Op<"ti.dynamicize",
+      [DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>,
+       DeclareOpInterfaceMethods<AllowsOutputTypesSpecializationInterface>,
+       DeclareOpInterfaceMethods<InferTypeOpInterface, ["inferReturnTypes"]>]> {
+  let description = [{
+    An op with particularly weird type inference, that progressively sets each
+    dimension to dynamic (`?`).
+  }];
+  let arguments = (ins AnyTensor);
+  let results = (outs AnyTensor);
+  let extraClassDeclaration = [{
+    static bool isCompatibleReturnTypes(::mlir::TypeRange lhs, ::mlir::TypeRange rhs);
+  }];
+}
+
 #endif // TEST_OPS

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestConstantFold.cpp
   TestInlining.cpp
   TestJoinMeetTypeInterface.cpp
   TestLoopFusion.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
+  TestTypeInference.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   LINK_LIBS PUBLIC
   MLIRTestDialect
   MLIRTransforms
   )
 
 target_include_directories(MLIRTestTransforms
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test
   )

mlir/test/lib/Transforms/TestTypeInference.cpp

+//===- TestTypeInference.cpp - Test Type Inference  -----------------------===//
+//
+// 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 "TestDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/TypeInference.h"
+
+using namespace mlir;
+
+namespace {
+Value typeAdapter(Value value, Type type) {
+  Type currentType = value.getType();
+  if (currentType == type)
+    return value;
+  OpBuilder builder(value.getContext());
+  builder.setInsertionPointAfterValue(value);
+  if (isLessSpecialized(currentType, type))
+    return builder.createOrFold<test::TISpecializeTypeOp>(value.getLoc(), value,
+                                                          type);
+  if (isMoreSpecialized(currentType, type))
+    return builder.createOrFold<test::TIRelaxTypeOp>(value.getLoc(), value,
+                                                     type);
+  return {};
+}
+
+struct TestTypeInference : public PassWrapper<TestTypeInference, Pass> {
+  StringRef getArgument() const final { return "test-type-inference"; }
+  StringRef getDescription() const final { return "Test type inference with"; }
+  void runOnOperation() override {
+    TypeInference typeInference(getOperation(), typeAdapter);
+    if (failed(typeInference.run()))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestTypeInference() { PassRegistration<TestTypeInference>(); }
+} // namespace test
+} // namespace mlir

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

 void registerTestMemRefStrideCalculation();
 void registerTestNumberOfBlockExecutionsPass();
 void registerTestNumberOfOperationExecutionsPass();
 void registerTestOpaqueLoc();
 void registerTestPDLByteCodePass();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
+void registerTestTypeInference();
 void registerTestVectorConversions();
 } // namespace test
 } // namespace mlir
 
 namespace test {
 void registerTestDialect(DialectRegistry &);
 } // namespace test
 
#endif
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestNumberOfBlockExecutionsPass();
   mlir::test::registerTestNumberOfOperationExecutionsPass();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
+  mlir::test::registerTestTypeInference();
   mlir::test::registerTestVectorConversions();
 }
 #endif
 
 int main(int argc, char **argv) {
   registerAllPasses();
 #ifdef MLIR_INCLUDE_TESTS
   registerTestPasses();