xref: /llvm-project/mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td (revision 5a9bdd85ee4d8527e2cedf44f3ce26ff414f9b6a)
1//===- MemRefTransformOps.td - MemRef transformation ops --*- 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 MEMREF_TRANSFORM_OPS
10#define MEMREF_TRANSFORM_OPS
11
12include "mlir/Dialect/Transform/IR/TransformDialect.td"
13include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
14include "mlir/Dialect/Transform/IR/TransformTypes.td"
15include "mlir/Interfaces/SideEffectInterfaces.td"
16include "mlir/IR/OpBase.td"
17
18def MemrefToLLVMTypeConverterOp : Op<Transform_Dialect,
19    "apply_conversion_patterns.memref.memref_to_llvm_type_converter",
20    [DeclareOpInterfaceMethods<TypeConverterBuilderOpInterface,
21                               ["getTypeConverter",
22                                "getTypeConverterType"]>]> {
23  let description = [{
24    This operation provides an "LLVMTypeConverter" that lowers memref types to
25    LLVM types.
26
27    The type converter can be customized as follows:
28    - `use_aligned_alloc`: Use aligned_alloc in place of malloc for heap
29      allocations.
30    - `index_bitwidth`: Bitwidth of the index type, "0" indicates the size of a
31      machine word.
32    - `use_generic_functions`: Use generic allocation and deallocation functions
33      instead of the classic "malloc", "aligned_alloc" and "free" functions.
34    // TODO: the following two options don't really make sense for
35    // memref_to_llvm_type_converter specifically.
36    // We should have a single to_llvm_type_converter.
37    - `use_bare_ptr_call_conv`: Replace FuncOp's MemRef arguments with bare
38      pointers to the MemRef element types.
39    - `data-layout`: String description (LLVM format) of the data layout that is
40      expected on the produced module.
41  }];
42
43  let arguments = (ins
44      DefaultValuedOptionalAttr<BoolAttr, "false">:$use_aligned_alloc,
45      DefaultValuedOptionalAttr<I64Attr, "64">:$index_bitwidth,
46      DefaultValuedOptionalAttr<BoolAttr, "false">:$use_generic_functions,
47      DefaultValuedOptionalAttr<BoolAttr, "false">:$use_bare_ptr_call_conv,
48      OptionalAttr<StrAttr>:$data_layout);
49  let assemblyFormat = "attr-dict";
50}
51
52def ApplyAllocToAllocaOp : Op<Transform_Dialect,
53    "apply_patterns.memref.alloc_to_alloca",
54    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface, ["populatePatternsWithState"]>]> {
55  let description = [{
56    Collects patterns to rewrite scoped dynamic allocation (`alloc`/`dealloc`
57    pairs) into automatic allocation (`alloca`) in the same scope, for memrefs
58    of static shape.
59
60    The `size_limit` attribute controls the maximum allocated memory (in bytes,
61    subject to data layout) for which the pattern applies.
62  }];
63
64  let arguments = (ins
65    OptionalAttr<I64Attr>:$size_limit);
66  let assemblyFormat = "(`size_limit` `(` $size_limit^ `)`)? attr-dict";
67}
68
69def ApplyExpandOpsPatternsOp : Op<Transform_Dialect,
70    "apply_patterns.memref.expand_ops",
71    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
72  let description = [{
73    Collects patterns to rewrite ops within the memref dialect.
74
75    - Converts `atomic_rmw` that cannot be lowered to a simple atomic op with
76      AtomicRMWOpLowering pattern, e.g. with "minf" or "maxf" attributes, to
77      `memref.generic_atomic_rmw` with the expanded code.
78    - Converts `memref.reshape` that has a target shape of a statically-known
79      size to `memref.reinterpret_cast`.
80  }];
81
82  let assemblyFormat = "attr-dict";
83}
84
85def ApplyExpandStridedMetadataPatternsOp : Op<Transform_Dialect,
86    "apply_patterns.memref.expand_strided_metadata",
87    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
88  let description = [{
89    Collects patterns for expanding memref operations that modify the metadata
90    (sizes, offset, strides) of a memref into easier to analyze constructs.
91  }];
92
93  let assemblyFormat = "attr-dict";
94}
95
96def ApplyExtractAddressComputationsPatternsOp : Op<Transform_Dialect,
97    "apply_patterns.memref.extract_address_computations",
98    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
99  let description = [{
100    Collects patterns for extracting address computations from operations
101    with memory accesses such that these memory accesses use only a base
102    pointer.
103
104    For instance,
105    ```mlir
106    memref.load %base[%off0, ...]
107    ```
108
109    Will be rewritten in:
110    ```mlir
111    %new_base = memref.subview %base[%off0,...][1,...][1,...]
112    memref.load %new_base[%c0,...]
113    ```
114  }];
115
116  let assemblyFormat = "attr-dict";
117}
118
119def ApplyFoldMemrefAliasOpsPatternsOp : Op<Transform_Dialect,
120    "apply_patterns.memref.fold_memref_alias_ops",
121    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
122  let description = [{
123    Collects patterns for folding memref aliasing ops (memref.subview) into
124    consumer load/store ops (affine.load, memref.load, nvgpu.ldmatrix,
125    vector.load, vector.transfer_read, affine.store, memref.store, etc.) and
126    other ops (e.g., memref.subview).
127  }];
128
129  let assemblyFormat = "attr-dict";
130}
131
132def ApplyResolveRankedShapedTypeResultDimsPatternsOp : Op<Transform_Dialect,
133    "apply_patterns.memref.resolve_ranked_shaped_type_result_dims",
134    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
135  let description = [{
136    Collects patterns that resolve `memref.dim` operations with values that are
137    defined by operations that implement the `ReifyRankedShapedTypeOpInterface`,
138    in terms of shapes of its input operands.
139  }];
140
141  let assemblyFormat = "attr-dict";
142}
143
144def Transform_MemRefAllocOp : Transform_ConcreteOpType<"memref.alloc">;
145def Transform_MemRefAllocaOp : Transform_ConcreteOpType<"memref.alloca">;
146
147def MemRefAllocaToGlobalOp :
148  Op<Transform_Dialect, "memref.alloca_to_global",
149     [TransformOpInterface,
150      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
151      DeclareOpInterfaceMethods<TransformOpInterface>]> {
152  let description = [{
153    Inserts a new `memref.global` for each provided `memref.alloca` into the
154    nearest symbol table (e.g., a `builtin.module`) and replaces it with a
155    `memref.get_global`. This is useful, for example, for allocations that
156    should reside in the shared memory of a GPU, which have to be declared as
157    globals.
158
159    #### Example
160
161    Consider the following transform op:
162
163    ```mlir
164    %get_global, %global =
165        transform.memref.alloca_to_global %alloca
166          : (!transform.op<"memref.alloca">)
167            -> (!transform.any_op, !transform.any_op)
168    ```
169
170    and the following input payload:
171
172    ```mlir
173    module {
174      func.func @func() {
175        %alloca = memref.alloca() : memref<2x32xf32>
176        // usages of %alloca...
177      }
178    }
179    ```
180
181    then applying the transform op to the payload would result in the following
182    output IR:
183
184    ```mlir
185    module {
186      memref.global "private" @alloc : memref<2x32xf32>
187      func.func @func() {
188        %alloca = memref.get_global @alloc : memref<2x32xf32>
189        // usages of %alloca...
190      }
191    }
192    ```
193
194    #### Return modes
195
196    Succeeds always. The returned handles refer to the `memref.get_global` and
197    `memref.global` ops that were inserted by the transformation.
198  }];
199
200  let arguments = (ins Transform_MemRefAllocaOp:$alloca);
201  let results = (outs TransformHandleTypeInterface:$getGlobal,
202                  TransformHandleTypeInterface:$global);
203
204  let assemblyFormat = [{
205    $alloca attr-dict `:` functional-type(operands, results)
206  }];
207}
208
209def MemRefMultiBufferOp : Op<Transform_Dialect, "memref.multibuffer",
210    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
211     DeclareOpInterfaceMethods<TransformOpInterface>]> {
212  let summary = "Multibuffers an allocation";
213  let description = [{
214     Transformation to do multi-buffering/array expansion to remove
215     dependencies on the temporary allocation between consecutive loop
216     iterations. This transform expands the size of an allocation by
217     a given multiplicative factor and fixes up any users of the
218     multibuffered allocation.
219     If skip analysis is not set the transformation will only apply
220     if it can prove that there is no data being carried across loop
221     iterations.
222
223     #### Return modes
224
225     This operation returns the new allocation if multi-buffering
226     succeeds, and failure otherwise.
227  }];
228
229  let arguments =
230      (ins Transform_MemRefAllocOp:$target,
231           ConfinedAttr<I64Attr, [IntPositive]>:$factor,
232           UnitAttr:$skip_analysis);
233
234  let results = (outs TransformHandleTypeInterface:$transformed);
235
236  let assemblyFormat =
237    "$target attr-dict `:` functional-type(operands, results)";
238}
239
240def MemRefEraseDeadAllocAndStoresOp
241    : Op<Transform_Dialect, "memref.erase_dead_alloc_and_stores", [
242      TransformEachOpTrait, TransformOpInterface,
243      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
244      ReportTrackingListenerFailuresOpTrait
245    ]> {
246  let description = [{
247    This applies memory optimization on memref. In particular it does store to
248    load forwarding, dead store elimination and dead alloc elimination.
249
250    #### Return modes
251
252    This operation applies a set of memory optimization on the whole region of
253    the operand.
254
255    The transformation does not consume the target handle. It modifies the
256    payload. Dead allocations, loads and stores are silently dropped from all
257    mappings.
258  }];
259
260  let arguments = (ins TransformHandleTypeInterface:$target);
261  let results = (outs);
262
263  let assemblyFormat = "$target attr-dict `:` functional-type($target, results)";
264
265  let skipDefaultBuilders = 1;
266  let builders = [
267    OpBuilder<(ins "Value":$target)>
268  ];
269  let extraClassDeclaration = [{
270    ::mlir::DiagnosedSilenceableFailure applyToOne(
271        ::mlir::transform::TransformRewriter &rewriter,
272        ::mlir::Operation *target,
273        ::mlir::transform::ApplyToEachResultList &results,
274        ::mlir::transform::TransformState &state);
275  }];
276}
277
278def MemRefMakeLoopIndependentOp
279    : Op<Transform_Dialect, "memref.make_loop_independent",
280         [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
281          TransformOpInterface, TransformEachOpTrait]> {
282  let description = [{
283    Rewrite the targeted ops such that their index-typed operands no longer
284    depend on any loop induction variable of the `num_loop` enclosing `scf.for`
285    loops. I.e., compute an upper bound that is independent of any such loop IV
286    for every tensor dimension. The transformed op could then be hoisted from
287    the `num_loop` enclosing loops. To preserve the original semantics, place a
288    `memref.subview` inside the loop.
289
290    Currently supported operations are:
291    - memref.alloca: Replaced with a new memref.alloca with upper bound sizes,
292      followed by a memref.subview.
293
294    #### Return modes
295
296    This operation fails if at least one induction variable could not be
297    eliminated. In case the targeted op is already independent of induction
298    variables, this transform succeeds and returns the unmodified target op.
299
300    Otherwise, the returned handle points to a subset of the produced ops:
301    - memref.alloca: The returned handle points to the memref.subview op.
302
303    This transform op consumes the target handle and produces a result handle.
304  }];
305
306  let arguments = (ins TransformHandleTypeInterface:$target, I64Attr:$num_loops);
307  let results = (outs TransformHandleTypeInterface:$transformed);
308  let assemblyFormat =
309      "$target attr-dict `:` functional-type($target, $transformed)";
310
311  let extraClassDeclaration = [{
312    ::mlir::DiagnosedSilenceableFailure applyToOne(
313        ::mlir::transform::TransformRewriter &rewriter,
314        ::mlir::Operation *target,
315        ::mlir::transform::ApplyToEachResultList &results,
316        ::mlir::transform::TransformState &state);
317  }];
318}
319
320#endif // MEMREF_TRANSFORM_OPS
321