[mlir] Create a generic reduction detection utility

[dcaballe]

Summary

This patch introduces a generic reduction detection utility that works
across different dialects. It is mostly a generalization of the reduction
detection algorithm in Affine. The reduction detection logic in Affine,
Linalg and SCFToOpenMP have been replaced with this new generic utility.

The utility takes some basic components of the potential reduction and
returns: 1) the reduced value, and 2) a list with the combiner operations.
The logic to match reductions involving multiple combiner operations disabled
until we can properly test it.

[bondhugula]

Nice to see this! It's a bit weird though that Conversion depends on Analysis - looks fine for now.

Commit summary:
dialecs <-- typo

mlir/test/lib/Analysis/TestMatchReduction.cpp

23–32

Can you use an emitRemark instead? This gives you a location to anchor to and in general better output. It wasn't a good pattern that some tests started of printing out text when the information could have been emitted using emitRemark and attached to an Operation. The memref dependence analysis stuff and loop parallel marking for reference uses emitRemark.

[bondhugula]

mlir/test/lib/Analysis/TestMatchReduction.cpp

67–68

Just an ArrayRef with a drop_front() should do?

[ftynse]

mlir/include/mlir/Analysis/LoopAnalysis.h

99–100

Should combiner operations also be side effect-free?

mlir/lib/Analysis/AffineAnalysis.cpp

43–46

Nit: please drop the explicit number of stack elements from SmallVector, it is now capable of selecting one by default. We only use explicit numbers if there is a specific reason to do so.

mlir/lib/Analysis/LoopAnalysis.cpp

402

I think you can get ancestorOp as iterCarriedArgs.front().getOwner().getParent().getParentOp. This looks a bit long, but avoids the case of iterCarriedArgs not being the arguments of the block being held by ancestorOp.

[nicolasvasilache]

Nice! Thanks @dcaballe

[sgrechanik]

mlir/include/mlir/Analysis/LoopAnalysis.h

115

In the future we might want to support reduction operations for which the reduced value does not exist explicitly (like a fused multiply-accumulate). I would recommend returning a bool to indicate whether reduction was detected, and making the reduced value an optional output parameter.

mlir/lib/Analysis/LoopAnalysis.cpp

467

Are there tests that will fail if this condition is removed?

[dcaballe]

• Addressed all the feedback.
• Extended 'operator<<' for diagnostic printing.
• Refactored code to a LinalgOp interface.

[dcaballe]

Nice to see this! It's a bit weird though that Conversion depends on Analysis - looks fine for now.

Yeah, I couldn't find a better place.

mlir/include/mlir/Analysis/LoopAnalysis.h

115

That's an interesting point! If we had a single fma op where only the addition part is the combiner operation in the reduction, I think some preprocessing should happen before using this utility. I don't think this utility should be responsible for, for example, splitting the fma. Returning the fma as a combiner operation would also be incorrect. In general, I think the mul+add->fma folding should be a peephole optimization that should happen later in the pipeline to avoid this kind of problems.

In any case, we can extend the utility when needed. The current version follows mostly the same API as the original utility in the Affine dialect. This is mostly a refactoring.

the combiner operation part if the combiner for the reduction? If that is the case, there

mlir/lib/Analysis/LoopAnalysis.cpp

402

I assumed ancestorOp could be a non-immediate parent but I can't think of a realistic use case for that. Ok, I'll simplify it but I'll move your suggestion to matchReduction. I'm adding a check to make sure all the combiner ops are immediately nested in this region

467

I thought it was covered by the OpenMP tests but it's not. I added a test. Thanks!

mlir/test/lib/Analysis/TestMatchReduction.cpp

23–32

Much better, thanks! I had to extend Diagnostics a bit, though .

[bondhugula]

mlir/lib/IR/Diagnostics.cpp

134

an -> a

mlir/test/lib/Analysis/TestMatchReduction.cpp

46

You can get rid of these as well? func->emitRemark.

[pifon2a]

Thank you, Diego!

[dcaballe]

Thank you all! I'll commit this in a while if no more comments.

