[mlir][linalg] expose convolution dimension classifier

[ftynse]

Summary

Make available through functions in the linalg::detail namespace the
classification of Linalg op dimensions as different kinds (batch, image,
channel, etc) of convolution dimensions. This is useful for identifying
which dimensions to target with transformations.

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

 } // namespace detail
 
 /// Checks whether `linalgOp` conforms to ContractionOpInterface.
 // TODO: embed within `isa<ContractionOpInterface>` if possible / natural.
 bool isaContractionOpInterface(LinalgOp linalgOp);
 
 namespace detail {
 
+/// Result of matching a Linalg generic against the predicates of it being a
+/// convolution.
+enum class MatchConvolutionResult;
+
+/// Positions of a Linalg op loops that correspond to different kinds of a
+/// convolution dimension.
+struct ConvolutionDimensions {
+  SmallVector<unsigned, 2> batch;
+  SmallVector<unsigned, 2> outputImage;
+  SmallVector<unsigned, 2> outputChannel;
+  SmallVector<unsigned, 2> filterLoop;
+  SmallVector<unsigned, 2> inputChannel;
+  SmallVector<unsigned, 2> depth;
+};
+
+/// Checks whether `op` conforms to ConvolutionOpInterface and populates
+/// `dimensions` with indexes of the different kinds of dimensions when present.
+MatchConvolutionResult
+isConvolutionInterfaceImpl(Operation *op,
+                           ConvolutionDimensions *dimensions = nullptr);
+
+/// Returns the error message corresponding to the convolution checking return
+/// code.
+StringRef getMatchConvolutionMessage(MatchConvolutionResult res);
+
 /// 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);

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

static llvm::SmallDenseSet<unsigned> getPreservedDims(AffineMap map) {
   assert(map.isProjectedPermutation() &&
          "expected map to have projected permutations");
   llvm::SmallDenseSet<unsigned> preservedDims;
   for (auto expr : map.getResults())
     preservedDims.insert(expr.cast<AffineDimExpr>().getPosition());
   return preservedDims;
 }
 
+namespace mlir::linalg::detail {
 enum class MatchConvolutionResult {
   Success = 0,
   NotLinalgOp,
   WrongNumOperands,
   WrongInputIndexingMap,
   NotProjectedPermutations,
   NonConvolutionLoop,
   OutputDimsNotParallel,
   NonOutputDimNotReduction
 };
+} // namespace mlir::linalg::detail
 
-static MatchConvolutionResult isConvolutionInterfaceImpl(Operation *op) {
+mlir::linalg::detail::MatchConvolutionResult
+mlir::linalg::detail::isConvolutionInterfaceImpl(
+    Operation *op, ConvolutionDimensions *dimensions) {
   auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
   if (!linalgOp)
     return MatchConvolutionResult::NotLinalgOp;
   if (linalgOp.getNumDpsInputs() < 2 || linalgOp.getNumDpsInits() != 1)
     return MatchConvolutionResult::WrongNumOperands;
 
   auto indexingMaps = linalgOp.getIndexingMapsArray();
 
if (!indexingMaps[1].isProjectedPermutation() ||
       !indexingMaps.back().isProjectedPermutation())
     return MatchConvolutionResult::NotProjectedPermutations;
 
   auto iteratorTypes = linalgOp.getIteratorTypesArray();
 
   llvm::SmallDenseSet<unsigned> outputDims =
       getPreservedDims(indexingMaps.back());
   llvm::SmallDenseSet<unsigned> filterDims = getPreservedDims(indexingMaps[1]);
-  // Make sure all loops are charecterized as one of:
+  // Make sure all loops are characterized as one of:
   // - Batch loop : present in output, as non-convolved in input, not present in
   //   filter.
   // - Output image dimension : present in output, convolved dims in input, not
   //   present in filter.
   // - Output channel dimension : present in output, not present in input,
   //   present in filter.
   // - Filter loop dimension : present in filter, convolved in input, not
   //   present in output.
   // - Input channel dimension : unconvolved in input, not present in output,
   //   present in filter.
   // - Depth multiplier : unconvolved in input, present in output, present in
   //   filter.
   llvm::SmallDenseSet<unsigned> allLoopDims;
   for (auto outputExpr : indexingMaps.back().getResults()) {
     unsigned outputDim = outputExpr.cast<AffineDimExpr>().getPosition();
     if (inputExprWalker.unConvolvedDims.count(outputDim) &&
         !filterDims.count(outputDim)) {
       // Batch dimension.
       if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
+      if (dimensions)
+        dimensions->batch.push_back(outputDim);
       continue;
     }
     if (inputExprWalker.convolvedDims.count(outputDim) &&
         !filterDims.count(outputDim)) {
       // Output image Loop dimension.
       if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
+      if (dimensions)
+        dimensions->outputImage.push_back(outputDim);
       continue;
     }
     if (!inputExprWalker.convolvedDims.count(outputDim) &&
         !inputExprWalker.unConvolvedDims.count(outputDim) &&
         filterDims.count(outputDim)) {
       // Output channel dimension.
       if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
+      if (dimensions)
+        dimensions->outputChannel.push_back(outputDim);
       continue;
     }
     if (inputExprWalker.unConvolvedDims.count(outputDim) &&
         filterDims.count(outputDim)) {
       // Depth multiplier.
       if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
+      if (dimensions)
+        dimensions->depth.push_back(outputDim);
       continue;
     }
     return MatchConvolutionResult::NonConvolutionLoop;
   }
   for (auto filterExpr : indexingMaps[1].getResults()) {
     unsigned filterDim = filterExpr.cast<AffineDimExpr>().getPosition();
     if (outputDims.count(filterDim) &&
         !inputExprWalker.unConvolvedDims.count(filterDim) &&
         !inputExprWalker.convolvedDims.count(filterDim)) {
-      // Output channel dimension. THis is already seen, continue;
+      // Output channel dimension. This is already seen, continue;
+      assert((!dimensions ||
+              llvm::is_contained(dimensions->outputChannel, filterDim)) &&
+             "expected output channel to have been found from output dims");
       continue;
     }
     if (inputExprWalker.convolvedDims.count(filterDim) &&
         !outputDims.count(filterDim)) {
       // Filter loop dimension.
       if (iteratorTypes[filterDim] != utils::IteratorType::reduction)
         return MatchConvolutionResult::NonOutputDimNotReduction;
       if (allLoopDims.count(filterDim))
         return MatchConvolutionResult::NonConvolutionLoop;
       allLoopDims.insert(filterDim);
+      if (dimensions)
+        dimensions->filterLoop.push_back(filterDim);
       continue;
     }
     if (inputExprWalker.unConvolvedDims.count(filterDim) &&
         !outputDims.count(filterDim)) {
       // Input channel dimension.
       if (iteratorTypes[filterDim] != utils::IteratorType::reduction)
         return MatchConvolutionResult::NonOutputDimNotReduction;
       if (allLoopDims.count(filterDim))
         return MatchConvolutionResult::NonConvolutionLoop;
       allLoopDims.insert(filterDim);
+      if (dimensions)
+        dimensions->inputChannel.push_back(filterDim);
       continue;
     }
     if (inputExprWalker.unConvolvedDims.count(filterDim) &&
         outputDims.count(filterDim)) {
       // Depthwise loop. Already seen.
+      assert(
+          (!dimensions || llvm::is_contained(dimensions->depth, filterDim)) &&
+          "expected depthwise dimension to have been found from output dims");
       continue;
     }
     return MatchConvolutionResult::NonConvolutionLoop;
   }
   // All loops must be covered now.
   if (allLoopDims.size() != linalgOp.getNumLoops())
     return MatchConvolutionResult::NonConvolutionLoop;
 
+  if (dimensions) {
+    assert(dimensions->batch.size() + dimensions->outputImage.size() +
+                   dimensions->outputChannel.size() +
+                   dimensions->filterLoop.size() +
+                   dimensions->inputChannel.size() + dimensions->depth.size() ==
+               linalgOp.getNumLoops() &&
+           "expected all loops to be classified");
+  }
+
   return 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");
-  if (res == MatchConvolutionResult::NotProjectedPermutations) {
-    return op->emitError(
-        "expected output/filter indexing maps to be projected permutations");
-  }
-  if (res == MatchConvolutionResult::NonConvolutionLoop) {
-    return op->emitError("unexpected loop dimension for convolution op");
-  }
-  if (res == MatchConvolutionResult::OutputDimsNotParallel) {
-    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");
+StringRef
+mlir::linalg::detail::getMatchConvolutionMessage(MatchConvolutionResult res) {
+  switch (res) {
+  case MatchConvolutionResult::NotLinalgOp:
+    return "expected a LinalgOp";
+  case MatchConvolutionResult::WrongNumOperands:
+    return "expected op with 2 inputs and 1 output";
+  case MatchConvolutionResult::WrongInputIndexingMap:
+    return "unexpected input index map for convolutions";
+  case MatchConvolutionResult::NotProjectedPermutations:
+    return "expected output/filter indexing maps to be projected permutations";
+  case MatchConvolutionResult::NonConvolutionLoop:
+    return "unexpected loop dimension for convolution op";
+  case MatchConvolutionResult::OutputDimsNotParallel:
+    return "expected all iterators used to access outputs to be parallel";
+  case MatchConvolutionResult::NonOutputDimNotReduction:
+    return "expected all iterators not used to access outputs to be reduction";
+  case MatchConvolutionResult::Success:
+    return "";
   }
+  llvm_unreachable("unhandled MatchConvolutionResult case");
+}
+
+LogicalResult mlir::linalg::detail::verifyConvolutionInterface(Operation *op) {
+  MatchConvolutionResult res = isConvolutionInterfaceImpl(op);
+  if (res != MatchConvolutionResult::Success)
+    return op->emitError(getMatchConvolutionMessage(res));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // FillOpInterface implementation
 //===----------------------------------------------------------------------===//
 
 enum class MatchFillResult {