[mlir][Linalg] Add ReduceOp to Linalg structured ops.

[akuegel]

Summary

This will allow to model (variadic) reductions with this special op instead of
using GenericOp.

RFC: https://discourse.llvm.org/t/rfc-primitive-ops-add-mapop-reductionop-transposeop-broadcastop-to-linalg/64184

[nicolasvasilache]

Thanks @akuegel for pushing in this direction!

I added a first round of comments.

Upon deeper inspection, I find that none of the tests exercise mixing ? and shape, I'd drop it from this PR for now and we can introduce later if needed (I suspect this would trigger discussions).

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

266

getIteratorTypes ?

270

getOutputs?

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1190

getIteratorTypes?

1208

I'd just start from getMultiDimIdentityMap and apply AffineMap dropResult(unsigned pos) for all getDimensions.

Bonus points: impl AffineMap dropResults(ArrayRef<unsigned> pos)

1245

do we want to just ignore stuff that is not what we expect ?

1306

you seem to want to allow the mixing .. plz see my message below.

1314

This allows mixing static and dynamic shapes, are you sure this is what you want?

If so, this needs to be well-documented in the op semantics.
It will be surprising if someone mixes 4 and ? and at runtime ? != 4 => this should be clearly spelled out as UB.

If this is not the behavior you want, just do a.getShape() == b.getShape() ?

Looking back at the RFC, we didn't discuss allowing implicit conversions to/from ?, could you comment on the tradeoffs?

1323

same comment

1337

can we use some Tablegen "confined" spec here and avoid reinventing?

1342

feels like a generally useful Tablegen-level property to add (can be left for later).

1375

if (getBody()->getNumArguments() != this->getNumOperands())
  return emitOpError() << "mismatching number of operands and block arguments";

1384

// Note: zip iterates until the min length.
for (auto [input, bbArg] :: llvm::zip(getInputs(), block->getArguments())) {
  if (input.getType().cast<ShapedType>().getElementType() != bbArg.getType())
    return ...
}

1398

Same