mlir/include/mlir/Analysis/LoopAnalysis.h

 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
 
 namespace mlir {
 
 class AffineExpr;
 class AffineForOp;
 class AffineMap;
+class BlockArgument;
 class MemRefType;
 class NestedPattern;
 class Operation;
 class Value;
 
 /// Returns the trip count of the loop as an affine map with its corresponding
 /// operands if the latter is expressible as an affine expression, and nullptr
 /// otherwise. This method always succeeds as long as the lower bound is not a
bool isVectorizableLoopBody(AffineForOp loop, int *memRefDim,
                             NestedPattern &vectorTransferMatcher);
 
 /// Checks where SSA dominance would be violated if a for op's body
 /// operations are shifted by the specified shifts. This method checks if a
 /// 'def' and all its uses have the same shift factor.
 // TODO: extend this to check for memory-based dependence violation when we have
 // the support.
 bool isOpwiseShiftValid(AffineForOp forOp, ArrayRef<uint64_t> shifts);
+
+/// Utility to match a generic reduction given a list of iteration-carried
+/// arguments, `iterCarriedArgs` and the position of the potential reduction
+/// argument within the list, `redPos`. If a reduction is matched, returns the
+/// reduced value and the topologically-sorted list of combiner operations
+/// involved in the reduction. Otherwise, returns a null value.
+///
+/// The matching algorithm relies on the following invariants, which are subject
+/// to change:
+///  1. The first combiner operation must be a binary operation with the
+///     iteration-carried value and the reduced value as operands.
+///  2. The iteration-carried value and combiner operations must be side
+///     effect-free, have single result and a single use.
+///  3. Combiner operations must be immediately nested in the region op
+///     performing the reduction.
+///  4. Reduction def-use chain must end in a terminator op that yields the
+///     next iteration/output values in the same order as the iteration-carried
+///     values in `iterCarriedArgs`.
+///  5. `iterCarriedArgs` must contain all the iteration-carried/output values
+///     of the region op performing the reduction.
+///
+/// This utility is generic enough to detect reductions involving multiple
+/// combiner operations (disabled for now) across multiple dialects, including
+/// Linalg, Affine and SCF. For the sake of genericity, it does not return
+/// specific enum values for the combiner operations since its goal is also
+/// matching reductions without pre-defined semantics in core MLIR. It's up to
+/// each client to make sense out of the list of combiner operations. It's also
+/// up to each client to check for additional invariants on the expected
+/// reductions not covered by this generic matching.
+Value matchReduction(ArrayRef<BlockArgument> iterCarriedArgs, unsigned redPos,
+                     SmallVectorImpl<Operation *> &combinerOps);
 } // end namespace mlir
 
 #endif // MLIR_ANALYSIS_LOOP_ANALYSIS_H

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

InterfaceMethod<
         if (auto shapedType =
               opOperand->get().getType().template dyn_cast<ShapedType>())
           return shapedType.getRank();
         return 0;
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the output block arguments of the region.
+      }],
+      /*retTy=*/"Block::BlockArgListType",
+      /*methodName=*/"getRegionOutputArgs",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        Block &entryBlock = this->getOperation()->getRegion(0).front();
+        return entryBlock.getArguments().take_back(this->getNumOutputs());
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the `opOperand` shape or an empty vector for scalars.
       }],
       /*retTy=*/"ArrayRef<int64_t>",
       /*methodName=*/"getShape",
       /*args=*/(ins "OpOperand*":$opOperand),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         assert(opOperand->getOwner() == this->getOperation());

