[mlir] Add a verifier for `shape.yield`.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let arguments = (ins Shape_SizeType:$arg);
   let results = (outs Index:$result);
 
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
 }
 
-def Shape_YieldOp : Shape_Op<"yield", [NoSideEffect, Terminator]> {
+def Shape_YieldOp : Shape_Op<"yield", 
+    [HasParent<"ReduceOp">,
+     NoSideEffect,
+     ReturnLike,
+     Terminator]> {
   let summary = "Returns the value to parent op";
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result", [{ build(b, result, llvm::None); }]
   >];
 
+  let verifier = [{ return ::verify(*this); }];
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 // TODO: Add Ops: if_static, if_ranked
 
 // For testing usage.
 def Shape_DebugPrintOp : Shape_Op<"debug_print", []> {
   let summary = "Prints the input shape or size";

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

OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])
     return arg;
   return {};
 }
 
 //===----------------------------------------------------------------------===//
+// YieldOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(YieldOp op) {
+  auto *parentOp = op.getParentOp();
+  auto results = parentOp->getResults();
+  auto operands = op.getOperands();
+
+  if (parentOp->getNumResults() != op.getNumOperands())
+    return op.emitOpError() << "number of operands does not match number of "
+                               "results of its parent";
+  for (auto e : llvm::zip(results, operands))
+    if (std::get<0>(e).getType() != std::get<1>(e).getType())
+      return op.emitOpError()
+             << "types mismatch between yield op and its parent";
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // SplitAtOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult SplitAtOp::fold(ArrayRef<Attribute> operands,
                               SmallVectorImpl<OpFoldResult> &results) {
   if (!operands[0] || !operands[1])
     return failure();
   auto shapeVec = llvm::to_vector<6>(

mlir/test/Dialect/Shape/invalid.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 func @reduce_op_args_num_mismatch(%shape : !shape.shape, %init : !shape.size) {
   // expected-error@+1 {{ReduceOp body is expected to have 3 arguments}}
   %num_elements = shape.reduce(%shape, %init) -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size):
-      "shape.yield"(%dim) : (!shape.size) -> ()
+      shape.yield %dim : !shape.size
   }
 }
 
 // -----
 
 func @reduce_op_arg0_wrong_type(%shape : !shape.shape, %init : !shape.size) {
   // expected-error@+1 {{argument 0 of ReduceOp body is expected to be of IndexType}}
   %num_elements = shape.reduce(%shape, %init) -> !shape.size {
-    ^bb0(%index: f32, %dim: !shape.size, %lci: !shape.size):
-      %acc = "shape.add"(%lci, %dim) : (!shape.size, !shape.size) -> !shape.size
-      "shape.yield"(%acc) : (!shape.size) -> ()
+    ^bb0(%index: f32, %dim: !shape.size, %acc: !shape.size):
+      %new_acc = "shape.add"(%acc, %dim)
+          : (!shape.size, !shape.size) -> !shape.size
+      shape.yield %new_acc : !shape.size
   }
 }
 
 // -----
 
 func @reduce_op_arg1_wrong_type(%shape : !shape.shape, %init : !shape.size) {
   // expected-error@+1 {{argument 1 of ReduceOp body is expected to be of SizeType}}
   %num_elements = shape.reduce(%shape, %init) -> !shape.size {
     ^bb0(%index: index, %dim: f32, %lci: !shape.size):
-      "shape.yield"() : () -> ()
+      shape.yield
   }
 }
 
 // -----
 
 func @reduce_op_init_type_mismatch(%shape : !shape.shape, %init : f32) {
   // expected-error@+1 {{type mismatch between argument 2 of ReduceOp body and initial value 0}}
   %num_elements = shape.reduce(%shape, %init) -> f32 {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
-      "shape.yield"() : () -> ()
+      shape.yield
+  }
+}
+
+// -----
+
+func @yield_op_args_num_mismatch(%shape : !shape.shape, %init : !shape.size) {
+  // expected-error@+3 {{number of operands does not match number of results of its parent}}
+  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+    ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
+      shape.yield %dim, %dim : !shape.size, !shape.size
+  }
+}
+
+// -----
+
+func @yield_op_type_mismatch(%shape : !shape.shape, %init : !shape.size) {
+  // expected-error@+4 {{types mismatch between yield op and its parent}}
+  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+    ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
+      %c0 = constant 1 : index
+      shape.yield %c0 : index
   }
 }

mlir/test/Dialect/Shape/ops.mlir

 // RUN: mlir-opt -split-input-file %s | mlir-opt | FileCheck %s --dump-input-on-failure
 // Verify the printed output can be parsed.
 // RUN: mlir-opt %s | mlir-opt | FileCheck %s --dump-input-on-failure
 // Verify the generic form can be parsed.
 // RUN: mlir-opt -mlir-print-op-generic %s | mlir-opt | FileCheck %s --dump-input-on-failure
 
 // CHECK-LABEL: shape_num_elements
 func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
   %init = shape.const_size 0
   %num_elements = shape.reduce(%shape, %init) -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
       %acc = "shape.add"(%lci, %dim) : (!shape.size, !shape.size) -> !shape.size
-      "shape.yield"(%acc) : (!shape.size) -> ()
+      shape.yield %acc : !shape.size
   }
   return %num_elements : !shape.size
 }
 
 func @test_shape_num_elements_unknown() {
   %0 = "shape.unknown_shape"() : () -> !shape.shape
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %2 = "shape.print"(%1) : (!shape.size) -> !shape.size