[mlir][Linalg] Add methods to charecterize the loops of a ConvolutionOpInterface.

[mravishankar]

Summary

The loops in convolution/pooling ops (as defined/used today in Linalg)
can have the loops charecterized as

• batch dimensions,
• loop over output image dimensions,
• loop over output channel dimensions,
• loop over convolved filter dimensions,
• loop over input channel dimension,
• loop over depth multiplier.

Add interface methods that given an op that implements the
ConvolutionOpInterface that returns the loop that fall into one of
these categories.

Depends On D109793

[nicolasvasilache]

I still fail to see why these methods are desirable?
I'd venture that if some compiler pass (e.g. vectorization) uses this information it is thinking about the wrong abstractions.

Another way to phrase this is that in the fullness of time, the generic vectorization + enough canonicalizations and foldings does what we want.
Wrting code based on this information seems wrong to me.

[nicolasvasilache]

vectorization works fine without any of this as I was hoping, declaring bankruptcy here, thanks!

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 loop dimension that iterates over the batch dimension for `op` that
+/// implements the `ConvolutionOpInterface`. Returns llvm::None if the `op` does
+/// not implement the `ConvolutionOpInterface`, or if no batch dimensions exist.
+Optional<unsigned> getConvolutionBatchLoopDim(Operation *op);
+
+/// Get the loop dimensions that iterate over the output image for `op` that
+/// implements the `ConvolutionOpInterface`. Returns `{}` if the `op` does not
+/// implement the `ConvolutionOpInterface`.
+SmallVector<unsigned> getConvolutionOutputImageLoopDims(Operation *op);
+
+/// Get the loop dimension that iterates over the output channel dimensions for
+/// `op` that implements the `ConvolutionOpInterface`.  Returns llvm::None if
+/// the `op` does not implement the `ConvolutionOpInterface`, or if no output
+/// channel dimensions exist.
+Optional<unsigned> getConvolutionOutputChannelLoopDim(Operation *op);
+
+/// Get the loop dimensions that iterate over the filter loops for `op` that
+/// implements the `ConvolutionOpInterface`. Returns `{}` if the `op` does not
+/// implement the `ConvolutionOpInterface`.
+SmallVector<unsigned> getConvolutionFilterLoopDims(Operation *op);
+
+/// Get the loop dimension that iterates over the input channel dimensions for
+/// `op` that implements the `ConvolutionOpInterface`.  Returns llvm::None if
+/// the `op` does not implement the `ConvolutionOpInterface`, or if no input
+/// channel dimensions exist.
+Optional<unsigned> getConvolutionInputChannelLoopDim(Operation *op);
+
+/// Get the loop dimension that iterates over the depthwise dimension for `op`
+/// that implements the `ConvolutionOpInterface`.  Returns llvm::None if the
+/// `op` does not implement the `ConvolutionOpInterface`, or is not a depthwise
+/// convolution.
+Optional<unsigned> getConvolutionDepthwiseLoopDim(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

InterfaceMethod<
       /*retTy=*/"Value",
       /*methodName=*/"filter",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.getOperation()->getOperand(1);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Return the loop over batch dimensions of the convolution operation.",
+      /*retTy=*/"Optional<unsigned>",
+      /*methodName=*/"getBatchLoopDim",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionBatchLoopDim($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the loops over output image dimensions of the convolution operation.",
+      /*retTy=*/"SmallVector<unsigned>",
+      /*methodName=*/"getOutputImageLoopDims",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionOutputImageLoopDims($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the loop over output channel dimensions of the convolution operation.",
+      /*retTy=*/"Optional<unsigned>",
+      /*methodName=*/"getOutputChannelLoopDim",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionOutputChannelLoopDim($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the loop over filter dimensions of the convolution operation.",
+      /*retTy=*/"SmallVector<unsigned>",
+      /*methodName=*/"getFilterLoopDims",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionFilterLoopDims($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the loop over input channel dimensions of the convolution operation.",
+      /*retTy=*/"Optional<unsigned>",
+      /*methodName=*/"getInputChannelLoopDim",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionInputChannelLoopDim($_op);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the depthwise loop of the convolution operation.",
+      /*retTy=*/"Optional<unsigned>",
+      /*methodName=*/"getDepthwiseLoopDim",
+      /*args=*/(ins),
+      /*methodBody=*/[{
+        return mlir::getConvolutionDepthwiseLoopDim($_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();
 }
+
+Optional<unsigned> mlir::getConvolutionBatchLoopDim(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 batch dimensions are part of unconvoled dimensions in the input, and
+  // not present in the filter.
+  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;
+}
+
+SmallVector<unsigned> mlir::getConvolutionOutputImageLoopDims(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> outputImageLoopDims;
+  for (AffineExpr expr : indexingMaps.back().getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    if (inputExprWalker.convolvedDims.count(dim))
+      outputImageLoopDims.push_back(dim);
+  }
+  return outputImageLoopDims;
+}
+
+SmallVector<unsigned> mlir::getConvolutionFilterLoopDims(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> filterLoopDims;
+  for (AffineExpr expr : indexingMaps[1].getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    if (inputExprWalker.convolvedDims.count(dim))
+      filterLoopDims.push_back(dim);
+  }
+  return filterLoopDims;
+}
+
+Optional<unsigned> mlir::getConvolutionOutputChannelLoopDim(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 channel dimensions are not part of convolved or unconvolved
+  // dimensions in the input, and present in the filter.
+  llvm::SmallDenseSet<unsigned> filterDims = getPreservedDims(indexingMaps[1]);
+  for (AffineExpr expr : indexingMaps.back().getResults()) {
+    unsigned dim = expr.cast<AffineDimExpr>().getPosition();
+    if (!inputExprWalker.convolvedDims.count(dim) &&
+        !inputExprWalker.unConvolvedDims.count(dim) && filterDims.count(dim))
+      return dim;
+  }
+  return llvm::None;
+}
+
+Optional<unsigned> mlir::getConvolutionInputChannelLoopDim(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::getConvolutionDepthwiseLoopDim(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),