mlir/include/mlir/IR/Diagnostics.h

 namespace mlir {
 class DiagnosticEngine;
 class Identifier;
 struct LogicalResult;
 class MLIRContext;
 class Operation;
 class OperationName;
 class Type;
+class Value;
 
 namespace detail {
 struct DiagnosticEngineImpl;
 } // end namespace detail
 
 /// Defines the different supported severity of a diagnostic.
 enum class DiagnosticSeverity {
   Note,
public:
   Diagnostic &operator<<(OperationName val);
 
   /// Stream in an Operation.
   Diagnostic &operator<<(Operation &val);
   Diagnostic &operator<<(Operation *val) {
     return *this << *val;
   }
 
+  /// Stream in a Value.
+  Diagnostic &operator<<(Value val);
+
   /// Stream in a range.
   template <typename T, typename ValueT = llvm::detail::ValueOfRange<T>>
   std::enable_if_t<!std::is_constructible<DiagnosticArgument, T>::value,
                    Diagnostic &>
   operator<<(T &&range) {
     return appendRange(range);
   }
 

mlir/lib/Analysis/AffineAnalysis.cpp

 //===- AffineAnalysis.cpp - Affine structures analysis routines -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements miscellaneous analysis routines for affine structures
 // (expressions, maps, sets), and other utilities relying on such analysis.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineAnalysis.h"
+#include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "affine-analysis"
 
 using namespace mlir;
 
 using llvm::dbgs;
 
-/// Returns true if `value` (transitively) depends on iteration arguments of the
-/// given `forOp`.
-static bool dependsOnIterArgs(Value value, AffineForOp forOp) {
-  // Compute the backward slice of the value.
-  SetVector<Operation *> slice;
-  getBackwardSlice(value, &slice,
-                   [&](Operation *op) { return !forOp->isAncestor(op); });
-
-  // Check that none of the operands of the operations in the backward slice are
-  // loop iteration arguments, and neither is the value itself.
-  auto argRange = forOp.getRegionIterArgs();
-  llvm::SmallPtrSet<Value, 8> iterArgs(argRange.begin(), argRange.end());
-  if (iterArgs.contains(value))
-    return true;
-
-  for (Operation *op : slice)
-    for (Value operand : op->getOperands())
-      if (iterArgs.contains(operand))
-        return true;
-
-  return false;
-}
-
 /// Get the value that is being reduced by `pos`-th reduction in the loop if
 /// such a reduction can be performed by affine parallel loops. This assumes
 /// floating-point operations are commutative. On success, `kind` will be the
 /// reduction kind suitable for use in affine parallel loop builder. If the
 /// reduction is not supported, returns null.
 static Value getSupportedReduction(AffineForOp forOp, unsigned pos,
                                    AtomicRMWKind &kind) {
-  auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->back());
-  Value yielded = yieldOp.operands()[pos];
-  Operation *definition = yielded.getDefiningOp();
-  if (!definition)
-    return nullptr;
-  if (!forOp.getRegionIterArgs()[pos].hasOneUse())
+  SmallVector<Operation *> combinerOps;
+  Value reducedVal =
+      matchReduction(forOp.getRegionIterArgs(), pos, combinerOps);
+  if (!reducedVal)
     return nullptr;
-  if (!yielded.hasOneUse())
+
+  // Expected only one combiner operation.
+  if (combinerOps.size() > 1)
     return nullptr;
 
+  Operation *combinerOp = combinerOps.back();
   Optional<AtomicRMWKind> maybeKind =
-      TypeSwitch<Operation *, Optional<AtomicRMWKind>>(definition)
+      TypeSwitch<Operation *, Optional<AtomicRMWKind>>(combinerOp)
           .Case<AddFOp>([](Operation *) { return AtomicRMWKind::addf; })
           .Case<MulFOp>([](Operation *) { return AtomicRMWKind::mulf; })
           .Case<AddIOp>([](Operation *) { return AtomicRMWKind::addi; })
           .Case<MulIOp>([](Operation *) { return AtomicRMWKind::muli; })
           .Default([](Operation *) -> Optional<AtomicRMWKind> {
             // TODO: AtomicRMW supports other kinds of reductions this is
             // currently not detecting, add those when the need arises.
             return llvm::None;
           });
   if (!maybeKind)
     return nullptr;
 
   kind = *maybeKind;
-  if (definition->getOperand(0) == forOp.getRegionIterArgs()[pos] &&
-      !dependsOnIterArgs(definition->getOperand(1), forOp))
-    return definition->getOperand(1);
-  if (definition->getOperand(1) == forOp.getRegionIterArgs()[pos] &&
-      !dependsOnIterArgs(definition->getOperand(0), forOp))
-    return definition->getOperand(0);
-
-  return nullptr;
+  return reducedVal;
 }
 
 /// Returns true if `forOp' is a parallel loop. If `parallelReductions` is
 /// provided, populates it with descriptors of the parallelizable reductions and
 /// treats them as not preventing parallelization.
 bool mlir::isLoopParallel(AffineForOp forOp,
                           SmallVectorImpl<LoopReduction> *parallelReductions) {
   unsigned numIterArgs = forOp.getNumIterOperands();

mlir/lib/Analysis/LoopAnalysis.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/LoopAnalysis.h"
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/NestedMatcher.h"
+#include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Support/MathExtras.h"
 
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallString.h"
 #include <type_traits>
 
 using namespace mlir;
 
 /// Returns the trip count of the loop as an affine expression if the latter is
 /// expressible as an affine expression, and nullptr otherwise. The trip count
 /// expression is simplified before returning. This method only utilizes map
for (unsigned i = 0, e = op.getNumResults(); i < e; ++i) {
           if (shift != forBodyShift[ancOp])
             return false;
         }
       }
     }
   }
   return true;
 }
