[mlir][tensor] Expose padding requirement of pack ops to a static method

[hanchung]

Summary

It also simplifies the implementation of the method. The map is not needed in the check.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let extraClassDeclaration = commonExtraClassDeclaration # [{
 
     // Method to get the `ShapedType` of the result based on the inner tiles,
     // position of the inner tiles (innerDimsPos)  and interchange vector of
     // outer loops (outerDimsPerm).
     static ShapedType inferPackedType(ShapedType sourceType,
         ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
         ArrayRef<int64_t> outerDimsPerm = {});
 
+    // Returns true if we have enough static information to catch undefined
+    // behavior when the tile size does not divide perfectly the dimension of
+    // the input tensor.
+    static bool requirePaddingValue(ArrayRef<int64_t> inputShape,
+                                    ArrayRef<int64_t> innerDimsPos,
+                                    ArrayRef<OpFoldResult> innerTiles);
+
     static Value createDestinationTensor(OpBuilder &b, Location loc,
         Value source, ArrayRef<OpFoldResult> innerTileSizes,
         ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm);
 
     /// Build and return a new PackOp that is a clone of the current PackOp with
     /// (innerDimsPos, innerTiles) (resp. outerDimsPerm) are permuted by 
     /// innerPermutation (resp. outerPermutation).
     /// A new `tensor.empty` of the proper shape is built in the process.

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 SmallVector<OpFoldResult> PackOp::getMixedTiles() {
   return getMixedTilesImpl(*this);
 }
 
 SmallVector<int64_t> PackOp::getStaticTiles() {
   return getStaticTilesImpl(*this);
 }
 
-/// Check if we have enough static information to catch undefined behavior when
-/// the tile size does not divide perfectly the dimension of the input tensor.
-static bool
-areNotFullTiles(ArrayRef<int64_t> inputShape,
-                DenseMap<int64_t, OpFoldResult> const &dimAndTileMapping) {
-  int64_t rank = inputShape.size();
-  for (int64_t dim = 0; dim < rank; dim++) {
-    if (ShapedType::isDynamic(inputShape[dim]))
-      continue;
-    auto it = dimAndTileMapping.find(dim);
-    if (it == dimAndTileMapping.end())
+bool PackOp::requirePaddingValue(ArrayRef<int64_t> inputShape,
+                                 ArrayRef<int64_t> innerDimsPos,
+                                 ArrayRef<OpFoldResult> innerTiles) {
+  for (auto [pos, tileSize] : llvm::zip_equal(innerDimsPos, innerTiles)) {
+    if (ShapedType::isDynamic(inputShape[pos]))
       continue;
-    std::optional<int64_t> constantTile = getConstantIntValue(it->second);
+    std::optional<int64_t> constantTile = getConstantIntValue(tileSize);
     if (!constantTile)
       continue;
-    if (inputShape[dim] % (*constantTile) != 0)
+    if (inputShape[pos] % (*constantTile) != 0)
       return true;
   }
   return false;
 }
 
 LogicalResult PackOp::verify() {
   if (failed(commonVerifierPackAndUnPackOp(*this)))
     return failure();
 
   // Verify padding value, and bail out if the tile does not divide the
   // dimension fully. In the case of dynamic tile factors or dimensions, having
   // a partial tile is undefined behavior.
   auto paddingValue = getPaddingValue();
   if (paddingValue &&
       paddingValue.getType() != getSourceType().getElementType()) {
     return emitOpError("expected padding_value has ")
            << getSourceType().getElementType()
            << " but got: " << paddingValue.getType();
   }
 
-  auto dimAndTileMapping = getDimAndTileMapping();
   if (!paddingValue &&
-      areNotFullTiles(getSourceType().getShape(), dimAndTileMapping)) {
+      requirePaddingValue(getSourceType().getShape(), getInnerDimsPos(),
+                          getMixedTiles())) {
     return emitOpError("invalid tile factor provided. Only full tiles are "
                        "supported when padding_value is not set");
   }
   return success();
 }
 
 /// Converts OpFoldResults to int64_t shape entries, unconditionally mapping all
 /// Value's to kDynamic, even if they are arith.constant values.