1//===- LoopExtensionOps.td - Transform dialect operations --*- tablegen -*-===// 2// 3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4// See https://llvm.org/LICENSE.txt for license information. 5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6// 7//===----------------------------------------------------------------------===// 8 9#ifndef MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS 10#define MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS 11 12include "mlir/Dialect/Transform/IR/TransformDialect.td" 13include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td" 14include "mlir/Interfaces/SideEffectInterfaces.td" 15 16def HoistLoopInvariantSubsetsOp 17 : TransformDialectOp<"loop.hoist_loop_invariant_subsets", 18 [TransformOpInterface, TransformEachOpTrait, 19 DeclareOpInterfaceMethods<MemoryEffectsOpInterface>, 20 ReportTrackingListenerFailuresOpTrait]> { 21 let summary = "Hoist loop invariant subset ops"; 22 let description = [{ 23 This transform hoists loop-invariant subset ops out of the targeted 24 loop-like op. It looks for matching subset extraction/insertion op pairs and 25 hoists them. The loop body operates on a newly introduced region iter_arg. 26 27 Subset ops are hoisted only from the targeted op. If subset ops should be 28 hoisted from an entire loop nest, this transformation must be applied to 29 each loop-like op of the loop nest, starting with the innermost loop and 30 ending with the outermost loop. 31 32 Example: 33 ``` 34 %r = scf.for ... iter_args(%t = %a) -> (tensor<?xf32>) { 35 %0 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32> 36 %1 = "test.foo"(%0) : (tensor<5xf32>) -> (tensor<5xf32>) 37 %2 = tensor.insert_slice %1 into %t[0][5][1] 38 : tensor<5xf32> into tensor<?xf32> 39 scf.yield %2 : tensor<?xf32> 40 } 41 ``` 42 Is transformed to: 43 ``` 44 %0 = tensor.extract_slice %a[0][5][1] : tensor<?xf32> to tensor<5xf32> 45 %new_loop:2 = scf.for ... iter_args(%t = %a, %h = %0) -> (tensor<?xf32>) { 46 %1 = "test.foo"(%h) : (tensor<5xf32>) -> (tensor<5xf32>) 47 scf.yield %t, %2 : tensor<?xf32>, tensor<5xf32> 48 } 49 %r = tensor.insert_slice %new_loop#1 into %new_loop#0 50 : tensor<5xf32> into tensor<?xf32> 51 ``` 52 53 Subset ops are hoisted only if there are no conflicting subset ops. E.g., 54 if there were a second overlapping extraction in the above example, no ops 55 could be hoisted safely. 56 57 This transform reads the target handle and modifies the payload. This 58 transform does not invalidate any handles, but loop-like ops are replaced 59 with new loop-like ops when a subset op is hoisted. The transform rewriter 60 updates all handles accordingly. 61 }]; 62 63 let arguments = (ins TransformHandleTypeInterface:$target); 64 let results = (outs); 65 let assemblyFormat = "$target attr-dict `:` type($target)"; 66 67 let extraClassDeclaration = [{ 68 ::mlir::DiagnosedSilenceableFailure applyToOne( 69 ::mlir::transform::TransformRewriter &rewriter, 70 ::mlir::LoopLikeOpInterface loopLikeOp, 71 ::mlir::transform::ApplyToEachResultList &results, 72 ::mlir::transform::TransformState &state); 73 }]; 74} 75 76#endif // MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS 77