Linalg interfaces

[nicolasvasilache]

@springerm

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

 namespace detail {
 /// Implementation of the method that that check if given operands
 /// can be dropped, i.e. the remaining operands can compute the loop
 /// bounds of the op.
 bool canOpOperandsBeDroppedImpl(linalg::LinalgOp linalgOp,
                                 ArrayRef<OpOperand *> droppedOperands);
 } // namespace detail
 
-/// Checks whether `linalgOp` conforms to ContractionOpInterface.
-// TODO: embed within `isa<ContractionOpInterface>` if possible / natural.
-bool isaContractionOpInterface(LinalgOp linalgOp);
+//===----------------------------------------------------------------------===//
+// Interfaces for common op classes.
+//===----------------------------------------------------------------------===//
+// TODO: embed within `isa<XXXInterface>` if possible / natural.
+
+/// Return true if `op` conforms to linalg::ContractionOpInterface.
+bool isaContractionOpInterface(Operation *op);
+
+/// Return true if `op` conforms to linalg::ConvolutionOpInterface.
+bool isaConvolutionOpInterface(Operation *op);
+
+/// Return true if `op` conforms to linalg::FillOpInterface.
+bool isaFillOpInterface(Operation *op);
+
+/// Return true if `op` conforms to linalg::ReduceOpInterface.
+bool isaReduceOpInterface(Operation *op);
+
+/// Return true if `op` conforms to one of the known interfaces above.
+inline bool isKnownInterface(Operation *op) {
+  return isaContractionOpInterface(op) || isaConvolutionOpInterface(op) ||
+         isaFillOpInterface(op) || isaReduceOpInterface(op);
+}
 
 namespace detail {
 
 /// Verify that `op` conforms to ContractionOpInterface.
 LogicalResult verifyContractionInterface(Operation *op);
 
 /// Verify that `op` conforms to the ConvolutionOpInterface.
 LogicalResult verifyConvolutionInterface(Operation *op);
 
 /// Verify that `op` conforms to the FillOpInterface.
 LogicalResult verifyFillInterface(Operation *op);
 
+/// Verify that `op` conforms to the ReduceOpInterface.
+LogicalResult verifyReduceInterface(Operation *op);
+
 /// Verify that `op` conforms to the invariants of StructuredOpInterface
 LogicalResult verifyStructuredOpInterface(Operation *op);
 
 } // namespace detail
 } // namespace linalg
 } // namespace mlir
 
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.h.inc"
 
 /// Include the generated interface declarations.
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h.inc"
 
 #endif // MLIR_DIALECT_LINALG_IR_LINALGINTERFACES_H_

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

InterfaceMethod<
         if ($_op.getOperation()->getResults().empty())
           return nullptr;
         return $_op.getOperation()->getResults().front();
       }]
     >,
   ];
 }
 
+def LinalgReduceOpInterface : OpInterface<"ReduceOpInterface"> {
+  let description = [{
+    A reduce operation is defined in general terms:
+      1. Has a single input and init operand. In the future, could be relaxed.
+      2. Has only parallel and reduction loops.
+      3. Has only projected permutation indexing maps. In the future, could also
+         support mixed broadcasts.
+      4. Has a single matched combinator op in its region.
+  }];
+  let cppNamespace = "::mlir::linalg";
+  let verify = [{ return detail::verifyReduceInterface($_op); }];
+}
+
 // The 'LinalgStructuredInterface' provides access to the 'LinalgOp' interface.
 def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
   let cppNamespace = "::mlir::linalg";
   let methods = [
     //===------------------------------------------------------------------===//
     // Loop types handling.
     //===------------------------------------------------------------------===//
     InterfaceMethod<

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

 
 //===----------------------------------------------------------------------===//
 // Reduce op.
 //===----------------------------------------------------------------------===//
 
 def ReduceOp : LinalgStructuredBase_Op<"reduce", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>,
+    LinalgReduceOpInterface,
     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.
 
let arguments = (ins
     // Output arg
     Variadic<TensorOrMemref>:$inits,
 
     ConfinedAttr<DenseI64ArrayAttr,
                  [DenseArrayStrictlySorted<DenseI64ArrayAttr>]>:$dimensions
   );
   let results = (outs Variadic<AnyTensor>);
   let regions = (region SizedRegion<1>:$combiner);
+  DerivedAttr input_rank = DerivedAttr<"int",
+    "return getInputs()[0].getType().cast<ShapedType>().getRank();",
+    "$_builder.getI64IntegerAttr($_self)">;
 
   let builders = [
     OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$inits,
       "ArrayRef<int64_t>":$dimensions,
       "function_ref<void(OpBuilder &, Location, ValueRange)>",
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
   ];
 

mlir/lib/Dialect/Linalg/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgDialect
   MLIRLinalgInterfacesIncGen
   MLIRLinalgOpsAttributesIncGen
   MLIRLinalgOpsEnumsIncGen
   MLIRLinalgOpsIncGen
   MLIRLinalgStructuredOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
+  MLIRAnalysis
   MLIRArithDialect
   MLIRArithUtils
   MLIRBufferizationDialect
   MLIRDestinationStyleOpInterface
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRParser

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

 //===- LinalgInterfaces.cpp - Linalg interfaces implementation ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 
+#include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
if (!isAddMul<arith::AddFOp, arith::MulFOp>(linalgOp->getRegion(0).front()) &&
       !isAddMul<arith::AddIOp, arith::MulIOp>(linalgOp->getRegion(0).front()) &&
       !isAddMul<complex::AddOp, complex::MulOp>(
           linalgOp->getRegion(0).front()) &&
       !isAddMul<arith::OrIOp, arith::AndIOp>(linalgOp->getRegion(0).front()))
     return MatchContractionResult::NotAddMul;
   return MatchContractionResult::Success;
 }
 
