[mlir] Change to requiring location for infer*

[jpienaar]

Summary

Previously it was optional to cater for potential future cases where we'd not
want errors reported while still verifying types. That lead to some
inconvenient usage patterns inside inference function. One can get the same
behavior via the diagnostic handler and given the vast majority of uses has a
location, flip it so that the common case is the easy case.

[rriddle]

LGTM in general, though one comment about the diagnostic handler.

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

103

I don't think we have a general utility for this. We would need to add something like a NullDiagnosticHandler or DiscardingDiagnosticHandler, that only ignores diagnostics from the current thread(or child threads perhaps?).

[jpienaar]

Not pursuing this one at the moment

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

 
 namespace detail {
 // Helper function to infer return tensor returns types given element and shape
 // inference function.
 //
 // TODO: Consider generating typedefs for trait member functions if this usage
 // becomes more common.
 LogicalResult inferReturnTensorTypes(
-    function_ref<LogicalResult(
-        MLIRContext *, Optional<Location> location, ValueRange operands,
+    function_ref<
+        LogicalResult(MLIRContext *, Location location, ValueRange operands,
                       DictionaryAttr attributes, RegionRange regions,
                       SmallVectorImpl<ShapedTypeComponents> &retComponents)>
         componentTypeFn,
-    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    MLIRContext *context, Location location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes);
 
 /// Verifies that the inferred result types match the actual result types for
 /// the op. Precondition: op implements InferTypeOpInterface.
 LogicalResult verifyInferredResultTypes(Operation *op);
 } // namespace detail
 
 namespace OpTrait {
 
 /// Tensor type inference trait that constructs a tensor from the inferred
 /// shape and elemental types.
 /// Requires: Op implements functions of InferShapedTypeOpInterface.
 template <typename ConcreteType>
 class InferTensorType : public TraitBase<ConcreteType, InferTensorType> {
 public:
   static LogicalResult
-  inferReturnTypes(MLIRContext *context, Optional<Location> location,
-                   ValueRange operands, DictionaryAttr attributes,
-                   RegionRange regions,
+  inferReturnTypes(MLIRContext *context, Location location, ValueRange operands,
+                   DictionaryAttr attributes, RegionRange regions,
                    SmallVectorImpl<Type> &inferredReturnTypes) {
     return ::mlir::detail::inferReturnTensorTypes(
         ConcreteType::inferReturnTypeComponents, context, location, operands,
         attributes, regions, inferredReturnTypes);
   }
 };
 
 } // namespace OpTrait
 } // namespace mlir
 
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/InferTypeOpInterface.h.inc"
 
 #endif // MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

let description = [{
     during op construction, verification or type inference.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     StaticInterfaceMethod<
       /*desc=*/[{Infer the return types that an op would generate.
 
-      The method takes an optional location which, if set, will be used to
-      report errors on. The operands and attributes correspond to those with
-      which an Operation would be created (e.g., as used in Operation::create)
-      and the regions of the op.
+      The method takes a location which 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.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"inferReturnTypes",
       /*args=*/(ins "::mlir::MLIRContext *":$context,
-                    "::llvm::Optional<::mlir::Location>":$location,
+                    "::mlir::Location":$location,
                     "::mlir::ValueRange":$operands,
                     "::mlir::DictionaryAttr":$attributes,
                     "::mlir::RegionRange":$regions,
                     "::llvm::SmallVectorImpl<::mlir::Type>&":$inferredReturnTypes)
     >,
     StaticInterfaceMethod<
       /*desc=*/"Returns whether two array of types are compatible result types"
                " for an op.",
let description = [{
     The components consists of element type, shape and raw attribute.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     StaticInterfaceMethod<
       /*desc=*/[{Infer the components of return type of shape containter.
 
-      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.
+      The method takes a location which 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.
 
       Unknown (e.g., unranked) shape and nullptrs for element type and attribute
       may be returned by this function while returning success. E.g., partial
       population of components is not error condition.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"inferReturnTypeComponents",
       /*args=*/(ins "::mlir::MLIRContext*":$context,
-                    "::mlir::Optional<::mlir::Location>":$location,
+                    "::mlir::Location":$location,
                     "::mlir::ValueRange":$operands,
                     "::mlir::DictionaryAttr":$attributes,
                     "::mlir::RegionRange":$regions,
                     "::mlir::SmallVectorImpl<::mlir::ShapedTypeComponents>&":
                       $inferredReturnShapes)
     >,
     InterfaceMethod<
       /*desc=*/[{Reify the shape computation for the operation.

mlir/lib/Interfaces/InferTypeOpInterface.cpp

 
 using namespace mlir;
 
 namespace mlir {
 #include "mlir/Interfaces/InferTypeOpInterface.cpp.inc"
 } // namespace mlir
 
 LogicalResult mlir::detail::inferReturnTensorTypes(
-    function_ref<LogicalResult(
-        MLIRContext *, Optional<Location> location, ValueRange operands,
+    function_ref<
+        LogicalResult(MLIRContext *, Location location, ValueRange operands,
                       DictionaryAttr attributes, RegionRange regions,
                       SmallVectorImpl<ShapedTypeComponents> &retComponents)>
         componentTypeFn,
-    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    MLIRContext *context, Location location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   SmallVector<ShapedTypeComponents, 2> retComponents;
   if (failed(componentTypeFn(context, location, operands, attributes, regions,
                              retComponents)))
     return failure();
   for (auto shapeAndType : retComponents) {
     assert(shapeAndType.getAttribute() == nullptr && "attribute not supported");

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

OpFoldResult TestOpInPlaceFold::fold(ArrayRef<Attribute> operands) {
   return {};
 }
 
 OpFoldResult TestPassthroughFold::fold(ArrayRef<Attribute> operands) {
   return getOperand();
 }
 
 LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(
-    MLIRContext *, Optional<Location> location, ValueRange operands,
+    MLIRContext *, Location location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   if (operands[0].getType() != operands[1].getType()) {
     return emitOptionalError(location, "operand type mismatch ",
                              operands[0].getType(), " vs ",
                              operands[1].getType());
   }
   inferredReturnTypes.assign({operands[0].getType()});
   return success();
 }
 
 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
-    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    MLIRContext *context, Location location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   // Create return type consisting of the last element of the first operand.
   auto operandType = *operands.getTypes().begin();
   auto sval = operandType.dyn_cast<ShapedType>();
   if (!sval) {
     return emitOptionalError(location, "only shaped type operands allowed");
   }

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

 // ReturnType Driver.
 //===----------------------------------------------------------------------===//
 
 namespace {
 // Generate ops for each instance where the type can be successfully inferred.
 template <typename OpTy>
 static void invokeCreateWithInferredReturnType(Operation *op) {
   auto *context = op->getContext();
+  // Squash error emission/propagation.
+  DiagnosticEngine::HandlerID diagHandler =
+      context->getDiagEngine().registerHandler([](Diagnostic &) {});
   auto fop = op->getParentOfType<FuncOp>();
   auto location = UnknownLoc::get(context);
   OpBuilder b(op);
   b.setInsertionPointAfter(op);
 
   // Use permutations of 2 args as operands.
   assert(fop.getNumArguments() >= 2);
   for (int i = 0, e = fop.getNumArguments(); i < e; ++i) {
     for (int j = 0; j < e; ++j) {
       std::array<Value, 2> values = {{fop.getArgument(i), fop.getArgument(j)}};
       SmallVector<Type, 2> inferredReturnTypes;
       if (succeeded(OpTy::inferReturnTypes(
-              context, llvm::None, values, op->getAttrDictionary(),
+              context, op->getLoc(), values, op->getAttrDictionary(),
               op->getRegions(), inferredReturnTypes))) {
         OperationState state(location, OpTy::getOperationName());
         // TODO: Expand to regions.
         OpTy::build(b, state, values, op->getAttrs());
         (void)b.createOperation(state);
       }
     }
   }
+  context->getDiagEngine().eraseHandler(diagHandler);
 }
 
 static void reifyReturnShape(Operation *op) {
   OpBuilder b(op);
 
   // Use permutations of 2 args as operands.
   auto shapedOp = cast<OpWithShapedTypeInferTypeInterfaceOp>(op);
   SmallVector<Value, 2> shapes;