+
+/// Returns true if `value` (transitively) depends on iteration-carried values
+/// of the given `ancestorOp`.
+static bool dependsOnCarriedVals(Value value,
+                                 ArrayRef<BlockArgument> iterCarriedArgs,
+                                 Operation *ancestorOp) {
+  // Compute the backward slice of the value.
+  SetVector<Operation *> slice;
+  getBackwardSlice(value, &slice,
+                   [&](Operation *op) { return !ancestorOp->isAncestor(op); });
+
+  // Check that none of the operands of the operations in the backward slice are
+  // loop iteration arguments, and neither is the value itself.
+  SmallPtrSet<Value, 8> iterCarriedValSet(iterCarriedArgs.begin(),
+                                          iterCarriedArgs.end());
+  if (iterCarriedValSet.contains(value))
+    return true;
+
+  for (Operation *op : slice)
+    for (Value operand : op->getOperands())
+      if (iterCarriedValSet.contains(operand))
+        return true;
+
+  return false;
+}
+
+/// Utility to match a generic reduction given a list of iteration-carried
+/// arguments, `iterCarriedArgs` and the position of the potential reduction
+/// argument within the list, `redPos`. If a reduction is matched, returns the
+/// reduced value and the topologically-sorted list of combiner operations
+/// involved in the reduction. Otherwise, returns a null value.
+///
+/// The matching algorithm relies on the following invariants, which are subject
+/// to change:
+///  1. The first combiner operation must be a binary operation with the
+///     iteration-carried value and the reduced value as operands.
+///  2. The iteration-carried value and combiner operations must be side
+///     effect-free, have single result and a single use.
+///  3. Combiner operations must be immediately nested in the region op
+///     performing the reduction.
+///  4. Reduction def-use chain must end in a terminator op that yields the
+///     next iteration/output values in the same order as the iteration-carried
+///     values in `iterCarriedArgs`.
+///  5. `iterCarriedArgs` must contain all the iteration-carried/output values
+///     of the region op performing the reduction.
+///
+/// This utility is generic enough to detect reductions involving multiple
+/// combiner operations (disabled for now) across multiple dialects, including
+/// Linalg, Affine and SCF. For the sake of genericity, it does not return
+/// specific enum values for the combiner operations since its goal is also
+/// matching reductions without pre-defined semantics in core MLIR. It's up to
+/// each client to make sense out of the list of combiner operations. It's also
+/// up to each client to check for additional invariants on the expected
+/// reductions not covered by this generic matching.
+Value mlir::matchReduction(ArrayRef<BlockArgument> iterCarriedArgs,
+                           unsigned redPos,
+                           SmallVectorImpl<Operation *> &combinerOps) {
+  assert(redPos < iterCarriedArgs.size() && "'redPos' is out of bounds");
+
+  BlockArgument redCarriedVal = iterCarriedArgs[redPos];
+  if (!redCarriedVal.hasOneUse())
+    return nullptr;
+
+  // For now, the first combiner op must be a binary op.
+  Operation *combinerOp = *redCarriedVal.getUsers().begin();
+  if (combinerOp->getNumOperands() != 2)
+    return nullptr;
+  Value reducedVal = combinerOp->getOperand(0) == redCarriedVal
+                         ? combinerOp->getOperand(1)
+                         : combinerOp->getOperand(0);
+
+  Operation *redRegionOp =
+      iterCarriedArgs.front().getOwner()->getParent()->getParentOp();
+  if (dependsOnCarriedVals(reducedVal, iterCarriedArgs, redRegionOp))
+    return nullptr;
+
+  // Traverse the def-use chain starting from the first combiner op until a
+  // terminator is found. Gather all the combiner ops along the way in
+  // topological order.
+  while (!combinerOp->mightHaveTrait<OpTrait::IsTerminator>()) {
+    if (!MemoryEffectOpInterface::hasNoEffect(combinerOp) ||
+        combinerOp->getNumResults() != 1 || !combinerOp->hasOneUse() ||
+        combinerOp->getParentOp() != redRegionOp)
+      return nullptr;
+
+    combinerOps.push_back(combinerOp);
+    combinerOp = *combinerOp->getUsers().begin();
+  }
+
+  // Limit matching to single combiner op until we can properly test reductions
+  // involving multiple combiners.
+  if (combinerOps.size() != 1)
+    return nullptr;
+
+  // Check that the yielded value is in the same position as in
+  // `iterCarriedArgs`.
+  Operation *terminatorOp = combinerOp;
+  if (terminatorOp->getOperand(redPos) != combinerOps.back()->getResults()[0])
+    return nullptr;
+
+  return reducedVal;
+}