for (auto [..., bbArg]

[pifon2a]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1314

That's because we never discussed variadic reductions in the RFC. HLO allows having just "compatible" shapes among inputs and among outputs, but it is definitely not clear if it's needed.

[akuegel]

Address review comments.

[akuegel]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

270

Done, but this requires changing the DestinationStyleOpInterface. Please let me know whether I should move that change to a separate patch. Changing the name of this interface method would also require updates to users outside MLIR, e.g. tHLO in Tensorflow.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1208

I have added a dropResults method, however to make it easily usable for my case, I needed to use ArrayRef<int64_t>. To make both methods consistent, I also changed dropResult() to take int64_t. But maybe unsigned types are preferred here, and I should instead cast my int64_t ArrayRef?

1245

Good point, I guess it makes sense to return a failure. Done.

1306

Yes, not really necessary. I have removed this method.

1314

Lets get rid of this. I think we don't need to allow mismatch of dynamic and static dimensions.

1337

I guess you mean specifying dimensions as a ConfinedAttr?

I tried the following:

class IndexInBounds<int upperBound>
    : AttrConstraint<CPred<"llvm::all_of($_self.cast<::mlir::DenseI64ArrayAttr>"
                           "().asArrayRef(), [](int64_t dim) { return dim >= 0 "
                           "&& dim < " # upperBound # "; })">,
                     "should be in the range [0, " # upperBound # "]">;

ConfinedAttr<DenseI64ArrayAttr, [IndexInBounds<2>]>:$dimensions

But in fact I cannot use a constant upper bound. How can I pass a dynamic upper bound (the input rank) to IndexInBounds? Is it even possible? And if it is not possible, should I just verify that each value is >= 0 and keep the other verification code inside the verify method?

1342

So here I can see how it can be done as AttrConstraint. If I understand correctly, you would prefer if such a constraint is available for reuse. Where would be a good place for it in the code base?

[akuegel]

Rebase.

[akuegel]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

270

Landed this change of the interface in https://reviews.llvm.org/rGdbf1fe5024e205f295ed0280d602d582774c39c9

[nicolasvasilache]

Nice!

Minor comments (and maybe a bug)

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

270

thanks for putting this in a separate patch!

mlir/include/mlir/IR/AffineMap.h

254

++

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1208

Generally we should avoid unsigned as much as possible (see e.g. https://www.youtube.com/watch?v=Puio5dly9N8), there are a bunch of older APIs based on unsigned that will need to be cleaned up but avoiding to propagate that mistake is better, this is great as is.

1337

Had the same issue in https://reviews.llvm.org/D130348 actually..

At the time, @Mogball wrote: This can be achieved with a derived attribute that uses the value of an optional attribute or some derived value
I wasn't able to immediately turn that into something actionnable; if you see a way to please go ahead otherwise it's fine with me to leave as a future cleanup.

1342

I would think something in OpBase.td would make sense to implement the "sorted" constraint.

1377

wrong input/output mixing and you are getting lucky with the current tests :) ?

[akuegel]

Fix copy/paste bug in verify.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1337

I also don't understand what is meant with derived attribute.
So I will leave this as a future cleanup.

1342

Ok, I will prepare a separate patch for that.

1377

Good catch! I have adjusted the test for mismatched outputs a bit so that it would catch this bug.

[nicolasvasilache]

Landed as fa79dff8bc8b2da923d4935f4b2035f7e1e11e0a

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

def GenericOp : LinalgStructuredBase_Op<"generic", [AttrSizedOperandSegments]> {
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 
 //===----------------------------------------------------------------------===//
+// Reduce op.
+//===----------------------------------------------------------------------===//
+
+def TensorOrMemref :
+  AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;
+
+def ReduceOp : LinalgStructuredBase_Op<"reduce", [
+      SameVariadicOperandSize, SingleBlockImplicitTerminator<"YieldOp">
+    ]> {
+  let summary = "Reduce operator";
+  let description = [{
+    Executes `combiner` on the `dimensions` of `inputs` and returns the
+    reduced result. The `dimensions` attribute needs to list the reduction
+    dimensions in increasing order.
+
+    Example:
+    ```
+      %reduce = linalg.reduce
+          ins(%input:tensor<16x32x64xf32>)
+          outs(%init:tensor<16x64xf32>)
+          dimensions = [1]
+          (%in: f32, %out: f32) {
+            %0 = arith.addf %in, %out: f32
+            linalg.yield %0: f32
+          }
+    ```
+  }];
+
+  let arguments = (ins
+    // Input arg
+    Variadic<TensorOrMemref>:$inputs,
+    // Output arg
+    Variadic<TensorOrMemref>:$inits,
+
+    DenseI64ArrayAttr:$dimensions
+  );
+  let results = (outs Variadic<TensorOrMemref>);
+  let regions = (region SizedRegion<1>:$combiner);
+
+  let extraClassDeclaration = structuredOpsBaseDecls # [{
+    // Declare functions necessary for LinalgStructuredInterface.
+    ArrayAttr getIteratorTypes();
+    ArrayAttr getIndexingMaps();
+
+    // Implement functions necessary for DestinationStyleOpInterface.
+    mlir::ValueRange getOutputs() { return getInits(); }
+    unsigned getNumInputs() { return getInputs().size(); };
+    unsigned getNumOutputs() { return getInits().size(); };
+    static std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
+                              mlir::ArrayRef<mlir::NamedAttribute>)>
+    getRegionBuilder() {
+      return nullptr;
+    }
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
+
+//===----------------------------------------------------------------------===//
 // Named Linalg ops, implemented as a declarative configurations of generic ops.
 //===----------------------------------------------------------------------===//
 
 include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.td"
 
 #endif // LINALG_STRUCTURED_OPS

mlir/include/mlir/IR/AffineMap.h

return AffineMap::get(getNumDims(), getNumSymbols() + shift,
                                 return e.shiftSymbols(getNumSymbols(), shift,
                                                       offset);
                               })),
                           getContext());
   }
 
   /// Returns a new AffineMap with the same number of dims and symbols and one
   /// less result at `pos`, dropped.
