1//===- TensorTransformOps.td - Tensor 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 TENSOR_TRANSFORM_OPS 10#define TENSOR_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 ApplyDecomposeTensorConcatPatternsOp : Op<Transform_Dialect, 19 "apply_patterns.tensor.decompose_concat", 20 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 21 let description = [{ 22 Indicates that tensor.concat ops should be decomposed into a chain of 23 tensor.insert_slice operations inserting into a materialized destination. 24 }]; 25 26 let assemblyFormat = "attr-dict"; 27} 28 29 30def ApplyDropRedundantInsertSliceRankExpansionPatternsOp : Op<Transform_Dialect, 31 "apply_patterns.tensor.drop_redundant_insert_slice_rank_expansion", 32 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 33 let description = [{ 34 Indicates that redundant tensor.insert_slice rank reductions should be 35 dropped. E.g., cases where a tensor.extract_slice rank reduction immediately 36 follows an inverse tensor.insert_slice rank expansion. 37 }]; 38 39 let assemblyFormat = "attr-dict"; 40} 41 42def ApplyFoldTensorEmptyPatternsOp : Op<Transform_Dialect, 43 "apply_patterns.tensor.fold_tensor_empty", 44 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 45 let description = [{ 46 Indicates that tensor.extract_slice and reassociative reshapes should be 47 folded into tensor.empty. 48 49 If `fold_single_use_only` is set to "true", only tensor.empty that have a 50 single use are folded. 51 }]; 52 53 let arguments = (ins DefaultValuedAttr<BoolAttr, "false">:$fold_single_use_only); 54 let assemblyFormat = "attr-dict"; 55} 56def ApplyFoldIntoPackAndUnpackPatternsOp : Op<Transform_Dialect, 57 "apply_patterns.tensor.fold_into_pack_and_unpack", 58 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 59 let description = [{ 60 Indicates that operations like tensor.pad and tensor.extract_slice should 61 be folded into tensor.pack and tensor.unpack operations, respectively. 62 }]; 63 64 let assemblyFormat = "attr-dict"; 65} 66 67def ApplyFoldTensorSubsetOpsPatternsOp : Op<Transform_Dialect, 68 "apply_patterns.tensor.fold_tensor_subset_ops", 69 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 70 let description = [{ 71 Indicates that tensor.empty should be folded with tensor.extract_slice, 72 tensor.expand_shape and tensor.collapse_shape. 73 }]; 74 75 let assemblyFormat = "attr-dict"; 76} 77 78def ApplyFoldTensorSubsetOpsIntoVectorTransfersPatternsOp : Op<Transform_Dialect, 79 "apply_patterns.tensor.fold_tensor_subset_ops_into_vector_transfers", 80 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 81 let description = [{ 82 Indicates that tensor.extract_slice -> vector.transfer_read and 83 vector.transfer_write -> tensor.insert_slice op chains should be folded into 84 vector tranfer read and write ops 85 }]; 86 87 let assemblyFormat = "attr-dict"; 88} 89 90def ApplyMergeConsecutiveInsertExtractSlicePatternsOp : Op<Transform_Dialect, 91 "apply_patterns.tensor.merge_consecutive_insert_extract_slice", 92 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 93 let description = [{ 94 Indicates that consecutive tensor.extract_slice/tensor.insert_slice ops 95 should be merged into a single op. These patterns are not canonicalizations 96 because the bufferization is sensitive to IR structure. 97 }]; 98 99 let assemblyFormat = "attr-dict"; 100} 101 102def ApplyReassociativeReshapeFoldingPatternsOp : Op<Transform_Dialect, 103 "apply_patterns.tensor.reassociative_reshape_folding", 104 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 105 let description = [{ 106 Indicates that reassociative reshapes (tensor.collapse_shape / 107 tensor.expand_shape) should be folded with inverse rank expansions / rank 108 reductions (via tensor.insert_slice / tensor.extract_slice). 109 }]; 110 111 let assemblyFormat = "attr-dict"; 112} 113 114def ApplyRewriteTensorOpsAsConstantPatternsOp : Op<Transform_Dialect, 115 "apply_patterns.tensor.rewrite_as_constant", 116 [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { 117 let arguments = (ins UnitAttr:$aggressive); 118 let description = [{ 119 Indicates that tensor ops (such as tensor.generate) should be replaced with 120 constants (arith.constant) when possible. 121 }]; 122 123 let assemblyFormat = 124 "(`aggressive` $aggressive^)? attr-dict"; 125} 126 127def Transform_TensorPadOp : Transform_ConcreteOpType<"tensor.pad">; 128 129def MakeLoopIndependentOp 130 : Op<Transform_Dialect, "tensor.make_loop_independent", 131 [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface, 132 TransformOpInterface, TransformEachOpTrait]> { 133 let description = [{ 134 Rewrite the targeted ops such that their index-typed operands no longer 135 depend on any loop induction variable of the `num_loop` enclosing `scf.for` 136 loops. I.e., compute an upper bound that is independent of any such loop IV 137 for every tensor dimension. The transformed op could then be hoisted from 138 the `num_loop` enclosing loops. To preserve the original semantics, place a 139 `tensor.extract_slice` inside the loop. 140 141 Currently supported operations are: 142 - tensor.empty: Replaced with a new tensor.empty with upper bound sizes, 143 followed by a tensor.extract_slice. 144 - tensor.pad: Replaced by an upper bound padding, followed by a 145 tensor.extract_slice. 146 147 #### Return modes 148 149 This operation fails if at least one induction variable could not be 150 eliminated. In case the targeted op is already independent of induction 151 variables, this transform succeeds and returns the unmodified target op. 152 153 Otherwise, the returned handle points to a subset of the produced ops: 154 - tensor.empty: The returned handle points to the tensor.extract_slice op. 155 - tensor.pad: The returned handle points to the tensor.extract_slice op. 156 157 This transform op consumes the target handle and produces a result handle. 158 }]; 159 160 let arguments = (ins TransformHandleTypeInterface:$target, I64Attr:$num_loops); 161 let results = (outs TransformHandleTypeInterface:$transformed); 162 let assemblyFormat = 163 "$target attr-dict `:` functional-type($target, $transformed)"; 164 165 let extraClassDeclaration = [{ 166 ::mlir::DiagnosedSilenceableFailure applyToOne( 167 ::mlir::transform::TransformRewriter &rewriter, 168 ::mlir::Operation *target, 169 ::mlir::transform::ApplyToEachResultList &results, 170 ::mlir::transform::TransformState &state); 171 }]; 172} 173 174def TypeConversionCastShapeDynamicDimsOp : Op<Transform_Dialect, 175 "type_conversion.tensor.cast_shape_dynamic_dims", 176 [DeclareOpInterfaceMethods<TypeConverterBuilderOpInterface, 177 ["populateTypeMaterializations"]>]> { 178 let description = [{ 179 Populates a type converter with conversion materialization functions that 180 cast a tensor value between two cast-compatible tensors. See `tensor.cast` 181 for more information on cast compatibility between tensors. 182 183 If `ignore_dynamic_info` is not set, this will set an additional constraint 184 that source materializations do not cast dynamic dimensions to static ones. 185 }]; 186 let arguments = (ins UnitAttr:$ignore_dynamic_info); 187 188 let assemblyFormat = 189 "(`ignore_dynamic_info` $ignore_dynamic_info^)? attr-dict"; 190} 191 192#endif // TENSOR_TRANSFORM_OPS 193