Add lit test and clean up

Author: Abhishek-Varma

mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp

@@ -237,6 +237,7 @@ static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
   return replaceWithMatmulVariant<MatmulOp>(rewriter, genericOp);
 }
 
+/// Utility to create a `genericOp` with a convolution op of type `ConvOpTy` with `dilations` and `strides`.
 template <typename ConvOpTy>
 static FailureOr<LinalgOp> specializeToConvOp(RewriterBase &rewriter, GenericOp genericOp, ArrayRef<int64_t> dilations, ArrayRef<int64_t> strides) {
   SmallVector<Value> inputs = genericOp.getDpsInputs();
@@ -380,7 +381,7 @@ static FailureOr<LinalgOp> inferAndSpecializeBasedOnRank9ConvIteratorTypes(Rewri
   return failure();
 }
 
-// Converts linalg.generic to named linalg.*conv* where possible.
+// Converts linalg.generic to named linalg.*conv/pooling* where possible. To improve the search speed, the convolution ops have been segregated based on the rank of iterator types array.
 static FailureOr<LinalgOp> inferAndSpecializeToConvolutionOp(RewriterBase &rewriter, GenericOp genericOp) {
   SmallVector<utils::IteratorType> iteratorTypes = genericOp.getIteratorTypesArray();
   unsigned totalIterators = iteratorTypes.size();
@@ -483,7 +484,7 @@ FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(RewriterBase &rewriter,
     return specializeLinalgContractions(rewriter, genericOp);
   }
 
-  // Convolution - e.g. *conv*
+  // Convolution - e.g. *conv/pooling*
   if (isaConvolutionOpInterface(genericOp)) {
     return inferAndSpecializeToConvolutionOp(rewriter, genericOp);
   }

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

@@ -319,6 +319,11 @@ static bool isDimTimesConstantOrDimOnly(AffineExpr expr, AffineExpr &dim, int64_
   return false;
 }
 
+/// Given an array of AffineMaps `indexingMaps` verify the following :-
+///   indexingMaps[0].getResult(iDim) == 
+///         indexingMaps[1].getResult(fDim) * <CST_1> +
+///         indexingMaps[n-1].getResult(oDim) * <CST_2>
+///  where, CST_1 and CST_2 can be any constant.
 static bool matchConvDimAddExprPattern(ArrayAttr indexingMaps, unsigned iDim, unsigned fDim, unsigned oDim,
   int64_t& dilation, int64_t& stride) {
   unsigned iIndex = 0, fIndex = 1, oIndex = indexingMaps.size() - 1;
@@ -348,10 +353,13 @@ static bool matchConvDimAddExprPattern(ArrayAttr indexingMaps, unsigned iDim, un
   return false;
 }
 
+/// Given an array of AffineMaps `indexingMaps` verify the following :-
+///   indexingMaps[aIndex].getResult(aDim) == indexingMaps[bIndex].getResult(bDim)
 static bool matchConvDimExprPattern(ArrayAttr indexingMaps, unsigned aIndex, unsigned aDim, unsigned bIndex, unsigned bDim) {
   return getAffineMapDim(indexingMaps, aIndex, aDim) == getAffineMapDim(indexingMaps, bIndex, bDim);
 }
 
+/// Give an array of AffineMaps, verify each map to be of the corresponding `expectedSize`.
 static bool verifyConvIndexingMapSizes(ArrayAttr indexingMaps, ArrayRef<int64_t> expectedSizes) {
   if (indexingMaps.size() != expectedSizes.size()) return false;
 
@@ -362,6 +370,7 @@ static bool verifyConvIndexingMapSizes(ArrayAttr indexingMaps, ArrayRef<int64_t>
   return true;
 }
 
+/// Utility to update `dilations` and `strides` by copy the corresponding data from `tempDilations` and `tempStrides`.
 static bool updateConvDilationsAndStrides(SmallVector<int64_t>* dilations, SmallVector<int64_t>* strides, ArrayRef<int64_t> tempDilations, ArrayRef<int64_t> tempStrides) {
   if (!(dilations && strides))
     return true;