[mlir][linalg] Add stride/dilation inference to convolution dimension classifier

[qedawkins]

Summary

Stride/dilation values are needed for certain transformations so this
infers them from the dim coefficients in the affine expressions for
the convolution. Adds the values to the classifier exposed in D143584.

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

 /// 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;
+  SmallVector<int64_t, 2> strides;
+  SmallVector<int64_t, 2> dilations;
 };
 
 /// 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);
 

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

 ///      (`+` AffineDimExpr (`*` (AffineSymbolExpr | AffineConstantExpr))?)*
 ///
 /// classifies the AffineDimExpr as convolved dimensions or unconvolved
 /// dimensions and verifies each dimension occurs only once.
 struct ConvAccessExprWalker
     : public AffineExprVisitor<ConvAccessExprWalker, LogicalResult> {
   llvm::SmallDenseSet<unsigned> convolvedDims;
   llvm::SmallDenseSet<unsigned> unConvolvedDims;
+  llvm::SmallDenseMap<unsigned, AffineExpr> strideAndDilationMapping;
 
   LogicalResult visitDimExpr(AffineDimExpr dimExpr) {
     unsigned position = dimExpr.getPosition();
     if (unConvolvedDims.count(position) || convolvedDims.count(position)) {
       return failure();
     }
     unConvolvedDims.insert(position);
     return success();
return success(succeeded(isDimExprOrMulExpr(binaryExpr.getLHS())) &&
                    succeeded(isDimExprOrMulExpr(binaryExpr.getRHS())));
   }
 
   LogicalResult isDimExprOrMulExpr(AffineExpr expr) {
     if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) {
       unsigned dim = dimExpr.getPosition();
       if (convolvedDims.count(dim) || unConvolvedDims.count(dim))
         return failure();
+      // Stride/dilation for this dim is implicitly 1.
+      strideAndDilationMapping[dim] =
+          getAffineConstantExpr(1, expr.getContext());
       convolvedDims.insert(dim);
       return success();
     }
     if (auto symbolMulExpr = expr.dyn_cast<AffineBinaryOpExpr>()) {
       if (symbolMulExpr.getKind() != AffineExprKind::Mul)
         return failure();
       auto lhsExpr = symbolMulExpr.getLHS();
       auto rhsExpr = symbolMulExpr.getRHS();
       // Check for symbol expression.
       AffineExpr mulExpr =
           getAffineExprOfType<AffineSymbolExpr>(lhsExpr, rhsExpr);
       // If there was no symbol expr, check for constant expression.
       if (!mulExpr) {
         mulExpr = getAffineExprOfType<AffineConstantExpr>(lhsExpr, rhsExpr);
       }
       auto dimExpr = getAffineExprOfType<AffineDimExpr>(lhsExpr, rhsExpr);
       if (!mulExpr || !dimExpr)
         return failure();
       unsigned dim = dimExpr.getPosition();
       if (convolvedDims.count(dim) || unConvolvedDims.count(dim))
         return failure();
+      strideAndDilationMapping[dim] = mulExpr;
       convolvedDims.insert(dim);
       return success();
     }
     return failure();
   }
 };
 } // namespace
 
 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;
 }
 
+static SmallVector<int64_t, 2>
+getConstantsFromExprList(SmallVector<AffineExpr, 2> exprs) {
+  SmallVector<int64_t, 2> vals;
+  for (auto e : exprs) {
+    auto constantExpr = e.dyn_cast<AffineConstantExpr>();
+    assert(constantExpr && "Found non-constant stride/dilation");
+    vals.push_back(constantExpr.getValue());
+  }
+  return vals;
+}
+
 namespace mlir::linalg::detail {
 enum class MatchConvolutionResult {
   Success = 0,
   NotLinalgOp,
   WrongNumOperands,
   WrongInputIndexingMap,
   NotProjectedPermutations,
   NonConvolutionLoop,
mlir::linalg::detail::isConvolutionInterfaceImpl(
   //   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;
+  llvm::SmallVector<AffineExpr, 2> strideExprs;
   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)
+      if (dimensions) {
+        strideExprs.push_back(
+            inputExprWalker.strideAndDilationMapping[outputDim]);
         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;
if (inputExprWalker.unConvolvedDims.count(outputDim) &&
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
       if (dimensions)
         dimensions->depth.push_back(outputDim);
       continue;
     }
     return MatchConvolutionResult::NonConvolutionLoop;
   }
+  llvm::SmallVector<AffineExpr, 2> dilationExprs;
   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;
       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)
+      if (dimensions) {
+        dilationExprs.push_back(
+            inputExprWalker.strideAndDilationMapping[filterDim]);
         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))
mlir::linalg::detail::isConvolutionInterfaceImpl(
 
   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");
+
+    // Use the op carried strides/dilations attribute if present.
+    auto nativeStrides = op->getAttrOfType<DenseIntElementsAttr>("strides");
+    dimensions->strides =
+        !nativeStrides ? getConstantsFromExprList(strideExprs)
+                       : llvm::to_vector<2>(nativeStrides.getValues<int64_t>());
+    auto nativeDilations = op->getAttrOfType<DenseIntElementsAttr>("dilations");
+    dimensions->dilations =
+        !nativeDilations
+            ? getConstantsFromExprList(dilationExprs)
+            : llvm::to_vector<2>(nativeDilations.getValues<int64_t>());
   }
 
   return MatchConvolutionResult::Success;
 }
 
 StringRef
 mlir::linalg::detail::getMatchConvolutionMessage(MatchConvolutionResult res) {
   switch (res) {