mlir/lib/Conversion/SCFToOpenMP/CMakeLists.txt

 add_mlir_conversion_library(MLIRSCFToOpenMP
   SCFToOpenMP.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/SCFToStandard
 
   DEPENDS
   MLIRConversionPassIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
+  MLIRAnalysis
   MLIRLLVMIR
   MLIROpenMP
   MLIRSCF
   MLIRStandard
   MLIRTransforms
   )

mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp

 //===- SCFToOpenMP.cpp - Structured Control Flow to OpenMP conversion -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to convert scf.parallel operations into OpenMP
 // parallel loops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
 #include "../PassDetail.h"
+#include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 /// Matches a block containing a "simple" reduction. The expected shape of the
 /// block is as follows.
 ///
 ///   ^bb(%arg0, %arg1):
 ///     %0 = OpTy(%arg0, %arg1)
 ///     scf.reduce.return %0
 template <typename... OpTy>
 static bool matchSimpleReduction(Block &block) {
   if (block.empty() || llvm::hasSingleElement(block) ||
       std::next(block.begin(), 2) != block.end())
     return false;
-  return isa<OpTy...>(block.front()) &&
+
+  if (block.getNumArguments() != 2)
+    return false;
+
+  SmallVector<Operation *, 4> combinerOps;
+  Value reducedVal = matchReduction({block.getArguments()[1]},
+                                    /*redPos=*/0, combinerOps);
+
+  if (!reducedVal || !reducedVal.isa<BlockArgument>() ||
+      combinerOps.size() != 1)
+    return false;
+
+  return isa<OpTy...>(combinerOps[0]) &&
          isa<scf::ReduceReturnOp>(block.back()) &&
-         block.front().getOperands() == block.getArguments() &&
-         block.back().getOperand(0) == block.front().getResult(0);
+         block.front().getOperands() == block.getArguments();
 }
 
 /// Matches a block containing a select-based min/max reduction. The types of
 /// select and compare operations are provided as template arguments. The
 /// comparison predicates suitable for min and max are provided as function
 /// arguments. If a reduction is matched, `ifMin` will be set if the reduction
 /// compute the minimum and unset if it computes the maximum, otherwise it
 /// remains unmodified. The expected shape of the block is as follows.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

 //===- Vectorization.cpp - Implementation of linalg Vectorization ---------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the linalg dialect Vectorization transformations.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"

 /// ShapedType of `v`.
 static VectorType extractVectorTypeFromShapedValue(Value v) {
   auto st = v.getType().cast<ShapedType>();
   if (st.isa<MemRefType>() && st.getShape().empty())
     return VectorType();
   return VectorType::get(st.getShape(), st.getElementType());
 }
 
-/// Given an `outputOperand` of a LinalgOp, compute the intersection of the
-/// forward slice starting from `outputOperand` and the backward slice
-/// starting from the corresponding linalg.yield operand.
-/// This intersection is assumed to have a single binary operation that is
-/// the reduction operation. Multiple reduction operations would impose an
+/// Check whether `outputOperand` is a reduction with a single combiner
+/// operation. Return the combiner operation of the reduction, which is assumed
+/// to be a binary operation. Multiple reduction operations would impose an
 /// ordering between reduction dimensions and is currently unsupported in
