[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

2282

more specialized?

[stephenneuendorffer]

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

2241

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

2278–2279

specialize_types?

2286–2287

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

Output?

707

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

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

2820

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

2241

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

2286–2287

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

"support for output types refinement" is implicit.

I rewrote to clarify.

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

2820

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

2021

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

2030

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/Dialect/StandardOps/IR/Ops.h

 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/StandardOps/IR/OpsEnums.h.inc"
 
 namespace mlir {
 class AffineMap;

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

 #define STANDARD_OPS
 
 include "mlir/Dialect/StandardOps/IR/StandardOpsBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/VectorInterfaces.td"
 
 def StandardOps_Dialect : Dialect {
   let name = "std";
   let cppNamespace = "::mlir";
   let hasConstantMaterializer = 1;
 }
let description = [{
     ```
   }];
 
   let arguments = (ins AnyTypeOf<[AnyRankedOrUnrankedMemRef, AnyTensor],
                                  "any memref or tensor type">:$memrefOrTensor);
   let results = (outs Index);
   let verifier = ?;
 
-  let builders = [
-    OpBuilder<(ins "Value":$tensor), [{
-      auto indexType = $_builder.getIndexType();
-      build($_builder, $_state, indexType, tensor);
-    }]>];
-
   let hasFolder = 1;
   let assemblyFormat = "$memrefOrTensor attr-dict `:` type($memrefOrTensor)";
 }
 
 //===----------------------------------------------------------------------===//
 // RemFOp
 //===----------------------------------------------------------------------===//
 
let description = [{
     ```
   }];
 
   let arguments = (ins BoolLike:$condition,
                        AnyType:$true_value,
                        AnyType:$false_value);
   let results = (outs AnyType:$result);
 
-  let builders = [
-    OpBuilder<(ins "Value":$condition, "Value":$trueValue,
-      "Value":$falseValue), [{
-      $_state.addOperands({condition, trueValue, falseValue});
-      $_state.addTypes(trueValue.getType());
-    }]>];
-
   let extraClassDeclaration = [{
       Value getCondition() { return condition(); }
       Value getTrueValue() { return true_value(); }
       Value getFalseValue() { return false_value(); }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
def SwitchOp : Std_Op<"switch",
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TruncateIOp
 //===----------------------------------------------------------------------===//
 
 def TruncateIOp : Std_Op<"trunci", [NoSideEffect,
     DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
     ElementwiseMappable.traits> {
   let summary = "integer truncation operation";
   let description = [{
     The integer truncation operation takes an integer input of
     width M and an integer destination type of width N. The destination
     bit-width must be smaller than the input bit-width (N < M).
     The top-most (N - M) bits of the input are discarded.
 
     Example:
 
     ```mlir
       %1 = constant 21 : i5           // %1 is 0b10101
       %2 = trunci %1 : i5 to i4       // %2 is 0b0101
       %3 = trunci %1 : i5 to i3       // %3 is 0b101
 
       %5 = trunci %0 : vector<2 x i32> to vector<2 x i16>
     ```
def ZeroExtendIOp : Std_Op<"zexti", [NoSideEffect,
   let parser = [{
     return impl::parseCastOp(parser, result);
   }];
   let printer = [{
     return printStandardCastOp(this->getOperation(), p);
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// RelaxTypeOp
+//===----------------------------------------------------------------------===//
+
+def RelaxTypeOp
+    : Std_Op<"relax_type",
+             [NoSideEffect,
+              DeclareOpInterfaceMethods<AllowsInputTypesSpecializationInterface>,
+              DeclareOpInterfaceMethods<InferTypeOpInterface, ["inferReturnTypes"]>]> {
+  let summary = "Relax the value type.";
+  let description = [{
+    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 = 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 "Value":$source, "Type":$destType), [{
+       impl::buildCastOp($_builder, $_state, source, destType);
+    }]>
+  ];
+
+  let extraClassDeclaration = [{
+    static bool isCompatibleReturnTypes(TypeRange lhs, TypeRange rhs);
+  }];
+
+  let parser = [{
+    return impl::parseCastOp(parser, result);
+  }];
+  let printer = [{
+    return printStandardCastOp(this->getOperation(), p);
+  }];
+
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// TypeSpecializationOp
+//===----------------------------------------------------------------------===//
+
+def SpecializeTypeOp : Std_Op<"specialize_type", [NoSideEffect]> {
+  let summary = "Specialize the value type.";
+  let description = [{
+    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 "Value":$source, "Type":$destType), [{
+       impl::buildCastOp($_builder, $_state, source, destType);
+    }]>
+  ];
+
+  let parser = [{
+    return impl::parseCastOp(parser, result);
+  }];
+  let printer = [{
+    return printStandardCastOp(this->getOperation(), p);
+  }];
+
+  let hasFolder = 1;
+}
+
+
 #endif // STANDARD_OPS

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 operand 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/Passes.h

 /// Creates a pass which delete symbol operations that are unreachable. This
 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();
 
 /// Creates an interprocedural pass to normalize memrefs to have a trivial
 /// (identity) layout map.
 std::unique_ptr<OperationPass<ModuleOp>> createNormalizeMemRefsPass();
 
+/// Creates a pass to infer operation types.
+std::unique_ptr<Pass> createTypeInferencePass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Transforms/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_PASSES_H

mlir/include/mlir/Transforms/Passes.td

let description = [{
     ```
 
     See [Symbols and SymbolTables](SymbolsAndSymbolTables.md) for more
     information on `Symbols`.
   }];
   let constructor = "mlir::createSymbolDCEPass()";
 }
 
+def TypeInferencePass : Pass<"infer-types"> {
+  let summary = "Infer the types of values";
+  let description = [{
+    This pass propagates type information through the IR to refine value types.
+
+    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.
+
+    For example, running type inference 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 = type_relaxation %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 = type_relaxation %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 = type_relaxation %2 : tensor<2xi32> to tensor<*xi32>
+      br ^bb3(%3 : tensor<*xi32>)
+    ^bb3(%4: tensor<*xi32>):  // 2 preds: ^bb1, ^bb2
+      %5 = type_specialization %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
+    ```
+    .
+  }];
+  let constructor = "mlir::createTypeInferencePass()";
+}
+
 def ViewOpGraphPass : Pass<"view-op-graph"> {
   let summary = "Print Graphviz visualization of an operation";
   let description = [{
     This pass prints a Graphviz graph of a module.
 
     - Operations are represented as nodes;
     - Uses (data flow) as edges;
     - Control flow as dashed edges;

mlir/lib/Dialect/StandardOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRStandard
   IR/Ops.cpp
   Utils/Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps
 
   DEPENDS
   MLIRStandardOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRCallInterfaces
   MLIRCastInterfaces
   MLIRControlFlowInterfaces
   MLIRIR
+  MLIRJoinMeetTypeInterface
   MLIRSideEffectInterfaces
   MLIRVectorInterfaces
   )
 
 add_subdirectory(Transforms)

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

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
+#include "mlir/Interfaces/JoinMeetTypeInterface.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 #include <numeric>
if (srcType.cast<IntegerType>().getWidth() >=
       dstType.cast<IntegerType>().getWidth())
     return op.emitError("result type ")
            << dstType << " must be wider than operand type " << srcType;
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// RelaxTypeOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult RelaxTypeOp::fold(ArrayRef<Attribute> operands) {
+  return impl::foldCastOp(*this);
+}
+
+static LogicalResult verify(RelaxTypeOp op) {
+  if (!isLessSpecializedOrSame(op.getType(), op.in().getType()))
+    return op.emitOpError(
+        "output type is not less specialized than the input type");
+  return success();
+}
+
+LogicalResult
+RelaxTypeOp::inferReturnTypes(MLIRContext *, Optional<Location> location,
+                              ValueRange operands, DictionaryAttr attributes,
+                              RegionRange regions,
+                              SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(operands[0].getType());
+  return success();
+}
+
+bool RelaxTypeOp::isCompatibleReturnTypes(TypeRange inferredReturnTypes,
+                                          TypeRange resultTypes) {
+  return llvm::all_of_zip(inferredReturnTypes, resultTypes,
+                          isMoreSpecializedOrSame);
+}
+
+//===----------------------------------------------------------------------===//
+// SpecializeTypeOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult SpecializeTypeOp::fold(ArrayRef<Attribute> operands) {
+  return impl::foldCastOp(*this);
+}
+
+static LogicalResult verify(SpecializeTypeOp op) {
+  if (!isMoreSpecializedOrSame(op.getType(), op.in().getType()))
+    return op.emitOpError(
+        "output type is not more specialized than the input type");
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/StandardOps/IR/Ops.cpp.inc"

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
+//
+//===----------------------------------------------------------------------===//
+//
+// The type inference analysis uses the dataflow analysis infrastructure to
+// infer MLIR value types over the IR.
+// The type inference pass uses the results of the analysis to update MLIR value
+// types, querying operation interfaces to appropriately update types in place
+// or insert type relaxation or specialization ops.
+//
+// The `join` relationship over the type lattice is 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>
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/JoinMeetTypeInterface.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Format.h"
+
+#define DEBUG_TYPE "infer-types"
+
+using namespace mlir;
+using llvm::dbgs;
+
+namespace {
+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;
+};
+
+struct 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:
+  ChangeResult visitInferTypeOpInterface(
+      Operation *op, ArrayRef<LatticeElement<TypeLatticeValue> *> operands);
+  ChangeResult visitSameOperandsAndResultShape(
+      Operation *op, ArrayRef<LatticeElement<TypeLatticeValue> *> operands);
+  ChangeResult visitSameOperandsAndResultType(
+      Operation *op, ArrayRef<LatticeElement<TypeLatticeValue> *> operands);
+
+  // A wrapper allowing tracking updates for debugging purposes.
+  ChangeResult 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;
+  }
+
+  LogicalResult status;
+};
+
+ChangeResult 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;
+
+  // Handle different ways to infer the result types.
+  // Process in order of "strongest" (e.g. full type vs shape) and "simplest"
+  // (e.g. `SameOperandsAndResultType` vs `InferTypeOpInterface`).
+  if (op->hasTrait<OpTrait::SameOperandsAndResultType>())
+    return visitSameOperandsAndResultType(op, operands);
+  if (isa<InferTypeOpInterface>(op))
+    return visitInferTypeOpInterface(op, operands);
+  if (op->hasTrait<OpTrait::SameOperandsAndResultShape>())
+    return visitSameOperandsAndResultShape(op, operands);
+  // TODO: Extend support to other interfaces.
+
+  // By default, simply use the current result types.
+  ChangeResult change = ChangeResult::NoChange;
+  for (Value result : op->getResults())
+    change |= updateLatticeElement(result, result.getType());
+  return change;
+}
+
+ChangeResult TypeInferenceAnalysis::visitInferTypeOpInterface(
+    Operation *op, ArrayRef<LatticeElement<TypeLatticeValue> *> operands) {
+  // Ephemerally override operand types with inferred types.
+  llvm::SmallVector<Type, 8> 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, 4> inferredReturnTypes;
+  status = cast<InferTypeOpInterface>(op).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;
+}
+
+ChangeResult TypeInferenceAnalysis::visitSameOperandsAndResultType(
+    Operation *op, ArrayRef<LatticeElement<TypeLatticeValue> *> operands) {
+  if (operands.empty())
+    return ChangeResult::NoChange;
+
+  Type ty = operands[0]->getValue().type;
+  for (auto *operand : operands)
+    ty = meetTypes(ty, operand->getValue().type);
+
+  ChangeResult change = ChangeResult::NoChange;
+  for (auto result : op->getResults())
+    change |= updateLatticeElement(result, ty);
+
+  return change;
+}
+
+ChangeResult TypeInferenceAnalysis::visitSameOperandsAndResultShape(
+    Operation *op, ArrayRef<LatticeElement<TypeLatticeValue> *> operands) {
+  if (operands.empty())
+    return ChangeResult::NoChange;
+
+  // TODO: Introduce helper so that we can handle scalars.
+  SmallVector<int64_t> shape;
+  ArrayRef<int64_t> initialShape =
+      operands[0]->getValue().type.cast<ShapedType>().getShape();
+  shape.insert(shape.end(), initialShape.begin(), initialShape.end());
+
+  for (auto *operand : operands)
+    shape = *meetShapes(shape,
+                        operand->getValue().type.cast<ShapedType>().getShape());
+
+  ChangeResult change = ChangeResult::NoChange;
+  for (auto result : op->getResults())
+    change |= updateLatticeElement(
+        result, result.getType().cast<ShapedType>().clone(shape));
+
+  return change;
+}
+
+} // end anonymous namespace
+
+namespace {
+class TypeInference {
+public:
+  explicit TypeInference(Operation *op) : op(op), analysis(op->getContext()) {}
+
+  /// 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;
+};
+
+LogicalResult TypeInference::updateOp(Operation *op) {
+  for (auto &region : op->getRegions()) {
+    for (auto &block : region) {
+      if (failed(updateBlock(block)))
+        return failure();
+    }
+  }
+  return updateOpResults(op);
+}
+
+LogicalResult 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 TypeInference::updateOpResults(Operation *op) {
+  for (auto result : op->getResults())
+    if (failed(updateType(result)))
+      return failure();
+  return success();
+}
+
+LogicalResult 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();
+
+  // Never specialize `type_relaxation`.
+  if (dyn_cast_or_null<RelaxTypeOp>(definingOp))
+    return success();
+
+  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.
+  auto getValueWithOriginalType = [&]() -> Value {
+    if (!valueWithOriginalType) {
+      OpBuilder builder(value.getContext());
+      builder.setInsertionPointAfterValue(value);
+      valueWithOriginalType =
+          builder.createOrFold<RelaxTypeOp>(value.getLoc(), value, originalTy);
+    }
+    return valueWithOriginalType;
+  };
+  auto getValueWithSpecializedType = [&]() -> Value {
+    if (!valueWithInferredType) {
+      OpBuilder builder(value.getContext());
+      builder.setInsertionPointAfterValue(value);
+      valueWithInferredType = builder.createOrFold<SpecializeTypeOp>(
+          value.getLoc(), value, inferredTy);
+    }
+    return valueWithInferredType;
+  };
+  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());
+    use.set(allowsInputTypeSpecialization ? getValueWithSpecializedType()
+                                          : getValueWithOriginalType());
+    nSpecializedUses += allowsInputTypeSpecialization;
+  }
+
+  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();
+}
+
+struct TypeInferencePass : public TypeInferencePassBase<TypeInferencePass> {
+  void runOnOperation() override {
+    TypeInference typeInference(getOperation());
+    if (failed(typeInference.run()))
+      signalPassFailure();
+  }
+};
+} // end anonymous namespace
+
+std::unique_ptr<Pass> mlir::createTypeInferencePass() {
+  return std::make_unique<TypeInferencePass>();
+}

mlir/test/IR/invalid-ops.mlir

^bb0(%cond : i32, %t : i32, %f : i32):
   // expected-error@+1 {{op operand #0 must be bool-like}}
   %r = "std.select"(%cond, %t, %f) : (i32, i32, i32) -> i32
 }
 
 // -----
 
 func @func_with_ops(i1, i32, i64) {
 ^bb0(%cond : i1, %t : i32, %f : i64):
-  // expected-error@+1 {{all of {true_value, false_value, result} have same type}}
+  // expected-error@+1 {{'std.select' op inferred type(s) 'i64' are incompatible with return type(s) of operation 'i32'}}
   %r = "std.select"(%cond, %t, %f) : (i1, i32, i64) -> i32
 }
 
 // -----
 
 func @func_with_ops(vector<12xi1>, vector<42xi32>, vector<42xi32>) {
 ^bb0(%cond : vector<12xi1>, %t : vector<42xi32>, %f : vector<42xi32>):
   // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type}}

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 -infer-types %s -split-input-file -o %t1
+// RUN: FileCheck %s < %t1
+// Verify idempotence.
+// RUN: mlir-opt -allow-unregistered-dialect -infer-types %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:%.+]] = 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:%.+]] = 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 = 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:%.+]] = 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:%.+]] = 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:%.+]] = 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:%.+]] = 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:%.+]] = 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:%.+]] = 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:%.+]] = relax_type [[JOIN]] : tensor<2xi32> to tensor<*xi32>
+// CHECK:   br ^bb3([[RELAXED_JOIN]] : tensor<*xi32>)
+// CHECK: ^bb3([[VAL:%.+]]: tensor<*xi32>):
+// CHECK:   [[SPECIALIZED_VAL:%.+]] = 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:%.+]] = 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:%.+]] = 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:%.+]] = 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:%.+]] = 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:%.+]] = 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 = 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:%.+]] = 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:%.+]] = relax_type [[JOIN]] : tensor<1x2x3xi32> to tensor<*xi32>
+// CHECK:     scf.yield [[RELAXED_JOIN]] : tensor<*xi32>
+// CHECK:   }
+// CHECK:   [[SPECIALIZED_WHILE:%.+]] = 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 = 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:%.+]] = 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:%.+]] = 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:%.+]] = relax_type [[DYN]] : tensor<?x?x?xi32> to tensor<*xi32>
+// CHECK:     scf.yield [[RELAXED_DYN]] : tensor<*xi32>
+// CHECK:   }
+// CHECK:   [[SPECIALIZED_WHILE:%.+]] = 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 "TestTypes.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
 //===----------------------------------------------------------------------===//
 
+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 targetOperandsMutable();
 }
 
-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(&joinRegion(), getJoinArgs()));
     else
       regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   // The then and else regions are the entry regions of this op.
   regions.push_back(RegionSuccessor(&thenRegion(), getThenArgs()));
   regions.push_back(RegionSuccessor(&elseRegion(), 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();
+  }
+
+  ArrayRef<int64_t> inputShape = inputTy.getShape();
+
+  // Set the first non-dynamic dimension to dynamic.
+  SmallVector<int64_t> shape;
+  bool modified = false;
+  for (int64_t dim : inputShape) {
+    if (!ShapedType::isDynamic(dim) && !modified) {
+      dim = ShapedType::kDynamicSize;
+      modified = true;
+    }
+    shape.push_back(dim);
+  }
+
+  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>:$operands);
-  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);
 }
 
let extraClassDeclaration = [{
 
     std::string getLibraryCallName() {
       return "";
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
-// Test Control-Flow Edge Operand Types
+// Test Type Inferences
 //===----------------------------------------------------------------------===//
 
-def CFEBranchOp : TEST_Op<"cfe.br",
-    [DeclareOpInterfaceMethods<BranchOpInterface,
+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->allowOutputTypeSpecialization();
+    }
+  }];
+}
+
+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