-  AffineMap dropResult(unsigned pos) {
+  AffineMap dropResult(int64_t pos) {
     auto exprs = llvm::to_vector<4>(getResults());
     exprs.erase(exprs.begin() + pos);
     return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
   }
 
+  // Returns a new AffineMap with the same number of dims and symbols, but all
+  // positions in `positions` dropped from results.
+  AffineMap dropResults(ArrayRef<int64_t> positions) {
+    AffineMap resultMap = *this;
+    for (int64_t pos : positions)
+      resultMap = resultMap.dropResult(pos);
+    return resultMap;
+  }
+
   /// Returns a new AffineMap with the same number of dims and symbols and an
   /// extra result inserted at `pos`.
   AffineMap insertResult(AffineExpr expr, unsigned pos) {
     auto exprs = llvm::to_vector<4>(getResults());
     exprs.insert(exprs.begin() + pos, expr);
     return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
   }
 

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

fillStructuredOpRegion(b, region, TypeRange(inputs), TypeRange(outputs),
                          state.attributes.getAttrs(), regionBuilder);
 }
 
 /// Common parsing used for both named structured ops created by ods-gen and by
 /// manually defined C++ ops. Does not handle regions.
 static ParseResult
 parseCommonStructuredOpParts(OpAsmParser &parser, OperationState &result,
                              SmallVectorImpl<Type> &inputTypes,
-                             SmallVectorImpl<Type> &outputTypes) {
+                             SmallVectorImpl<Type> &outputTypes,
+                             bool addOperandSegmentSizes = true) {
   SMLoc inputsOperandsLoc, outputsOperandsLoc;
   SmallVector<OpAsmParser::UnresolvedOperand, 4> inputsOperands,
       outputsOperands;
 
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   if (succeeded(parser.parseOptionalKeyword("ins"))) {
parseCommonStructuredOpParts(OpAsmParser &parser, OperationState &result,
   }
 
   if (parser.resolveOperands(inputsOperands, inputTypes, inputsOperandsLoc,
                              result.operands) ||
       parser.resolveOperands(outputsOperands, outputTypes, outputsOperandsLoc,
                              result.operands))
     return failure();
 
+  if (addOperandSegmentSizes) {
     result.addAttribute("operand_segment_sizes",
                         parser.getBuilder().getDenseI32ArrayAttr(
                             {static_cast<int32_t>(inputsOperands.size()),
                              static_cast<int32_t>(outputsOperands.size())}));
+  }
   return success();
 }
 
 static void printCommonStructuredOpParts(OpAsmPrinter &p, ValueRange inputs,
                                          ValueRange outputs) {
   if (!inputs.empty())
     p << " ins(" << inputs << " : " << inputs.getTypes() << ")";
   if (!outputs.empty())

 }
 
 LogicalResult GenericOp::fold(ArrayRef<Attribute>,
                               SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
+// ReduceOp
+//===----------------------------------------------------------------------===//
+
+ArrayAttr ReduceOp::getIteratorTypes() {
+  int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
+  SmallVector<StringRef> iteratorTypes(inputRank,
+                                       getParallelIteratorTypeName());
+  for (int64_t reductionDim : getDimensions())
+    iteratorTypes[reductionDim] = getReductionIteratorTypeName();
+  return Builder(getContext()).getStrArrayAttr(iteratorTypes);
+}
+
+ArrayAttr ReduceOp::getIndexingMaps() {
+  int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
+  SmallVector<AffineMap> affineMaps(
+      getNumInputs(),
+      AffineMap::getMultiDimIdentityMap(inputRank, getContext()));
+  AffineMap resultMap =
+      AffineMap::getMultiDimIdentityMap(inputRank, getContext())
+          .dropResults(getDimensions());
+  for (int64_t i = 0, e = getNumOutputs(); i < e; ++i)
+    affineMaps.push_back(resultMap);
+  return Builder(getContext()).getAffineMapArrayAttr(affineMaps);
+}
+
+void ReduceOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  SmallVector<Value> inputBuffers = getInputBufferOperands();
+  SmallVector<Value> outputBuffers = getOutputBufferOperands();
+  getGenericEffectsImpl(effects, getOperation()->getResults(), inputBuffers,
+                        outputBuffers);
+}
+
+static ParseResult parseDstStyleOp(
+    OpAsmParser &parser, OperationState &result,
+    function_ref<ParseResult(OpAsmParser &, NamedAttrList &)> parseAttrsFn =
+        nullptr) {
+  // Parse `ins` and `outs`.
+  SmallVector<Type, 4> inputTypes, outputTypes;
+  if (parseCommonStructuredOpParts(parser, result, inputTypes, outputTypes,
+                                   /*addOperandSegmentSizes=*/false))
+    return failure();
+
+  // Add result types.
+  for (Type outputType : outputTypes) {
+    if (!outputType.isa<RankedTensorType>())
+      return failure();
+    result.addTypes(outputType);
+  }
+
+  // Parse required attributes.
+  if (parseAttrsFn && failed(parseAttrsFn(parser, result.attributes)))
+    return failure();
+
+  // Parse optional attributes.
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+  return success();
+}
+
+static ParseResult parseDenseI64ArrayAttr(OpAsmParser &parser,
+                                          NamedAttrList &attributes,
+                                          StringRef attributeName) {
+  if (parser.parseKeyword(attributeName) || parser.parseEqual())
+    return failure();
+
+  attributes.set(attributeName, DenseI64ArrayAttr::parse(parser, Type{}));
+  return success();
+}
+
+ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) {
+  if (parseDstStyleOp(
+          parser, result, [&](OpAsmParser &parser, NamedAttrList &attributes) {
+            return parseDenseI64ArrayAttr(parser, attributes, "dimensions");
+          }))
+    return failure();
+
+  SmallVector<OpAsmParser::Argument> regionArgs;
+  if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
+                               /*allowType=*/true, /*allowAttrs=*/true)) {
+    return failure();
+  }
+
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, regionArgs))
+    return failure();
+
+  return success();
+}
+
+static void printDenseI64ArrayAttr(OpAsmPrinter &p, StringRef attributeName,
+                                   ArrayRef<int64_t> attributeValue) {
+  p << " " << attributeName << " = [" << attributeValue << "] ";
+}
+
+void ReduceOp::print(OpAsmPrinter &p) {
+  printCommonStructuredOpParts(p, getInputs(), getOutputs());
+  printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
+  p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
+
+  p << "(";
+  llvm::interleaveComma(getCombiner().getArguments(), p,
+                        [&](auto arg) { p.printRegionArgument(arg); });
+  p << ") ";
+
+  p.printRegion(getCombiner(), /*printEntryBlockArgs=*/false);
+}
+
+LogicalResult ReduceOp::verify() {
+  ArrayRef<int64_t> dimensionsRef = getDimensions();
+
+  for (int64_t i = 1; i < getNumInputs(); ++i) {
+    if (getInputs()[i].getType().cast<ShapedType>().getShape() !=
+        getInputs()[0].getType().cast<ShapedType>().getShape()) {
+      return emitOpError() << "expects all inputs to have the same shapes. "
+                              "Shape at input-index "
+                           << i
+                           << " is not equal to the shape at input-index 0.";
+    }
+  }
+  for (int64_t i = 1; i < getNumOutputs(); ++i) {
+    if (getInits()[i].getType().cast<ShapedType>().getShape() !=
+        getInits()[0].getType().cast<ShapedType>().getShape()) {
+      return emitOpError() << "expects all outputs to have the same shapes. "
+                              "Shape at output-index "
+                           << i
+                           << " is not equal to the shape at output-index 0.";
+    }
+  }
+  auto inputType = getInputs()[0].getType().cast<ShapedType>();
+  auto initType = getInits()[0].getType().cast<ShapedType>();
+
+  DenseSet<int64_t> dimensionsToReduce;
+  int64_t lastDimension = -1;
+  for (int64_t dimension : dimensionsRef) {
+    if (dimension < 0 || dimension >= inputType.getRank()) {
+      return emitOpError()
+             << "dimensions for reduction should be in the range [0, "
+             << inputType.getRank() - 1 << "].";
+    }
+    if (dimension <= lastDimension) {
+      return emitOpError()
+             << "reduction dimensions are not in increasing order: "
+             << dimensionsRef;
+    }
+
+    lastDimension = dimension;
+    dimensionsToReduce.insert(dimension);
+  }
+
+  auto inputDims = inputType.getShape();
+  auto initDims = initType.getShape();
+
+  // Input dimensions that will be left after the reduction.
+  SmallVector<int64_t> reducedInputDims;
+  for (const auto &en : llvm::enumerate(inputDims)) {
+    if (!dimensionsToReduce.count(en.index()))
+      reducedInputDims.push_back(en.value());
+  }
+
+  if (reducedInputDims.size() != initType.getRank()) {
+    return emitOpError() << "number of dimensions after reduction "
+                         << reducedInputDims.size()
+                         << " doesn't match the init rank "
+                         << initType.getRank();
+  }
+
+  if (reducedInputDims != initDims)
+    return emitOpError() << "init dimensions [" << initDims
+                         << "] doesn't match input dimensions after reduction ["
+                         << reducedInputDims << "]";
+
+  Block *block = getBody();
+  if (block->getNumArguments() != this->getNumOperands())
+    return emitOpError()
+           << "mismatching number of operands and block arguments";
+
+  // Check that the first block arguments match the element type of the inputs.
+  for (auto [input, bbArg] : llvm::zip(getInputs(), block->getArguments())) {
+    Type inputElementType = input.getType().cast<ShapedType>().getElementType();
+    if (inputElementType != bbArg.getType())
+      return emitOpError()
+             << "input element type " << inputElementType
+             << " does not match corresponding block argument type "
+             << bbArg.getType();
+  }
+
+  // Check that the last block arguments match the element type of the outputs.
+  for (auto [output, bbArg] : llvm::zip(
+           getOutputs(), block->getArguments().take_back(getNumOutputs()))) {
+    auto outputElementType =
+        output.getType().cast<ShapedType>().getElementType();
+    if (outputElementType != bbArg.getType())
+      return emitOpError()
+             << "output element type " << outputElementType
+             << " does not match corresponding block argument type "
+             << bbArg.getType();
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // InitTensorOp
 //===----------------------------------------------------------------------===//
 
 void InitTensorOp::build(OpBuilder &b, OperationState &result,
                          ArrayRef<OpFoldResult> sizes, Type elementType,
                          ArrayRef<NamedAttribute> attrs) {
   SmallVector<Value, 4> dynamicSizes;
   SmallVector<int64_t, 4> staticSizes;
   dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
                              ShapedType::kDynamicSize);
   auto resultType = RankedTensorType ::get(staticSizes, elementType);
   build(b, result, resultType, dynamicSizes, b.getI64ArrayAttr(staticSizes));
   result.addAttributes(attrs);
 }
 
 LogicalResult InitTensorOp::verify() {
   RankedTensorType resultType = getType();

mlir/test/Dialect/Linalg/invalid.mlir

 func.func @invalid_reverse(%A: memref<5xf32>, %B: memref<5xf32>) {
   // expected-error @+1 {{unexpected result less than 0 at expression #0 in}}
   linalg.generic #attrs ins(%A: memref<5xf32>) outs(%B: memref<5xf32>) {
                 ^bb0(%a: f32, %b: f32):
                 linalg.yield %a : f32
         }
         return
 }
+
+// -----
+
+func.func @reduce_input_vs_init_dimension_mismatch(
+    %input: tensor<16x32x64xf32>,
+    %init: tensor<16x64xf32>)  -> tensor<16x64xf32> {
+  // expected-error @+1 {{'linalg.reduce' op init dimensions [16, 64] doesn't match input dimensions after reduction [16, 32]}}
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16x64xf32>)
+      dimensions = [2]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16x64xf32>
+}
+
+// -----
+
+func.func @reduce_dimensions_out_of_range(%input: tensor<16x32x64xf32>,
+    %init: tensor<16x64xf32>)  -> tensor<16x64xf32> {
+  // expected-error @+1 {{'linalg.reduce' op dimensions for reduction should be in the range [0, 2].}}
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16x64xf32>)
+      dimensions = [3]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16x64xf32>
+}
+
+// -----
+
+func.func @reduce_duplicate_dimensions(%input: tensor<16x32x64xf32>,
+    %init: tensor<16xf32>)  -> tensor<16xf32> {
+  // expected-error @+1 {{'linalg.reduce' op reduction dimensions are not in increasing order: 1, 1}}
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16xf32>)
+      dimensions = [1, 1]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16xf32>
+}
+
+// -----
+
+func.func @reduce_non_increasing_dimensions(%input: tensor<16x32x64xf32>,
+    %init: tensor<16xf32>)  -> tensor<16xf32> {
+  // expected-error @+1 {{'linalg.reduce' op reduction dimensions are not in increasing order: 2, 1}}
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16xf32>)
+      dimensions = [2, 1]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16xf32>
+}
+
+// -----
+
+func.func @reduce_reduced_input_init_rank_mismatch(%input: tensor<16x32x64xf32>,
+    %init: tensor<16x64xf32>)  -> tensor<16x64xf32> {
+  // expected-error @+1 {{'linalg.reduce' op number of dimensions after reduction 1 doesn't match the init rank 2}}
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16x64xf32>)
+      dimensions = [1, 2]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16x64xf32>
+}
+
+// -----
+
+func.func @reduce_wrong_number_of_block_arguments(
+    %input1: tensor<16x32x64xf32>,
+    %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
+    %init2: tensor<16x64xf32>)  -> (tensor<16x64xf32>, tensor<16x64xf32>) {
+  // expected-error @+1{{'linalg.reduce' op mismatching number of operands and block arguments}}
+  %reduce, %reduce2 = linalg.reduce
+      ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
+      outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf32>)
+      dimensions = [1]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf32>
+}
+
+// -----
+
+func.func @reduce_wrong_block_argument_input_type(
+    %input1: tensor<16x32x64xf32>,
+    %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
+    %init2: tensor<16x64xf32>)  -> (tensor<16x64xf32>, tensor<16x64xf32>) {
+  // expected-error @+1{{'linalg.reduce' op input element type 'f32' does not match corresponding block argument type 'f64'}}
+  %reduce, %reduce2 = linalg.reduce
+      ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
+      outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf32>)
+      dimensions = [1]
+      (%in1: f32, %in2: f64, %out1: f32, %out2: f64) {
+        %0 = arith.addf %in1, %out1: f32
+        %1 = arith.addf %in2, %out2: f64
+        linalg.yield %0, %1: f32, f64
+      }
+  func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf32>
+}
+
+// -----
+
+func.func @reduce_wrong_block_argument_output_type(
+    %input1: tensor<16x32x64xf32>,
+    %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
+    %init2: tensor<16x64xf64>)  -> (tensor<16x64xf32>, tensor<16x64xf32>) {
+  // expected-error @+1{{'linalg.reduce' op output element type 'f64' does not match corresponding block argument type 'f32'}}
+  %reduce, %reduce2 = linalg.reduce
+      ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
+      outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf64>)
+      dimensions = [1]
+      (%in1: f32, %in2: f32, %out1: f32, %out2: f32) {
+        %0 = arith.addf %in1, %out1: f32
+        linalg.yield %0, %out2: f32, f32
+      }
+  func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf64>
+}
+
+// -----
+
+func.func @reduce_different_input_shapes(%input1: tensor<16x32x64xf32>,
+    %init1: tensor<16x64xf32>, %input2: tensor<17x32x64xf32>,
+    %init2: tensor<17x64xf32>)  -> (tensor<16x64xf32>, tensor<17x64xf32>) {
+  // expected-error @+1{{'linalg.reduce' op expects all inputs to have the same shapes. Shape at input-index 1 is not equal to the shape at input-index 0.}}
+  %reduce, %reduce2 = linalg.reduce
+      ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<17x32x64xf32>)
+      outs(%init1, %init2 : tensor<16x64xf32>, tensor<17x64xf32>)
+      dimensions = [1]
+      (%in1: f32, %in2: f32, %out1: f32, %out2: f32) {
+        %0 = arith.addf %in1, %out1: f32
+        %1 = arith.addf %in2, %out2: f32
+        linalg.yield %0, %1: f32, f32
+      }
+  func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<17x64xf32>
+}
+
+// -----
+
+func.func @reduce_different_output_shapes(%input1: tensor<16x32x64xf32>,
+    %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
+    %init2: tensor<17x64xf32>)  -> (tensor<16x64xf32>, tensor<17x64xf32>) {
+  // expected-error @+1{{'linalg.reduce' op expects all outputs to have the same shapes. Shape at output-index 1 is not equal to the shape at output-index 0.}}
+  %reduce, %reduce2 = linalg.reduce
+      ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
+      outs(%init1, %init2 : tensor<16x64xf32>, tensor<17x64xf32>)
+      dimensions = [1]
+      (%in1: f32, %in2: f32, %out1: f32, %out2: f32) {
+        %0 = arith.addf %in1, %out1: f32
+        %1 = arith.addf %in2, %out2: f32
+        linalg.yield %0, %1: f32, f32
+      }
+  func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<17x64xf32>
+}