-bool mlir::linalg::isaContractionOpInterface(LinalgOp linalgOp) {
-  if (!linalgOp)
-    return false;
-  Operation *op = linalgOp.getOperation();
+bool mlir::linalg::isaContractionOpInterface(Operation *op) {
   return isa<ContractionOpInterface>(op) ||
          (isContractionInterfaceImpl(op) == MatchContractionResult::Success);
 }
 
 /// Verify that a LinalgOp `op` is a contraction.
 /// A Linalg contraction is defined in general terms:
 ///   1. Has 2 input and 1 output shapes.
 ///   2. Has at least one reduction dimension.
static MatchConvolutionResult isConvolutionInterfaceImpl(Operation *op) {
   }
   // All loops must be covered now.
   if (allLoopDims.size() != linalgOp.getNumLoops())
     return MatchConvolutionResult::NonConvolutionLoop;
 
   return MatchConvolutionResult::Success;
 }
 
+bool mlir::linalg::isaConvolutionOpInterface(Operation *op) {
+  return isa<ConvolutionOpInterface>(op) ||
+         (isConvolutionInterfaceImpl(op) == MatchConvolutionResult::Success);
+}
+
 LogicalResult mlir::linalg::detail::verifyConvolutionInterface(Operation *op) {
   auto res = isConvolutionInterfaceImpl(op);
   if (res == MatchConvolutionResult::NotLinalgOp)
     return op->emitError("expected a LinalgOp");
   if (res == MatchConvolutionResult::WrongNumOperands)
     return op->emitError("expected op with 2 inputs and 1 output");
   if (res == MatchConvolutionResult::WrongInputIndexingMap)
     return op->emitError("unexpected input index map for convolutions");
static MatchFillResult isFillInterfaceImpl(Operation *op) {
 
   OpOperand *value = linalgOp.getDpsInputOperand(0);
   if (!linalgOp.isScalar(value))
     return MatchFillResult::NotScalarInput;
 
   return MatchFillResult::Success;
 }
 
+bool mlir::linalg::isaFillOpInterface(Operation *op) {
+  return isa<FillOpInterface>(op) ||
+         (isFillInterfaceImpl(op) == MatchFillResult::Success);
+}
+
 LogicalResult mlir::linalg::detail::verifyFillInterface(Operation *op) {
   auto res = isFillInterfaceImpl(op);
   if (res == MatchFillResult::NotLinalgOp)
     return op->emitError("expected a LinalgOp");
   if (res == MatchFillResult::WrongNumOperands)
     return op->emitError("expected op with 1 input and 1 output");
   if (res == MatchFillResult::NotScalarInput)
     return op->emitError("expected op with scalar input");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// ReduceOpInterface implementation
+//===----------------------------------------------------------------------===//
+
+enum class MatchReduceResult {
+  Success = 0,
+  NotLinalgOp,
+  WrongIteratorType,
+  WrongInputIndexingMap,
+  NotSingleCombiner
+};
+
+static MatchReduceResult isReduceInterfaceImpl(Operation *op) {
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+  if (!linalgOp)
+    return MatchReduceResult::NotLinalgOp;
+
+  SmallVector<unsigned> reductionDims;
+  linalgOp.getReductionDims(reductionDims);
+  if (linalgOp.getNumLoops() !=
+      linalgOp.getNumReductionLoops() + linalgOp.getNumParallelLoops())
+    return MatchReduceResult::WrongIteratorType;
+
+  // Only support projected permutation for now.
+  // TODO: could add support for mixed broadcasts too.
+  if (llvm::any_of(linalgOp.getDpsInputOperands(), [&](OpOperand *input) {
+        return !linalgOp.getMatchingIndexingMap(input).isProjectedPermutation();
+      }))
+    return MatchReduceResult::WrongInputIndexingMap;
+
+  // Only support single combiner operations for now.
+  SmallVector<Operation *, 4> combinerOps;
+  if (!matchReduction(linalgOp.getRegionOutputArgs(), 0, combinerOps) ||
+      combinerOps.size() != 1)
+    return MatchReduceResult::NotSingleCombiner;
+
+  return MatchReduceResult::Success;
+}
+
+bool mlir::linalg::isaReduceOpInterface(Operation *op) {
+  return isa<ReduceOpInterface>(op) ||
+         (isReduceInterfaceImpl(op) == MatchReduceResult::Success);
+}
+
+LogicalResult mlir::linalg::detail::verifyReduceInterface(Operation *op) {
+  auto res = isReduceInterfaceImpl(op);
+  if (res == MatchReduceResult::NotLinalgOp)
+    return op->emitError("expected a LinalgOp");
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // StructuredOpInterface implementation
 //===----------------------------------------------------------------------===//
 
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
                                int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())