[mlir][linalg] Add ConvolutionOpInterface methods for iterators

[antiagainst]

Summary

This commit adds a few methods to ConvolutionOpInterface to query
the iterator indices for different flavors of dimensions. These
methods are used to implement convolution vectorization patterns
based on ConvolutionOpInterface.

This is based on Mahesh's https://reviews.llvm.org/D110031, with
modifications to use parallel/window/reduction instead of ML/image
specific terminologies.

Co-authored-by: Mahesh Ravishankar <[e-mail]>

[nicolasvasilache]

I still view these as undesirable ML leaking abstractions into MLIR.
Can you submit a PR with both this and the Conv vectorization in a single PR ?
I'll rebase on top and experiment a bit.

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

 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
+
+/// Get the iterator indices that index into the output window dimensions for
+/// `op` that implements the `ConvolutionOpInterface`. Returns `{}` if the `op`
+/// does not implement the `ConvolutionOpInterface`.
+SmallVector<unsigned> getConvolutionOutputWindowIteratorIndices(Operation *op);
+
+/// Get the iterator indices that index into the filter window dimensions for
+/// `op` that implements the `ConvolutionOpInterface`. Returns `{}` if the `op`
+/// does not implement the `ConvolutionOpInterface`.
+SmallVector<unsigned> getConvolutionFilterWindowIteratorIndices(Operation *op);
+
+/// Get the iterator index that indexes the input reduction dimension for
+/// `op` that implements the `ConvolutionOpInterface`.  Returns llvm::None if
+/// the `op` does not implement the `ConvolutionOpInterface`, or if no input
+/// reduction dimensions exist.
+Optional<unsigned> getConvolutionInputReductionIteratorIndex(Operation *op);
+
+/// Get the iterator index that indexes the filter parallel dimension for `op`
+/// that implements the `ConvolutionOpInterface`.  Returns llvm::None if the
+/// `op` does not implement the `ConvolutionOpInterface`, or does not have a
+/// standalone parallel dimension.
+Optional<unsigned> getConvolutionFilterParallelIteratorIndex(Operation *op);
+
 namespace linalg {
 class LinalgOp;
 
 /// OpOperand vector that implicitly converts to a Value vector.
 struct OpOperandVector : public SmallVector<OpOperand *> {
   operator SmallVector<Value>();
 };
 

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

let description = [{
     3. The indexing maps of the input have expressions that satisfy
     ```
        AffineExpr ::== AffineDimExpr | ConvolvedExpr
        ConvolvedExpr ::== MulExpr (`+` MulExpr)+
        MulExpr ::== AffineDimExpr (`*` (AffineConstantExpr | AffineSymbolExpr))?
     ```
     4. The filter and the output have projected permutation maps.
     5. Each of the loops can be qualified as one of,
-       - Loop over batch dimension,
-       - Loop over output image dimensions,
-       - Loop over output channel dimensions,
-       - Loop over convolved filter dimensions,
-       - Loop over input channel dimension.
+       - Loop over output parallel dimensions (e.g., batch, output channel),
+       - Loop over output window dimensions (e.g., output height/width),
+       - Loop over input reduction dimension (e.g., input channel),
+       - Loop over convolved filter parallel dimensions (e.g., filter channel multiplier),
+       - Loop over convolved filter window dimensions (e.g., filter height/width).
   }];
   let cppNamespace = "::mlir::linalg";
   let verify = [{ return detail::verifyConvolutionInterface($_op); }];
   let methods = [
     InterfaceMethod<
       /*desc=*/"Return the image operand.",
       /*retTy=*/"Value",
       /*methodName=*/"image",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.getOperation()->getOperand(0);
       }]
     >,
     InterfaceMethod<
       /*desc=*/"Return the filter operand.",
       /*retTy=*/"Value",
       /*methodName=*/"filter",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.getOperation()->getOperand(1);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Return the output operand.",
+      /*retTy=*/"Value",
+      /*methodName=*/"output",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getOperation()->getOperand(2);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the strides.",
+      /*retTy=*/"DenseIntElementsAttr",
+      /*methodName=*/"strides",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        if (auto strides = $_op->template
+              getAttrOfType<DenseIntElementsAttr>("strides")) {
+          return strides;
+        }
+        return nullptr;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the dilations.",
+      /*retTy=*/"DenseIntElementsAttr",
+      /*methodName=*/"dilations",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        if (auto dilations = $_op->template
+              getAttrOfType<DenseIntElementsAttr>("dilations")) {
+          return dilations;
+        }
+        return nullptr;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the indices for the iterators that index into output window dimensions of the convolution operation.",
+      /*retTy=*/"SmallVector<unsigned>",
+      /*methodName=*/"getOutputWindowIteratorIndices",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionOutputWindowIteratorIndices($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the indices for the iterators that index into filter window dimensions of the convolution operation.",
+      /*retTy=*/"SmallVector<unsigned>",
+      /*methodName=*/"getFilterWindowIteratorIndices",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionFilterWindowIteratorIndices($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the index for the iterator that index input reduction dimensions of the convolution operation.",
+      /*retTy=*/"Optional<unsigned>",
+      /*methodName=*/"getInputRecutionIteratorIndex",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionInputReductionIteratorIndex($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the index for the iterator that index filter parallel dimensions of the convolution operation.",
+      /*retTy=*/"Optional<unsigned>",
+      /*methodName=*/"getFilterParallelIteratorIndex",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionFilterParallelIteratorIndex($_op);
+      }]
+    >,
   ];
 }
 
 // The 'LinalgStructuredInterface' provides access to the 'LinalgOp' interface.
 def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
   let cppNamespace = "::mlir::linalg";
   let methods = [
     //===------------------------------------------------------------------===//

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

return op->emitError(
         "expected all iterators used to access outputs to be parallel");
   }
   if (res == MatchConvolutionResult::NonOutputDimNotReduction) {
     return op->emitError(
         "expected all iterators not used to access outputs to be reduction");
   }
   return success();
 }