-/// Linalg. This limitation is motivated by the fact that e.g.
-/// min(max(X)) != max(min(X))
+/// Linalg. This limitation is motivated by the fact that e.g. min(max(X)) !=
+/// max(min(X))
 // TODO: use in LinalgOp verification, there is a circular dependency atm.
 static Operation *getSingleBinaryOpAssumedReduction(OpOperand *outputOperand) {
   auto linalgOp = cast<LinalgOp>(outputOperand->getOwner());
-  auto yieldOp = cast<YieldOp>(linalgOp->getRegion(0).front().getTerminator());
-  unsigned yieldNum =
+  unsigned outputPos =
       outputOperand->getOperandNumber() - linalgOp.getNumInputs();
-  llvm::SetVector<Operation *> backwardSlice, forwardSlice;
-  BlockArgument bbArg = linalgOp->getRegion(0).front().getArgument(
-      outputOperand->getOperandNumber());
-  Value yieldVal = yieldOp->getOperand(yieldNum);
-  getBackwardSlice(yieldVal, &backwardSlice, [&](Operation *op) {
-    return op->getParentOp() == linalgOp;
-  });
-  backwardSlice.insert(yieldVal.getDefiningOp());
-  getForwardSlice(bbArg, &forwardSlice,
-                  [&](Operation *op) { return op->getParentOp() == linalgOp; });
-  // Search for the (assumed unique) elementwiseMappable op at the intersection
-  // of forward and backward slices.
-  Operation *reductionOp = nullptr;
-  for (Operation *op : llvm::reverse(backwardSlice)) {
-    if (!forwardSlice.contains(op))
-      continue;
-    if (OpTrait::hasElementwiseMappableTraits(op)) {
-      if (reductionOp) {
-        // Reduction detection fails: found more than 1 elementwise-mappable op.
-        return nullptr;
-      }
-      reductionOp = op;
-    }
-  }
+  SmallVector<Operation *, 4> combinerOps;
+  if (!matchReduction(linalgOp.getRegionOutputArgs(), outputPos, combinerOps) ||
+      combinerOps.size() != 1)
+
   // TODO: also assert no other subsequent ops break the reduction.
-  return reductionOp;
+  return combinerOps[0];
 }
 
 /// If `value` of assumed VectorType has a shape different than `shape`, try to
 /// build and return a new vector.broadcast to `shape`.
 /// Otherwise, just return `value`.
 // TODO: this is best effort atm and there is currently no guarantee of
 // correctness for the broadcast semantics.
 static Value broadcastIfNeeded(OpBuilder &b, Value value,

mlir/lib/IR/Diagnostics.cpp

 /// Stream in an Operation.
 Diagnostic &Diagnostic::operator<<(Operation &val) {
   std::string str;
   llvm::raw_string_ostream os(str);
   val.print(os, OpPrintingFlags().useLocalScope().elideLargeElementsAttrs());
   return *this << os.str();
 }
 
+/// Stream in a Value.
+Diagnostic &Diagnostic::operator<<(Value val) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  val.print(os);
+  return *this << os.str();
+}
+
 /// Outputs this diagnostic to a stream.
 void Diagnostic::print(raw_ostream &os) const {
   for (auto &arg : getArguments())
     arg.print(os);
 }
 
 /// Convert the diagnostic to a string.
 std::string Diagnostic::str() const {

mlir/test/Analysis/test-match-reduction.mlir

+// RUN: mlir-opt %s -test-match-reduction -verify-diagnostics -split-input-file
+
+// Verify that the generic reduction detection utility works on different
+// dialects.
+
+// expected-remark@below {{Testing function}}
+func @linalg_red_add(%in0t : tensor<?xf32>, %out0t : tensor<1xf32>) {
+  // expected-remark@below {{Reduction found in output #0!}}
+  // expected-remark@below {{Reduced Value: <block argument> of type 'f32' at index: 0}}
+  // expected-remark@below {{Combiner Op: %1 = addf %arg2, %arg3 : f32}}
+  %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
+                                          affine_map<(d0) -> (0)>],
+                                          iterator_types = ["reduction"]}
+   ins(%in0t : tensor<?xf32>)
+   outs(%out0t : tensor<1xf32>) {
+    ^bb0(%in0: f32, %out0: f32):
+      %add = addf %in0, %out0 : f32
+      linalg.yield %add : f32
+    } -> tensor<1xf32>
+  return
+}
+
+// -----
+
+// expected-remark@below {{Testing function}}
+func @affine_red_add(%in: memref<256x512xf32>, %out: memref<256xf32>) {
+ %cst = constant 0.000000e+00 : f32
+ affine.for %i = 0 to 256 {
+   // expected-remark@below {{Reduction found in output #0!}}
+   // expected-remark@below {{Reduced Value: %1 = affine.load %arg0[%arg2, %arg3] : memref<256x512xf32>}}
+   // expected-remark@below {{Combiner Op: %2 = addf %arg4, %1 : f32}}
+   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
+     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
+     %add = addf %red_iter, %ld : f32
+     affine.yield %add : f32
+   }
+   affine.store %final_red, %out[%i] : memref<256xf32>
+ }
+ return
+}
+
+// -----
+
+// TODO: Iteration-carried values with multiple uses are not supported yet.
+// expected-remark@below {{Testing function}}
+func @linalg_red_max(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) {
+  // expected-remark@below {{Reduction NOT found in output #0!}}
+  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                          affine_map<(d0, d1) -> (d0)>],
+   iterator_types = ["parallel", "reduction"]}
+   ins(%in0t : tensor<4x4xf32>)
+   outs(%out0t : tensor<4xf32>) {
+    ^bb0(%in0: f32, %out0: f32):
+      %cmp = cmpf ogt, %in0, %out0 : f32
+      %sel = select %cmp, %in0, %out0 : f32
+      linalg.yield %sel : f32
+    } -> tensor<4xf32>
+  return
+}
+
+// -----
+
+// expected-remark@below {{Testing function}}
+func @linalg_fused_red_add(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) {
+  // expected-remark@below {{Reduction found in output #0!}}
+  // expected-remark@below {{Reduced Value: %2 = subf %1, %arg2 : f32}}
+  // expected-remark@below {{Combiner Op: %3 = addf %2, %arg3 : f32}}
+  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                          affine_map<(d0, d1) -> (d0)>],
+   iterator_types = ["parallel", "reduction"]}
+   ins(%in0t : tensor<4x4xf32>)
+   outs(%out0t : tensor<4xf32>) {
+    ^bb0(%in0: f32, %out0: f32):
+      %mul = mulf %in0, %in0 : f32
+      %sub = subf %mul, %in0 : f32
+      %add = addf %sub, %out0 : f32
+      linalg.yield %add : f32
+    } -> tensor<4xf32>
+  return
+}
+
+// -----
+
+// expected-remark@below {{Testing function}}
+func @affine_no_red_rec(%in: memref<512xf32>) {
+ %cst = constant 0.000000e+00 : f32
+ // %rec is the value loaded in the previous iteration.
+ // expected-remark@below {{Reduction NOT found in output #0!}}
+ %final_val = affine.for %j = 0 to 512 iter_args(%rec = %cst) -> (f32) {
+   %ld = affine.load %in[%j] : memref<512xf32>
+   %add = addf %ld, %rec : f32
+   affine.yield %ld : f32
+ }
+ return
+}
+
+// -----
+
+// expected-remark@below {{Testing function}}
+func @affine_output_dep(%in: memref<512xf32>) {
+ %cst = constant 0.000000e+00 : f32
+ // Reduction %red is not supported because it depends on another
+ // loop-carried dependence.
+ // expected-remark@below {{Reduction NOT found in output #0!}}
+ // expected-remark@below {{Reduction NOT found in output #1!}}
+ %final_red, %final_dep = affine.for %j = 0 to 512
+  iter_args(%red = %cst, %dep = %cst) -> (f32, f32) {
+   %ld = affine.load %in[%j] : memref<512xf32>
+   %add = addf %dep, %red : f32
+   affine.yield %add, %ld : f32, f32
+ }
+ return
+}
+

