[mlir] Require explicit casts when using TypedValue

[rkayaith]

Summary

Currently TypedValue can be constructed directly from Value, hiding
errors that could be caught at compile time. For example the following
will compile, but crash/assert at runtime:

void foo(TypedValue<IntegerType>);
void bar(TypedValue<FloatType> v) {
  foo(v);
}

This change removes the constructors and replaces them with explicit
llvm casts.

Depends on D142852

[rkayaith]

Fix toy tutorial

mlir/docs/Tutorials/Toy/Ch-4.md

 
 Each of these operations will then need to provide a definition for the
 `inferShapes()` method. As an example, for the mul op, the result shape is
 inferred as the shape of the inputs.
 
 ```c++
 /// Infer the output shape of the MulOp, this is required by the shape inference
 /// interface.
-void MulOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
 ```
 
 At this point, each of the necessary Toy operations provide a mechanism by which
 to infer their output shapes. The ShapeInferencePass will operate on functions:
 it will run on each function in isolation. MLIR also supports general
 [OperationPasses](../../PassManagement.md#operation-pass) that run on any
 isolated operation, but here our module only contains functions, so there is no
 need to generalize to all operations.

mlir/examples/toy/Ch4/mlir/Dialect.cpp

mlir::ParseResult AddOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void AddOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the AddOp, this is required by the shape inference
 /// interface.
-void AddOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void AddOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 /// Infer the output shape of the CastOp, this is required by the shape
 /// inference interface.
-void CastOp::inferShapes() { getResult().setType(getOperand().getType()); }
+void CastOp::inferShapes() { getResult().setType(getInput().getType()); }
 
 /// Returns true if the given set of input and result types are compatible with
 /// this cast operation. This is required by the `CastOpInterface` to verify
 /// this operation and provide other additional utilities.
 bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (inputs.size() != 1 || outputs.size() != 1)
     return false;
   // The inputs must be Tensors with the same element type.
mlir::ParseResult MulOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void MulOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the MulOp, this is required by the shape inference
 /// interface.
-void MulOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult ReturnOp::verify() {
   // We know that the parent operation is a function, because of the 'HasParent'
   // trait attached to the operation definition.

mlir/examples/toy/Ch5/mlir/Dialect.cpp

mlir::ParseResult AddOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void AddOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the AddOp, this is required by the shape inference
 /// interface.
-void AddOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void AddOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 /// Infer the output shape of the CastOp, this is required by the shape
 /// inference interface.
-void CastOp::inferShapes() { getResult().setType(getOperand().getType()); }
+void CastOp::inferShapes() { getResult().setType(getInput().getType()); }
 
 /// Returns true if the given set of input and result types are compatible with
 /// this cast operation. This is required by the `CastOpInterface` to verify
 /// this operation and provide other additional utilities.
 bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (inputs.size() != 1 || outputs.size() != 1)
     return false;
   // The inputs must be Tensors with the same element type.
mlir::ParseResult MulOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void MulOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the MulOp, this is required by the shape inference
 /// interface.
-void MulOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult ReturnOp::verify() {
   // We know that the parent operation is a function, because of the 'HasParent'
   // trait attached to the operation definition.

mlir/examples/toy/Ch6/mlir/Dialect.cpp

mlir::ParseResult AddOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void AddOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the AddOp, this is required by the shape inference
 /// interface.
-void AddOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void AddOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 /// Infer the output shape of the CastOp, this is required by the shape
 /// inference interface.
-void CastOp::inferShapes() { getResult().setType(getOperand().getType()); }
+void CastOp::inferShapes() { getResult().setType(getInput().getType()); }
 
 /// Returns true if the given set of input and result types are compatible with
 /// this cast operation. This is required by the `CastOpInterface` to verify
 /// this operation and provide other additional utilities.
 bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (inputs.size() != 1 || outputs.size() != 1)
     return false;
   // The inputs must be Tensors with the same element type.
mlir::ParseResult MulOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void MulOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the MulOp, this is required by the shape inference
 /// interface.
-void MulOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult ReturnOp::verify() {
   // We know that the parent operation is a function, because of the 'HasParent'
   // trait attached to the operation definition.

mlir/examples/toy/Ch7/mlir/Dialect.cpp

 }
 
 mlir::LogicalResult StructConstantOp::verify() {
   return verifyConstantForType(getResult().getType(), getValue(), *this);
 }
 
 /// Infer the output shape of the ConstantOp, this is required by the shape
 /// inference interface.
-void ConstantOp::inferShapes() { getResult().setType(getValue().getType()); }
+void ConstantOp::inferShapes() {
+  getResult().setType(cast<TensorType>(getValue().getType()));
+}
 
 //===----------------------------------------------------------------------===//
 // AddOp
 //===----------------------------------------------------------------------===//
 
 void AddOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                   mlir::Value lhs, mlir::Value rhs) {
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
   state.addOperands({lhs, rhs});
 }
 
 mlir::ParseResult AddOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void AddOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the AddOp, this is required by the shape inference
 /// interface.
-void AddOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void AddOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 /// Infer the output shape of the CastOp, this is required by the shape
 /// inference interface.
-void CastOp::inferShapes() { getResult().setType(getOperand().getType()); }
+void CastOp::inferShapes() { getResult().setType(getInput().getType()); }
 
 /// Returns true if the given set of input and result types are compatible with
 /// this cast operation. This is required by the `CastOpInterface` to verify
 /// this operation and provide other additional utilities.
 bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (inputs.size() != 1 || outputs.size() != 1)
     return false;
   // The inputs must be Tensors with the same element type.
mlir::ParseResult MulOp::parse(mlir::OpAsmParser &parser,
                                mlir::OperationState &result) {
   return parseBinaryOp(parser, result);
 }
 
 void MulOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
 
 /// Infer the output shape of the MulOp, this is required by the shape inference
 /// interface.
-void MulOp::inferShapes() { getResult().setType(getOperand(0).getType()); }
+void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult ReturnOp::verify() {
   // We know that the parent operation is a function, because of the 'HasParent'
   // trait attached to the operation definition.

mlir/include/mlir/IR/OpDefinition.h

 public:
   /// This class provides return value APIs for ops that are known to have a
   /// single result.  ResultType is the concrete type returned by getType().
   template <typename ConcreteType>
   class Impl
       : public TraitBase<ConcreteType, OneTypedResult<ResultType>::Impl> {
   public:
     TypedValue<ResultType> getResult() {
-      return this->getOperation()->getResult(0);
+      return cast<TypedValue<ResultType>>(this->getOperation()->getResult(0));
     }
 
     /// If the operation returns a single value, then the Op can be implicitly
     /// converted to a Value. This yields the value of the only result.
     operator TypedValue<ResultType>() { return getResult(); }
 
     ResultType getType() { return getResult().getType(); }
   };

mlir/include/mlir/IR/Value.h

if (const auto *outOfLineResult = dyn_cast<OutOfLineOpResult>(this))
     return outOfLineResult->getResultNumber();
   return cast<InlineOpResult>(this)->getResultNumber();
 }
 
 /// TypedValue is a Value with a statically know type.
 /// TypedValue can be null/empty
 template <typename Ty>
 struct TypedValue : Value {
+  using Value::Value;
+
+  static bool classof(Value value) { return llvm::isa<Ty>(value.getType()); }
+
   /// Return the known Type
   Ty getType() { return Value::getType().template cast<Ty>(); }
-  void setType(mlir::Type ty) {
-    assert(ty.template isa<Ty>());
-    Value::setType(ty);
-  }
-
-  TypedValue(Value val) : Value(val) {
-    assert(!val || val.getType().template isa<Ty>());
-  }
-  TypedValue &operator=(const Value &other) {
-    assert(!other || other.getType().template isa<Ty>());
-    Value::operator=(other);
-    return *this;
-  }
+  void setType(Ty ty) { Value::setType(ty); }
 };
 
 } // namespace detail
 
 /// This is a value defined by a result of an operation.
 class OpResult : public Value {
 public:
   using Value::Value;

mlir/include/mlir/TableGen/Class.h

public:
   /// A method that makes another method redundant with a different return type
   /// can replace the other, the assumption being that the subsuming method
   /// provides a more resolved return type, e.g. IntegerAttr vs. Attribute.
   bool makesRedundant(const MethodSignature &other) const;
 
   /// Get the name of the method.
   StringRef getName() const { return methodName; }
 
+  /// Get the return type of the method
+  StringRef getReturnType() const { return returnType; }
+
   /// Get the number of parameters.
   unsigned getNumParameters() const { return parameters.getNumParameters(); }
 
   /// Write the signature as part of a method declaration.
   void writeDeclTo(raw_indented_ostream &os) const;
 
   /// Write the signature as part of a method definition. `namePrefix` is to be
   /// prepended to the method name (typically namespaces for qualifying the
public:
   bool isConstructor() const { return properties & Constructor; }
 
   /// Returns true if this class method is const.
   bool isConst() const { return properties & Const; }
 
   /// Returns the name of this method.
   StringRef getName() const { return methodSignature.getName(); }
 
+  /// Returns the return type of this method
+  StringRef getReturnType() const { return methodSignature.getReturnType(); }
+
   /// Returns if this method makes the `other` method redundant.
   bool makesRedundant(const Method &other) const {
     return methodSignature.makesRedundant(other.methodSignature);
   }
 
   /// Write the method declaration, including the definition if inline.
   void writeDeclTo(raw_indented_ostream &os) const override;
 

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

SmallVector<AffineMap, 2> affineMaps = {
         rewriter.getMultiDimIdentityMap(resultTy.getRank())};
 
     rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         op, resultTy, ArrayRef<Value>({}), ValueRange{emptyTensor}, affineMaps,
         getNParallelLoopsAttrs(resultTy.getRank()),
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           llvm::SmallVector<Value> indices;
           for (unsigned int i = 0; i < inputTy.getRank(); i++) {
-            auto index =
+            Value index =
                 rewriter.create<linalg::IndexOp>(nestedLoc, i).getResult();
             if (i == axis) {
               auto one = rewriter.create<arith::ConstantIndexOp>(nestedLoc, 1);
               auto sizeMinusOne =
                   rewriter.create<arith::SubIOp>(nestedLoc, axisDimSize, one);
               index = rewriter.create<arith::SubIOp>(nestedLoc, sizeMinusOne,
                                                      index);
             }

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

LogicalResult matchAndRewrite(TransferWriteOp xferOp,
     // Transfer ops that modify the element type are not supported atm.
     if (xferOp.getVectorType().getElementType() !=
         xferOp.getShapedType().getElementType())
       return failure();
 
     auto vec = getDataVector(xferOp);
     auto xferVecType = xferOp.getVectorType();
     int64_t dimSize = xferVecType.getShape()[0];
-    auto source = xferOp.getSource(); // memref or tensor to be written to.
+    Value source = xferOp.getSource(); // memref or tensor to be written to.
     auto sourceType = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();
 
     // Generate fully unrolled loop of transfer ops.
     Location loc = xferOp.getLoc();
     for (int64_t i = 0; i < dimSize; ++i) {
       Value iv = rewriter.create<arith::ConstantIndexOp>(loc, i);
 
       auto updatedSource = generateInBoundsCheck(

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

LogicalResult matchAndRewrite(NewOp op,
     //   %tmp = bufferization.alloc_tensor : an unordered COO with identity
     //                                       storage ordering
     //   for i = 0 to nnz
     //     get the next element from the input file
     //     insert the element to %tmp
     //   %t = sparse_tensor.ConvertOp %tmp
     RankedTensorType cooTp =
         getUnorderedCOOFromTypeWithOrdering(dstTp, encDst.getDimOrdering());
-    auto cooBuffer =
+    Value cooBuffer =
         rewriter.create<AllocTensorOp>(loc, cooTp, dynSizesArray).getResult();
 
     Value c0 = constantIndex(rewriter, loc, 0);
     Value c1 = constantIndex(rewriter, loc, 1);
     Value nnz = createFuncCall(rewriter, loc, "getSparseTensorReaderNNZ",
                                {indexTp}, {reader}, EmitCInterface::Off)
                     .getResult(0);
     Value symmetric;

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

public:
   }
 
 private:
   SparseTensorEncodingAttr enc;
 };
 
 class SparseTensorSpecifier {
 public:
-  explicit SparseTensorSpecifier(Value specifier) : specifier(specifier) {}
+  explicit SparseTensorSpecifier(Value specifier)
+      : specifier(cast<TypedValue<StorageSpecifierType>>(specifier)) {}
 
   // Undef value for dimension sizes, all zero value for memory sizes.
   static Value getInitValue(OpBuilder &builder, Location loc,
                             RankedTensorType rtp);
 
   /*implicit*/ operator Value() { return specifier; }
 
   Value getSpecifierField(OpBuilder &builder, Location loc,

mlir/test/mlir-tblgen/op-operand.td

 def OpD : NS_Op<"mix_variadic_and_normal_inputs_op", [SameVariadicOperandSize]> {
   let arguments = (ins Variadic<AnyTensor>:$input1, AnyTensor:$input2, Variadic<AnyTensor>:$input3);
 }
 
 // CHECK-LABEL: ::mlir::Operation::operand_range OpD::getInput1
 // CHECK-NEXT: return getODSOperands(0);
 
 // CHECK-LABEL: ::mlir::TypedValue<::mlir::TensorType> OpD::getInput2
-// CHECK-NEXT: return *getODSOperands(1).begin();
+// CHECK-NEXT: return ::llvm::cast<::mlir::TypedValue<::mlir::TensorType>>(*getODSOperands(1).begin());
 
 // CHECK-LABEL: OpD::build
 // CHECK-NEXT: odsState.addOperands(input1);
 // CHECK-NEXT: odsState.addOperands(input2);
 // CHECK-NEXT: odsState.addOperands(input3);

mlir/test/mlir-tblgen/op-result.td

 def OpI : NS_Op<"mix_variadic_and_normal_results_op", [SameVariadicResultSize]> {
   let results = (outs Variadic<AnyTensor>:$output1, AnyTensor:$output2, Variadic<AnyTensor>:$output3);
 }
 
 // CHECK-LABEL: ::mlir::Operation::result_range OpI::getOutput1
 // CHECK-NEXT:    return getODSResults(0);
 
 // CHECK-LABEL: ::mlir::TypedValue<::mlir::TensorType> OpI::getOutput2
-// CHECK-NEXT:    return *getODSResults(1).begin();
+// CHECK-NEXT:    return ::llvm::cast<::mlir::TypedValue<::mlir::TensorType>>(*getODSResults(1).begin());
 
 // CHECK-LABEL: OpI::build
 // CHECK-NEXT:    odsState.addTypes(output1);
 // CHECK-NEXT:    odsState.addTypes(output2);
 // CHECK-NEXT:    odsState.addTypes(output3);
 
 // Test that if the only operand is variadic, we access the first value in the
 // pack to set result type

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

if (operand.name.empty())
       continue;
     std::string name = op.getGetterName(operand.name);
     if (operand.isOptional()) {
       m = opClass.addMethod(isGenericAdaptorBase
                                 ? rangeElementType
                                 : generateTypeForGetter(operand),
                             name);
       ERROR_IF_PRUNED(m, name, op);
-      m->body().indent() << formatv(
-          "auto operands = getODSOperands({0});\n"
-          "return operands.empty() ? {1}{{} : *operands.begin();",
-          i, rangeElementType);
+      m->body().indent() << formatv("auto operands = getODSOperands({0});\n"
+                                    "return operands.empty() ? {1}{{} : ",
+                                    i, m->getReturnType());
+      if (!isGenericAdaptorBase)
+        m->body() << llvm::formatv("::llvm::cast<{0}>", m->getReturnType());
+      m->body() << "(*operands.begin());";
     } else if (operand.isVariadicOfVariadic()) {
       std::string segmentAttr = op.getGetterName(
           operand.constraint.getVariadicOfVariadicSegmentSizeAttr());
       if (genericAdaptorBase) {
         m = opClass.addMethod("::llvm::SmallVector<" + rangeType + ">", name);
         ERROR_IF_PRUNED(m, name, op);
         m->body() << llvm::formatv(variadicOfVariadicAdaptorCalcCode,
                                    segmentAttr, i, rangeType);
if (operand.isOptional()) {
       ERROR_IF_PRUNED(m, name, op);
       m->body() << "  return getODSOperands(" << i << ");";
     } else {
       m = opClass.addMethod(isGenericAdaptorBase
                                 ? rangeElementType
                                 : generateTypeForGetter(operand),
                             name);
       ERROR_IF_PRUNED(m, name, op);
-      m->body() << "  return *getODSOperands(" << i << ").begin();";
+      m->body().indent() << "return ";
+      if (!isGenericAdaptorBase)
+        m->body() << llvm::formatv("::llvm::cast<{0}>", m->getReturnType());
+      m->body() << llvm::formatv("(*getODSOperands({0}).begin());", i);
     }
   }
 }
 
 void OpEmitter::genNamedOperandGetters() {
   // Build the code snippet used for initializing the operand_segment_size)s
   // array.
   std::string attrSizeInitCode;
void OpEmitter::genNamedResultGetters() {
   for (int i = 0; i != numResults; ++i) {
     const auto &result = op.getResult(i);
     if (result.name.empty())
       continue;
     std::string name = op.getGetterName(result.name);
     if (result.isOptional()) {
       m = opClass.addMethod(generateTypeForGetter(result), name);
       ERROR_IF_PRUNED(m, name, op);
-      m->body()
-          << "  auto results = getODSResults(" << i << ");\n"
-          << "  return results.empty() ? ::mlir::Value() : *results.begin();";
+      m->body() << "  auto results = getODSResults(" << i << ");\n"
+                << llvm::formatv("  return results.empty()"
+                                 " ? {0}()"
+                                 " : ::llvm::cast<{0}>(*results.begin());",
+                                 m->getReturnType());
     } else if (result.isVariadic()) {
       m = opClass.addMethod("::mlir::Operation::result_range", name);
       ERROR_IF_PRUNED(m, name, op);
       m->body() << "  return getODSResults(" << i << ");";
     } else {
       m = opClass.addMethod(generateTypeForGetter(result), name);
       ERROR_IF_PRUNED(m, name, op);
-      m->body() << "  return *getODSResults(" << i << ").begin();";
+      m->body() << llvm::formatv(
+          "  return ::llvm::cast<{0}>(*getODSResults({1}).begin());",
+          m->getReturnType(), i);
     }
   }
 }
 
 void OpEmitter::genNamedRegionGetters() {
   unsigned numRegions = op.getNumRegions();
   for (unsigned i = 0; i < numRegions; ++i) {
     const auto &region = op.getRegion(i);

mlir/unittests/IR/IRMapping.cpp

TEST(IRMapping, TypedValue) {
 
   Value i64Val = builder.create<test::TestOpConstant>(
       loc, builder.getI64Type(), builder.getI64IntegerAttr(0));
   Value f64Val = builder.create<test::TestOpConstant>(
       loc, builder.getF64Type(), builder.getF64FloatAttr(0.0));
 
   IRMapping mapping;
   mapping.map(i64Val, f64Val);
-  TypedValue<IntegerType> typedI64Val = i64Val;
+  auto typedI64Val = cast<TypedValue<IntegerType>>(i64Val);
   EXPECT_EQ(mapping.lookup(typedI64Val), f64Val);
 }
 
 TEST(IRMapping, OperationClone) {
   MLIRContext ctx;
   ctx.allowUnregisteredDialects();
 
   OperationState state(UnknownLoc::get(&ctx), "no_results");