+
+SmallVector<unsigned>
+mlir::getConvolutionOutputWindowIteratorIndices(Operation *op) {
+  auto linalgOp = cast<linalg::LinalgOp>(op);
+  auto indexingMaps = linalgOp.getIndexingMaps();
+  ConvAccessExprWalker inputExprWalker;
+  if (llvm::any_of(indexingMaps[0].getResults(),
+                   [&inputExprWalker](AffineExpr expr) {
+                     return failed(inputExprWalker.visit(expr));
+                   })) {
+    llvm_unreachable("illegal input indexing map expression");
+  }
+  // The output image loops are dims that are convolved in the input, and
+  // present in the output.
+  SmallVector<unsigned> outputImageIteratorIndices;
+  for (AffineExpr expr : indexingMaps.back().getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    if (inputExprWalker.convolvedDims.count(dim))
+      outputImageIteratorIndices.push_back(dim);
+  }
+  return outputImageIteratorIndices;
+}
+
+SmallVector<unsigned>
+mlir::getConvolutionFilterWindowIteratorIndices(Operation *op) {
+  auto linalgOp = cast<linalg::LinalgOp>(op);
+  auto indexingMaps = linalgOp.getIndexingMaps();
+  ConvAccessExprWalker inputExprWalker;
+  if (llvm::any_of(indexingMaps[0].getResults(),
+                   [&inputExprWalker](AffineExpr expr) {
+                     return failed(inputExprWalker.visit(expr));
+                   })) {
+    llvm_unreachable("illegal input indexing map expression");
+  }
+  // The filter loops are dims that are convolved in the input, and present in
+  // the output.
+  SmallVector<unsigned> filterIteratorIndices;
+  for (AffineExpr expr : indexingMaps[1].getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    if (inputExprWalker.convolvedDims.count(dim))
+      filterIteratorIndices.push_back(dim);
+  }
+  return filterIteratorIndices;
+}
+
+Optional<unsigned>
+mlir::getConvolutionInputReductionIteratorIndex(Operation *op) {
+  auto linalgOp = cast<linalg::LinalgOp>(op);
+  auto indexingMaps = linalgOp.getIndexingMaps();
+  ConvAccessExprWalker inputExprWalker;
+  if (llvm::any_of(indexingMaps[0].getResults(),
+                   [&inputExprWalker](AffineExpr expr) {
+                     return failed(inputExprWalker.visit(expr));
+                   })) {
+    llvm_unreachable("illegal input indexing map expression");
+  }
+  // The input channel dimensions are part of unconvolved dimensions in the
+  // input, and not present in the output.
+  llvm::SmallDenseSet<unsigned> outputDims =
+      getPreservedDims(indexingMaps.back());
+  for (AffineExpr expr : indexingMaps[1].getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    if (inputExprWalker.unConvolvedDims.count(dim) && !outputDims.count(dim))
+      return dim;
+  }
+  return llvm::None;
+}
+
+Optional<unsigned>
+mlir::getConvolutionFilterParallelIteratorIndex(Operation *op) {
+  auto linalgOp = cast<linalg::LinalgOp>(op);
+  auto indexingMaps = linalgOp.getIndexingMaps();
+  ConvAccessExprWalker inputExprWalker;
+  if (llvm::any_of(indexingMaps[0].getResults(),
+                   [&inputExprWalker](AffineExpr expr) {
+                     return failed(inputExprWalker.visit(expr));
+                   })) {
+    llvm_unreachable("illegal input indexing map expression");
+  }
+  // The input channel dimensions are part of unconvolved dimensions in the
+  // input, and not present in the output.
+  llvm::SmallDenseSet<unsigned> filterDims = getPreservedDims(indexingMaps[1]);
+  for (AffineExpr expr : indexingMaps.back().getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    if (inputExprWalker.unConvolvedDims.count(dim) && filterDims.count(dim))
+      return dim;
+  }
+  return llvm::None;
+}
+
 //===----------------------------------------------------------------------===//
 // StructuredOpInterface implementation
 //===----------------------------------------------------------------------===//
 
 OpOperandVector::operator SmallVector<Value>() {
   SmallVector<Value> result;
   result.reserve(this->size());
   llvm::transform(*this, std::back_inserter(result),