mlir/test/lib/Analysis/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestAnalysis
   TestAliasAnalysis.cpp
   TestCallGraph.cpp
   TestLiveness.cpp
+  TestMatchReduction.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestNumberOfExecutions.cpp
 
 
   EXCLUDE_FROM_LIBMLIR
 

mlir/test/lib/Analysis/TestMatchReduction.cpp

+//===- TestMatchReduction.cpp - Test the match reduction utility ----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a test pass for the match reduction utility.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/LoopAnalysis.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+void printReductionResult(Operation *redRegionOp, unsigned numOutput,
+                          Value reducedValue,
+                          ArrayRef<Operation *> combinerOps) {
+  if (reducedValue) {
+    redRegionOp->emitRemark("Reduction found in output #") << numOutput << "!";
+    redRegionOp->emitRemark("Reduced Value: ") << reducedValue;
+    for (Operation *combOp : combinerOps)
+      redRegionOp->emitRemark("Combiner Op: ") << *combOp;
+
+    return;
+  }
+
+  redRegionOp->emitRemark("Reduction NOT found in output #")
+      << numOutput << "!";
+}
+
+struct TestMatchReductionPass
+    : public PassWrapper<TestMatchReductionPass, FunctionPass> {
+  StringRef getArgument() const final { return "test-match-reduction"; }
+  StringRef getDescription() const final {
+    return "Test the match reduction utility.";
+  }
+
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+    func->emitRemark("Testing function");
+
+    func.walk<WalkOrder::PreOrder>([](Operation *op) {
+      if (isa<FuncOp>(op))
+        return;
+
+      // Limit testing to ops with only one region.
+      if (op->getNumRegions() != 1)
+        return;
+
+      Region &region = op->getRegion(0);
+      if (!region.hasOneBlock())
+        return;
+
+      // We expect all the tested region ops to have 1 input by default. The
+      // remaining arguments are assumed to be outputs/reductions and there must
+      // be at least one.
+      // TODO: Extend it to support more generic cases.
+      Block &regionEntry = region.front();
+      auto args = regionEntry.getArguments();
+      if (args.size() < 2)
+        return;
+
+      auto outputs = args.drop_front();
+      for (int i = 0, size = outputs.size(); i < size; ++i) {
+        SmallVector<Operation *, 4> combinerOps;
+        Value reducedValue = matchReduction(outputs, i, combinerOps);
+        printReductionResult(op, i, reducedValue, combinerOps);
+      }
+    });
+  }
+};
+
+} // end anonymous namespace
+
+namespace mlir {
+namespace test {
+void registerTestMatchReductionPass() {
+  PassRegistration<TestMatchReductionPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
 void registerTestLinalgTileAndFuseSequencePass();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
+void registerTestMatchReductionPass();
 void registerTestMathAlgebraicSimplificationPass();
 void registerTestMathPolynomialApproximationPass();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestNumberOfBlockExecutionsPass();
 void registerTestNumberOfOperationExecutionsPass();
 void registerTestOpaqueLoc();
 void registerTestPDLByteCodePass();
#endif
   mlir::test::registerTestLinalgGreedyFusion();
   mlir::test::registerTestLinalgHoisting();
   mlir::test::registerTestLinalgTileAndFuseSequencePass();
   mlir::test::registerTestLinalgTransforms();
   mlir::test::registerTestLivenessPass();
   mlir::test::registerTestLoopFusion();
   mlir::test::registerTestLoopMappingPass();
   mlir::test::registerTestLoopUnrollingPass();
+  mlir::test::registerTestMatchReductionPass();
   mlir::test::registerTestMathAlgebraicSimplificationPass();
   mlir::test::registerTestMathPolynomialApproximationPass();
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestNumberOfBlockExecutionsPass();
   mlir::test::registerTestNumberOfOperationExecutionsPass();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPDLByteCodePass();

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     name = "SCFToOpenMP",
     srcs = [
         "lib/Conversion/PassDetail.h",
         "lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp",
     ],
     hdrs = ["include/mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"],
     includes = ["include"],
     deps = [
+        ":Analysis",
         ":ConversionPassIncGen",
         ":IR",
         ":LLVMDialect",
         ":OpenMPDialect",
         ":Pass",
         ":SCFDialect",
         ":StandardOps",
         ":Support",