[Mlir-commits] [mlir] [mlir][linalg] Support pack consumer fusion with padding semantic for perfect tiling. (PR #149600)

[Han-Chung Wang]

================
@@ -932,13 +932,18 @@ struct PackOpTiling
           continue;
         }
 
-        // If the dimension needs padding, it is not supported because there are
-        // iterations that only write padding values to the whole tile. The
-        // consumer fusion is driven by the source, so it is not possible to map
-        // an empty slice to the tile.
-        bool needExtraPadding =
-            ShapedType::isDynamic(destDimSize) || !cstInnerSize ||
-            destDimSize * cstInnerSize.value() != srcDimSize;
+        // If the dimension needs extra padding, it is not supported because
+        // there are iterations that only write padding values to the whole
+        // tile. The consumer fusion is driven by the source, so it is not
+        // possible to map an empty slice to the tile. Extra padding is not a
+        // regular form, and the implementation is being conversative.
+        bool needExtraPadding = true;
+        if (!ShapedType::isDynamic(srcDimSize) &&
+            !ShapedType::isDynamic(destDimSize) && cstInnerSize) {
+          needExtraPadding =
+              destDimSize >
+              (srcDimSize + cstInnerSize.value() - 1) / cstInnerSize.value();
+        }
----------------
hanhanW wrote:

Upstream prototype: https://github.com/llvm/llvm-project/pull/149624

Testing IREE downstream project here: https://github.com/iree-org/iree/pull/21424

I'll take a look at the result on Monday.

https://github.com/llvm/llvm-project/pull/149600