[mlir] Add apply_patterns.linalg.generalize_pack_unpack TD Op

[banach-space]

This PR introduces populateGeneralizePatterns, which collects the
following patterns:

• GeneralizeOuterUnitDimsPackOpPattern,
• GeneralizeOuterUnitDimsUnPackOpPattern (currently a TODO).

These patterns are wrapped in a new Transform Dialect Op:
apply_patterns.linalg.generalize_pack_unpack. This Op facilitates
creating more involved end-to-end compilation pipelines for
tensor.pack and tensor.unpack operations. It will be required in an
upcoming PR building on top of #115698.

No new tests are added in this PR. Instead, existing tests from:

• "generalize-tensor-pack.mlir"

are reused. To achieve this:

• I've updated the test to use
  transform.apply_patterns.linalg.generalize_pack_unpack instead of
  the flag
  --test-linalg-transform-patterns="test-generalize-tensor-pack",
  avoiding artificial tests solely for the TD Op.
• The TD sequence is saved to a new file, "generalize_pack.mlir", and
  pre-loaded using the option:
  --transform-preload-library='transform-library-paths=%p/td/generalize_pack.mlir'
  This avoids duplicating the sequence for every "split" in the input
  file.
• Added "lit.local.cfg" to exclude the "test/Dialect/Linalg/td"
  directory from test discovery, ensuring "generalize_pack.mlir" is
  not treated as a test file.

[hanhanW]

Is it an accidental change?

[banach-space]

No :)

• I am updating this test to use transform.apply_patterns.linalg.generalize_pack_unpack instead of --test-linalg-transform-patterns="test-generalize-tensor-pack". This way I avoid adding artificial tests to exercise the TD Op.
• To avoid duplicating the TD sequence for every "split" in the input file, I am using --transform-preload-library='transform-library-paths=%p/td/generalize_pack.mlir.
• To make sure that "generalize_pack.mlir" (used above) is not treated as a test file, I am adding "lit.local.cfg" to exclude the "test/Dialect/Linalg/td" directory.

That's actually quite a few steps :) I will add this to the commit summary.

[hanhanW]

I see, this is clever. It does simplify the test a lot, nice improvement.

[hanhanW]

Where is the pack op, commands and the checks?

[banach-space]

So this is not a test file, see my reply for lit.local.cfg

[hanhanW]

Sorry for asking this question while the naming is following the existing style. I'm confused about the difference between Generalization and Decomposition. Should we unify the terminology? How do you interpret them?

(I could be the root cause who invented the terms when I wrote the patterns. And I feel that they could be renamed to decomposition now.)

[banach-space]

I'm confused about the difference between Generalization and Decomposition.

Same and I am glad that you've raised this :) "Decompose" would make more sense to me - "Generalization" tends to be use in the context of Linalg named Ops (e.g. linalg.matmul -> linalg.generic). I am happy to rename these as "Decomposition" 👍🏻

[hanhanW]

Oh, that's probably why it put me to use generalization in the first place -- because they were prototyped in linalg_ext, and I could be thinking that linalg_ext -> linalg is a generalization. Anyway, I'm +1 on renaming it to decomposition!

[banach-space]

Renamed in the 2nd commit ... that should probably become a separate PR. WDYT?

[hanhanW]

Having a separate PR sounds good to me. I was in meetings, and I'll take a look at your other replies after lunch.

[banach-space]

#116439

[hanhanW]

Can you revert the last two commits in this case?

[hanhanW]

I'm not familiar with how transform-preload-library works in this case. Do we need the --test-transform-dialect-erase-schedule flag?

[banach-space]

Do we need the --test-transform-dialect-erase-schedule flag?

We do not - old habits :) I will remove it before landing.

[hanhanW]

LGTM, please revert the "generalize -> decompose" commits before landing the PR.

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Andrzej Warzyński (banach-space)

Changes

This PR simply creates populateGeneralizePatterns to collect the
following patterns:

• GeneralizeOuterUnitDimsPackOpPattern,
• GeneralizeOuterUnitDimsUnPackOpPattern (left as a TODO),

and wraps the newly created populate* method into a new Transform
Dialect Op: apply_patterns.linalg.generalize_pack_unpack. This allows
creating more involved e2e compilation pipelines for tensor.pack and
tensor.unpack Ops. I will require this this Op for a forthcoming
PR that will build on top of #115698.

No new tests are added. Instead, I've re-used existing tests:

• "generalize-tensor-pack.mlir".

---

Full diff: https://github.com/llvm/llvm-project/pull/116373.diff

7 Files Affected:

• (modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+12)
• (modified) mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h (+7-2)
• (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+5)
• (modified) mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp (+5)
• (modified) mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir (+1-2)
• (added) mlir/test/Dialect/Linalg/lit.local.cfg (+2)
• (added) mlir/test/Dialect/Linalg/td/generalize_pack.mlir (+12)

diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index f256af2f6b12b8..42057d8d0c9105 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -41,6 +41,18 @@ def ApplyEraseUnnecessaryInputsPatternsOp : Op<Transform_Dialect,
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyGeneralizeTensorPackUnpackPatternsOp
+    : Op<Transform_Dialect, "apply_patterns.linalg.generalize_pack_unpack",
+         [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Collect patterns to generalize tensor.pack and tensor.unpack (i.e. to
+    decompose it into e.g. tensor::PadOp, linalg::transposeOp etc). Requires
+    all outer dims to be unit.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 def ApplyFoldUnitExtentDimsViaReshapesPatternsOp : Op<Transform_Dialect,
     "apply_patterns.linalg.fold_unit_extent_dims_via_reshapes",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 89e9a3b70d2ab3..0b55a76f884331 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1516,8 +1516,8 @@ struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
 };
 
 /// Rewrites a tensor::PackOp into a sequence of:
-///   * tensor::PadOp + linalg::TransposeOp +
-///     tensor::EmptyOp + tensor::InsertSliceOp ops.
+///   * tensor::PadOp + linalg::TransposeOp + tensor::EmptyOp +
+///     tensor::InsertSliceOp ops.
 ///
 /// Required that all the outer dims of the input tensor::PackOp are 1.
 ///
@@ -1683,6 +1683,11 @@ void populateLinalgGenericOpsSpecializationPatterns(
 void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
                                           PatternBenefit benefit = 1);
 
+/// Populates patterns to decompose tensor.pack and tensor.unpack Ops into e.g.
+/// tensor.pad, linalg.transpose, tensor.{insert|extract}_slice. Require all
+/// outer dims to be unit.
+void populateGeneralizePatterns(RewritePatternSet &patterns);
+
 /// Populates patterns to transform linalg.conv_2d_xxx operations into
 /// linalg.generic (for img2col packing) and linalg.matmul.
 /// \see rewriteInIm2Col for more details.
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 1956fc634ef395..a00c609779c3a7 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -229,6 +229,11 @@ void transform::ApplyEraseUnnecessaryInputsPatternsOp::populatePatterns(
   linalg::populateEraseUnnecessaryInputsPatterns(patterns);
 }
 
+void transform::ApplyGeneralizeTensorPackUnpackPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  linalg::populateGeneralizePatterns(patterns);
+}
+
 void transform::ApplyFoldUnitExtentDimsViaReshapesPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   linalg::ControlDropUnitDims options;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index ed9ebca4f306a4..c9eac663675599 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -1618,3 +1618,8 @@ void linalg::populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
       DownscaleSizeOneWindowed2DConvolution<PoolingNchwMaxOp, PoolingNcwMaxOp>>(
       patterns.getContext(), benefit);
 }
+
+void linalg::populateGeneralizePatterns(RewritePatternSet &patterns) {
+  // TODO: Add and test patterns for tensor.unpack
+  patterns.add<GeneralizeOuterUnitDimsPackOpPattern>(patterns.getContext());
+}
diff --git a/mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir b/mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir
index f4b1d9a55f0914..c0c6f82a38179b 100644
--- a/mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir
+++ b/mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir
@@ -1,5 +1,4 @@
-// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-generalize-tensor-pack"  %s | FileCheck %s
-
+// RUN: mlir-opt  --transform-preload-library='transform-library-paths=%p/td/generalize_pack.mlir' -split-input-file  --transform-interpreter --test-transform-dialect-erase-schedule %s | FileCheck %s
 
 func.func @simple_KCRS_to_KCRSsr(%arg0: tensor<?x?xi32>, %arg1: tensor<1x1x?x1xi32>) -> tensor<1x1x?x1xi32> {
   %c8 = arith.constant 8 : index
diff --git a/mlir/test/Dialect/Linalg/lit.local.cfg b/mlir/test/Dialect/Linalg/lit.local.cfg
new file mode 100644
index 00000000000000..62743008a3e3a6
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/lit.local.cfg
@@ -0,0 +1,2 @@
+# Skip the directory with input TD sequences
+config.excludes = ["td"]
diff --git a/mlir/test/Dialect/Linalg/td/generalize_pack.mlir b/mlir/test/Dialect/Linalg/td/generalize_pack.mlir
new file mode 100644
index 00000000000000..62e5b779ff361a
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/td/generalize_pack.mlir
@@ -0,0 +1,12 @@
+module @transforms attributes { transform.with_named_sequence } {
+  transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+    %pack = transform.structured.match ops{["tensor.pack"]} in %module : (!transform.any_op) -> !transform.any_op
+
+    %1 = transform.get_parent_op %pack {isolated_from_above} : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %1 {
+      transform.apply_patterns.linalg.generalize_pack_unpack
+    } : !transform.any_op
+
+    transform.yield
+  }
+}