[llvm] [mlir][transform] LISH: Add transform op (PR #70630)

Oleksandr Alex Zinenko via llvm-commits llvm-commits at lists.llvm.org
Mon Oct 30 02:58:48 PDT 2023


================
@@ -691,6 +691,65 @@ def GetTypeOp : TransformDialectOp<"get_type",
                        "functional-type(operands, results)";
 }
 
+def HoistLoopInvariantSubsetsOp
+    : TransformDialectOp<"hoist_loop_invariant_subsets",
+        [TransformOpInterface, TransformEachOpTrait,
+         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+         ReportTrackingListenerFailuresOpTrait]> {
+  let summary = "Hoist loop invariant subset ops";
+  let description = [{
+    This transform hoist loop-invariant subset ops out of loop-like ops. It
+    looks for matching subset extraction/insertion op pairs and hoists them. The
+    loop body operates on a newly introduced region iter_arg.
+
+    Example:
+    ```
+    %r = scf.for ... iter_args(%t = %a) -> (tensor<?xf32>) {
+      %0 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
+      %1 = "test.foo"(%0) : (tensor<5xf32>) -> (tensor<5xf32>)
+      %2 = tensor.insert_slice %1 into %t[0][5][1]
+          : tensor<5xf32> into tensor<?xf32>
+      scf.yield %2 : tensor<?xf32>
+    }
+    ```
+    Is transformed to:
+    ```
+    %0 = tensor.extract_slice %a[0][5][1] : tensor<?xf32> to tensor<5xf32>
+    %new_loop:2 = scf.for ... iter_args(%t = %a, %h = %0) -> (tensor<?xf32>) {
+      %1 = "test.foo"(%h) : (tensor<5xf32>) -> (tensor<5xf32>)
+      scf.yield %t, %2 : tensor<?xf32>, tensor<5xf32>
+    }
+    %r = tensor.insert_slice %new_loop#1 into %new_loop#0
+        : tensor<5xf32> into tensor<?xf32>
+    ```
+
+    Subset ops are hoisted only if there are no conflicting subset ops. E.g.,
+    if there were a second overlapping extraction in the above example, no ops
+    could be hoisted safely.
+
+    This transform looks for `LoopLikeOpInterface` ops within the targeted op,
+    including the target op itself. It attempts hoisting on all found loop-like
+    ops.
+
+    This transform reads the target handle and modifies the payload.
+
+    TODO: Make this op more targeted if needed. I.e., apply the transformation
+    only to the targeted `LoopLikeOpInterface` op.
----------------
ftynse wrote:

Should we just go ahead and actually do this? It looks like we want to expose `walk` as a separate transform op, which could be useful beyond this transform. 

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


More information about the llvm-commits mailing list