[Mlir-commits] [mlir] [MLIR][Transform] Add hoist_extract_to_argument transformation (PR #175004)

[Rolf Morel]

================
@@ -2995,4 +2995,54 @@ def DecomposeWinogradOp : Op<Transform_Dialect,
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Hoist Extract to Linalg Argument
+//===----------------------------------------------------------------------===//
+
+def HoistExtractToArgumentOp : Op<Transform_Dialect, "structured.hoist_extract_to_argument",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     TransformOpInterface, TransformEachOpTrait,
+     ReportTrackingListenerFailuresOpTrait]> {
+  let summary = "Hoists a tensor.extract operation from the body of a linalg.generic to an input argument";
+  let description = [{
+    This transform hoists a `tensor.extract` operation from the body of a `linalg.generic`
+    operation to become an input argument of the generic operation. This transformation is
+    applicable when the extract operation's indices can be expressed as affine expressions
+    of the loop induction variables.
+
+    An extract operation is hoistable if all its indices are produced by:
+    - `linalg.index` operations (mapping to loop dimensions)
+    - `arith.constant_index` operations (constant indices)
+    - `arith.index_cast` operations that cast from block arguments corresponding to
+      constant input operands
+
+    #### Return modes
+
+    This operation only applies to `linalg.generic` operations. If the target is not a
+    `linalg.generic`, the transform succeeds but returns the original operation unchanged.
----------------
rolfmorel wrote:

Have you considered exposing this transform as a pattern rewrite (and expose the pattern via a `transform.apply_patterns.*` op)?

The main difference is that patterns can be easily applied to each matching op in a subtree. The advantage your transform op has is that you can match and filter before applying this specific hoist.

Given that the transformation has a canonicalization flavour to it, do you expect anyone will want to selectively apply the transform only to _some_ `linalg.generic`s (rather than, e.g., all in a function)?

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