mlir/test/Dialect/Linalg/roundtrip.mlir

^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32):
       %1 = arith.mulf %b0, %b1 : f32
       %2 = arith.addf %1, %b3 : f32
       linalg.yield %1, %2 : f32, f32
   } -> (tensor<?x?x?xf32>, tensor<?x?xf32>)
   return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?xf32>
 }
 // CHECK-LABEL: func @mixed_parallel_reduced_results
 //       CHECK:     linalg.generic
+
+// -----
+
+func.func @reduce(%input: tensor<16x32x64xf32>,
+                     %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16x64xf32>)
+      dimensions = [1]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16x64xf32>
+}
+// CHECK-LABEL: func @reduce
+//       CHECK:     linalg.reduce
+
+// -----
+
+func.func @variadic_reduce(%input1: tensor<16x32x64xf32>,
+    %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xi64>,
+    %init2: tensor<16x64xi64>)  -> (tensor<16x64xf32>, tensor<16x64xi64>) {
+  %reduce, %reduce2 = linalg.reduce
+      ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xi64>)
+      outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xi64>)
+      dimensions = [1]
+      (%in1: f32, %in2: i64, %out1: f32, %out2: i64) {
+        %0 = arith.addf %in1, %out1: f32
+        %1 = arith.addi %in2, %out2: i64
+        linalg.yield %0, %1: f32, i64
+      }
+  func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xi64>
+}
+// CHECK-LABEL: func @variadic_reduce
+//       CHECK:     linalg.reduce