//===----------------------------------------------------------------------===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // //===----------------------------------------------------------------------===// #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Arith/Utils/Utils.h" #include "mlir/Dialect/Complex/IR/Complex.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Utils/IndexingUtils.h" #include "mlir/Dialect/Utils/ReshapeOpsUtils.h" #include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributeInterfaces.h" #include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/OpDefinition.h" #include "mlir/IR/TypeUtilities.h" #include "mlir/Interfaces/DestinationStyleOpInterface.h" #include "mlir/Interfaces/LoopLikeInterface.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/MathExtras.h" #include <algorithm> #include <optional> usingnamespacemlir; usingnamespacemlir::tensor; divideCeilSigned; divideFloorSigned; mod; /// Materialize a single constant operation from a given attribute value with /// the desired resultant type. Operation *TensorDialect::materializeConstant(OpBuilder &builder, Attribute value, Type type, Location loc) { … } OpFoldResult tensor::getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim) { … } SmallVector<OpFoldResult> tensor::getMixedSizes(OpBuilder &builder, Location loc, Value value) { … } FailureOr<Value> tensor::getOrCreateDestination(OpBuilder &b, Location loc, OpResult opResult) { … } LogicalResult tensor::getOrCreateDestinations(OpBuilder &b, Location loc, Operation *op, SmallVector<Value> &result) { … } bool tensor::isSameTypeWithoutEncoding(Type tp1, Type tp2) { … } /// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or /// rank-extending tensor.insert_slice op. static llvm::SmallBitVector getDroppedDims(ArrayRef<int64_t> reducedShape, ArrayRef<OpFoldResult> mixedSizes) { … } /// Given a ranked tensor type and a range of values that defines its dynamic /// dimension sizes, turn all dynamic sizes that have a constant value into /// static dimension sizes. static RankedTensorType foldDynamicToStaticDimSizes(RankedTensorType type, ValueRange dynamicSizes, SmallVector<Value> &foldedDynamicSizes) { … } //===----------------------------------------------------------------------===// // BitcastOp //===----------------------------------------------------------------------===// bool BitcastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) { … } namespace { /// Replaces chains of two tensor.bitcast operations by a single tensor.bitcast /// operation. struct ChainedTensorBitcast : public OpRewritePattern<BitcastOp> { … }; } // namespace void BitcastOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // CastOp //===----------------------------------------------------------------------===// void CastOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) { … } /// Returns true if `target` is a ranked tensor type that preserves static /// information available in the `source` ranked tensor type. bool mlir::tensor::preservesStaticInformation(Type source, Type target) { … } /// Determines whether tensor::CastOp casts to a more dynamic version of the /// source tensor. This is useful to fold a tensor.cast into a consuming op and /// implement canonicalization patterns for ops in different dialects that may /// consume the results of tensor.cast operations. Such foldable tensor.cast /// operations are typically inserted as `slice` ops and are canonicalized, /// to preserve the type compatibility of their uses. /// /// Returns true when all conditions are met: /// 1. source and result are ranked tensors with same element type and rank. /// 2. the tensor type has more static information than the result /// /// Example: /// ```mlir /// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32> /// %2 = consumer %1 ... : tensor<?x?xf32> ... /// ``` /// /// folds into: /// /// ```mlir /// %2 = consumer %0 ... : tensor<8x16xf32> ... /// ``` bool mlir::tensor::canFoldIntoConsumerOp(CastOp castOp) { … } /// Determines whether the tensor::CastOp casts to a more static version of the /// source tensor. This is useful to fold into a producing op and implement /// canonicaliation patterns with the `tensor.cast` op as the root, but producer /// being from different dialects. Returns true when all conditions are met: /// 1. source and result and ranked tensors with same element type and rank. /// 2. the result type has more static information than the source. /// /// Example: /// ```mlir /// %1 = producer ... : tensor<?x?xf32> /// %2 = tensor.cast %1 : tensor<?x?xf32> to tensor<8x16xf32> /// ``` /// /// can be canonicalized to : /// /// ```mlir /// %2 = producer ... : tensor<8x16xf32> /// ``` /// Not all ops might be canonicalizable this way, but for those that can be, /// this method provides a check that it is worth doing the canonicalization. bool mlir::tensor::canFoldIntoProducerOp(CastOp castOp) { … } /// Performs folding of any operand of `op` if it comes from a tensor::CastOp /// that can be folded. LogicalResult mlir::tensor::foldTensorCast(Operation *op) { … } bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) { … } /// Compute a TensorType that has the joined shape knowledge of the two /// given TensorTypes. The element types need to match. static TensorType joinShapes(TensorType one, TensorType two) { … } namespace { /// Replaces chains of two tensor.cast operations by a single tensor.cast /// operation if doing so does not remove runtime constraints. struct ChainedTensorCast : public OpRewritePattern<CastOp> { … }; /// Fold tensor.cast into tesor.extract_slice producer. /// Example: /// ``` /// %0 = tensor.extract_slice %arg0[%o, 0] [%s, 512] [1, 1] : /// tensor<128x512xf32> to tensor<?x512xf32> /// %1 = tensor.cast %0 : tensor<?x512xf32> to tensor<16x512xf32> /// ``` /// -> /// ``` /// %1 = tensor.extract_slice %arg0[%o, 0] [16, 512] [1, 1] : /// tensor<128x512xf32> to tensor<16x512xf32> /// ``` struct TensorCastExtractSlice : public OpRewritePattern<CastOp> { … }; } // namespace void CastOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // ConcatOp //===----------------------------------------------------------------------===// RankedTensorType ConcatOp::inferResultType(int64_t dim, TypeRange inputTypes) { … } void ConcatOp::build(OpBuilder &builder, OperationState &result, int64_t dim, ValueRange inputs) { … } LogicalResult ConcatOp::verify() { … } LogicalResult ConcatOp::reifyResultShapes(OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } void ConcatOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } OpFoldResult ConcatOp::fold(FoldAdaptor) { … } namespace { /// Fold a concat op with a single input to a cast. struct SingleInputConcatOp : public OpRewritePattern<ConcatOp> { … }; } // namespace void ConcatOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // DimOp //===----------------------------------------------------------------------===// void DimOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) { … } void DimOp::build(OpBuilder &builder, OperationState &result, Value source, int64_t index) { … } std::optional<int64_t> DimOp::getConstantIndex() { … } Speculation::Speculatability DimOp::getSpeculatability() { … } OpFoldResult DimOp::fold(FoldAdaptor adaptor) { … } namespace { /// Fold dim of a cast into the dim of the source of the tensor cast. struct DimOfCastOp : public OpRewritePattern<DimOp> { … }; /// Fold dim of a destination passing style op into the dim of the corresponding /// init. struct DimOfDestStyleOp : public OpRewritePattern<DimOp> { … }; /// Fold dim of a tensor reshape operation to a extract into the reshape's shape /// operand. struct DimOfReshapeOp : public OpRewritePattern<DimOp> { … }; } // namespace void DimOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // EmptyOp //===----------------------------------------------------------------------===// void EmptyOp::build(OpBuilder &builder, OperationState &result, ArrayRef<int64_t> staticShape, Type elementType, Attribute encoding) { … } void EmptyOp::build(OpBuilder &builder, OperationState &result, ArrayRef<int64_t> staticShape, Type elementType, ValueRange dynamicSizes, Attribute encoding) { … } void EmptyOp::build(OpBuilder &builder, OperationState &result, ArrayRef<OpFoldResult> sizes, Type elementType, Attribute encoding) { … } LogicalResult EmptyOp::verify() { … } LogicalResult EmptyOp::reifyResultShapes(OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } Value EmptyOp::getDynamicSize(unsigned idx) { … } SmallVector<OpFoldResult> EmptyOp::getMixedSizes() { … } namespace { /// Change the type of the result of a `tensor.empty` by making the result /// type statically sized along dimensions that in the original operation were /// defined as dynamic, but the size was defined using a `constant` op. For /// example /// /// %c5 = arith.constant 5: index /// %0 = tensor.empty(%arg0, %c5) : tensor<?x?xf32> /// /// to /// /// %0 = tensor.empty(%arg0) : tensor<?x5xf32> struct ReplaceEmptyTensorStaticShapeDims : OpRewritePattern<EmptyOp> { … }; struct FoldEmptyTensorWithDimOp : public OpRewritePattern<DimOp> { … }; /// Canonicalize /// /// ```mlir /// %0 = tensor.empty(%d0, %d1) : tensor<?x?xf32> /// %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<4x?xf32> /// ``` /// /// into /// /// ```mlir /// %0 = tensor.empty(%d1) : tensor<4x?xf32> /// ``` /// /// This assumes the input program is correct in terms of its shape. So it is /// safe to assume that `%d0` is in fact 4. struct FoldEmptyTensorWithCastOp : public OpRewritePattern<CastOp> { … }; } // namespace void EmptyOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } /// Try to remove a tensor operation if it would only reshape a constant. /// Removes the op and replaces the constant with a new constant of the result /// shape. When an optional cst attribute is passed, it is reshaped only if the /// splat value matches the value in the attribute. static OpFoldResult reshapeConstantSource(DenseElementsAttr source, TensorType result, std::optional<Attribute> cst = std::nullopt) { … } //===----------------------------------------------------------------------===// // ExtractOp //===----------------------------------------------------------------------===// namespace { /// Canonicalizes the pattern of the form /// /// %val = tensor.cast %source : : tensor<?xi32> to tensor<2xi32> /// %extracted_element = tensor.extract %val[%c0] : tensor<2xi32> /// /// to /// /// %extracted_element = tensor.extract %source[%c0] : tensor<?xi32> struct ExtractFromTensorCast : public OpRewritePattern<tensor::ExtractOp> { … }; } // namespace void ExtractOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } LogicalResult ExtractOp::verify() { … } OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) { … } void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // FromElementsOp //===----------------------------------------------------------------------===// void FromElementsOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } void FromElementsOp::build(OpBuilder &builder, OperationState &result, ValueRange elements) { … } OpFoldResult FromElementsOp::fold(FoldAdaptor adaptor) { … } namespace { // Pushes the index_casts that occur before extractions to after the extract. // This minimizes type conversion in some cases and enables the extract // canonicalizer. This changes: // // %cast = arith.index_cast %tensor : tensor<1xi32> to tensor<1xindex> // %extract = tensor.extract %cast[%index] : tensor<1xindex> // // to the following: // // %extract = tensor.extract %tensor[%index] : tensor<1xindex> // %cast = arith.index_cast %extract : i32 to index // // to just %element. // // Consider expanding this to a template and handle all tensor cast // operations. struct ExtractElementFromIndexCast : public OpRewritePattern<tensor::ExtractOp> { … }; } // namespace void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // GatherOp //===----------------------------------------------------------------------===// void GatherOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } /// Return the inferred result type for a gatherOp where: /// - sourceType is the type of the source tensor gathered from /// - indicesType is the type of the indices used to gather /// - gatherDims are the dims along which the gather occurs. /// Return a full rank or ranked-reduced variant of the type depending on /// the value of rankReduced. /// /// The leading dimensions of the index tensor give the result tensor its /// leading dimensions. /// The trailing dimensions of the result tensor are obtained from the source /// tensor by setting the dimensions specified in gather_dims to `1` (if /// rankedReduced is false), or skipping them (otherwise). RankedTensorType GatherOp::inferResultType(RankedTensorType sourceType, RankedTensorType indicesType, ArrayRef<int64_t> gatherDims, bool rankReduced) { … } static LogicalResult verifyGatherOrScatterDims(Operation *op, ArrayRef<int64_t> dims, int64_t rank, StringRef gatherOrScatter, StringRef sourceOrDest) { … } LogicalResult GatherOp::verify() { … } OpFoldResult GatherOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // InsertOp //===----------------------------------------------------------------------===// void InsertOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } LogicalResult InsertOp::verify() { … } OpFoldResult InsertOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // GenerateOp //===----------------------------------------------------------------------===// void GenerateOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } LogicalResult GenerateOp::reifyResultShapes( OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } LogicalResult GenerateOp::verify() { … } LogicalResult GenerateOp::verifyRegions() { … } void GenerateOp::build( OpBuilder &b, OperationState &result, Type resultTy, ValueRange dynamicExtents, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) { … } namespace { /// Canonicalizes tensor.generate operations with a constant /// operand into the equivalent operation with the operand expressed in the /// result type, instead. We also insert a type cast to make sure that the /// resulting IR is still well-typed. struct StaticTensorGenerate : public OpRewritePattern<GenerateOp> { … }; /// Canonicalizes the pattern of the form /// /// %tensor = tensor.generate %x { /// ^bb0(%arg0: index): /// <computation> /// yield %1 : index /// } : tensor<?xindex> /// %extracted_element = tensor.extract %tensor[%c0] : tensor<?xi32> /// /// to just <computation> with %arg0 replaced by %c0. We only do this if the /// tensor.generate operation has no side-effects. struct ExtractFromTensorGenerate : public OpRewritePattern<tensor::ExtractOp> { … }; } // namespace void GenerateOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // RankOp //===----------------------------------------------------------------------===// void RankOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) { … } OpFoldResult RankOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // ReshapeOp //===----------------------------------------------------------------------===// void ReshapeOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } static int64_t getNumElements(ShapedType type) { … } LogicalResult ReshapeOp::verify() { … } OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // Reassociative reshape ops //===----------------------------------------------------------------------===// void CollapseShapeOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } void ExpandShapeOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } int64_t ExpandShapeOp::getCorrespondingSourceDim(int64_t resultDim) { … } FailureOr<SmallVector<OpFoldResult>> ExpandShapeOp::inferOutputShape(OpBuilder &b, Location loc, RankedTensorType expandedType, ArrayRef<ReassociationIndices> reassociation, ArrayRef<OpFoldResult> inputShape) { … } void ExpandShapeOp::build(OpBuilder &builder, OperationState &result, Type resultType, Value src, ArrayRef<ReassociationIndices> reassociation, ArrayRef<OpFoldResult> outputShape) { … } void ExpandShapeOp::build(OpBuilder &builder, OperationState &result, Type resultType, Value src, ArrayRef<ReassociationIndices> reassociation) { … } SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() { … } SmallVector<ReassociationExprs, 4> CollapseShapeOp::getReassociationExprs() { … } SmallVector<AffineMap, 4> ExpandShapeOp::getReassociationMaps() { … } SmallVector<ReassociationExprs, 4> ExpandShapeOp::getReassociationExprs() { … } RankedTensorType CollapseShapeOp::inferCollapsedType( RankedTensorType type, SmallVector<ReassociationIndices> reassociation) { … } /// Compute the RankedTensorType obtained by applying `reassociation` to /// `type`. RankedTensorType CollapseShapeOp::inferCollapsedType(RankedTensorType type, ArrayRef<AffineMap> reassociation) { … } void CollapseShapeOp::build(OpBuilder &b, OperationState &result, Value src, ArrayRef<ReassociationIndices> reassociation, ArrayRef<NamedAttribute> attrs) { … } template <typename TensorReshapeOp, bool isExpansion = std::is_same< TensorReshapeOp, ExpandShapeOp>::value> static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op, RankedTensorType expandedType, RankedTensorType collapsedType) { … } LogicalResult ExpandShapeOp::verify() { … } LogicalResult CollapseShapeOp::verify() { … } namespace { /// Reshape of a splat constant can be replaced with a constant of the result /// type. template <typename TensorReshapeOp> struct FoldReshapeWithConstant : OpRewritePattern<TensorReshapeOp> { … }; // Folds TensorReshapeOp(splat x : src_type) : res_type into splat x : res_type. template <typename TensorReshapeOp> class FoldReshapeWithSplat : public OpRewritePattern<TensorReshapeOp> { … }; /// Reshape of a FromElements can be replaced with a FromElements of the /// result type template <typename TensorReshapeOp> struct FoldReshapeWithFromElements : OpRewritePattern<TensorReshapeOp> { … }; // Fold CastOp into CollapseShapeOp when adding static information. struct FoldCollapseOfCastOp : public OpRewritePattern<CollapseShapeOp> { … }; struct FoldDimOfExpandShape : public OpRewritePattern<DimOp> { … }; struct FoldDimOfCollapseShape : public OpRewritePattern<DimOp> { … }; } // namespace void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } OpFoldResult ExpandShapeOp::fold(FoldAdaptor adaptor) { … } OpFoldResult CollapseShapeOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // ExtractSliceOp //===----------------------------------------------------------------------===// void ExtractSliceOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } /// An extract_slice result type can be inferred, when it is not /// rank-reduced, from the source type and the static representation of /// offsets, sizes and strides. Special sentinels encode the dynamic case. RankedTensorType ExtractSliceOp::inferResultType( RankedTensorType sourceTensorType, ArrayRef<int64_t> staticOffsets, ArrayRef<int64_t> staticSizes, ArrayRef<int64_t> staticStrides) { … } RankedTensorType ExtractSliceOp::inferResultType( RankedTensorType sourceTensorType, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) { … } /// If the rank is reduced (i.e. the desiredResultRank is smaller than the /// number of sizes), drop as many size 1 as needed to produce an inferred /// type with the desired rank. /// /// Note that there may be multiple ways to compute this rank-reduced type: /// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. /// /// To disambiguate, this function always drops the first 1 sizes occurrences. RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides) { … } RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) { … } /// Build an ExtractSliceOp with mixed static and dynamic entries and custom /// result type. If the type passed is nullptr, it is inferred. void ExtractSliceOp::build(OpBuilder &b, OperationState &result, RankedTensorType resultType, Value source, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides, ArrayRef<NamedAttribute> attrs) { … } /// Build an ExtractSliceOp with mixed static and dynamic entries and inferred /// result type. void ExtractSliceOp::build(OpBuilder &b, OperationState &result, Value source, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides, ArrayRef<NamedAttribute> attrs) { … } /// Build an ExtractSliceOp with mixed static and dynamic entries packed into /// a Range vector. void ExtractSliceOp::build(OpBuilder &b, OperationState &result, Value source, ArrayRef<Range> ranges, ArrayRef<NamedAttribute> attrs) { … } /// Build an ExtractSliceOp with dynamic entries and custom result type. If /// the type passed is nullptr, it is inferred. void ExtractSliceOp::build(OpBuilder &b, OperationState &result, RankedTensorType resultType, Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef<NamedAttribute> attrs) { … } /// Build an ExtractSliceOp with dynamic entries and inferred result type. void ExtractSliceOp::build(OpBuilder &b, OperationState &result, Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef<NamedAttribute> attrs) { … } static LogicalResult produceSliceErrorMsg(SliceVerificationResult result, Operation *op, RankedTensorType expectedType) { … } /// Verifier for ExtractSliceOp. LogicalResult ExtractSliceOp::verify() { … } llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { … } FailureOr<Value> ExtractSliceOp::rankReduceIfNeeded(OpBuilder &b, Location loc, Value value, ArrayRef<int64_t> desiredShape) { … } LogicalResult ExtractSliceOp::reifyResultShapes( OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } namespace { /// Pattern to rewrite an extract_slice op with tensor::Cast arguments. /// This essentially pushes memref_cast past its consuming slice when /// `canFoldIntoConsumerOp` is true. /// /// Example: /// ``` /// %0 = tensor.cast %V : tensor<16x16xf32> to tensor<?x?xf32> /// %1 = tensor.extract_slice %0[0, 0][3, 4][1, 1] : tensor<?x?xf32> to /// tensor<3x4xf32> /// ``` /// is rewritten into: /// ``` /// %0 = tensor.extract_slice %V[0, 0][3, 4][1, 1] : tensor<16x16xf32> to /// tensor<3x4xf32> %1 = tensor.cast %0: tensor<3x4xf32> to tensor<3x4xf32> /// ``` class ExtractSliceOpCastFolder final : public OpRewritePattern<ExtractSliceOp> { … }; /// Slice elements from `values` into `outValues`. `counts` represents the /// numbers of elements to stride in the original values for each dimension. /// The output values can be used to construct a DenseElementsAttr. template <typename IterTy, typename ElemTy> static void sliceElements(IterTy values, ArrayRef<int64_t> counts, ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides, llvm::SmallVectorImpl<ElemTy> *outValues) { … } /// Fold arith.constant and tensor.extract_slice into arith.constant. The /// folded operation might introduce more constant data; Users can control /// their heuristics by the control function. class ConstantOpExtractSliceFolder final : public OpRewritePattern<ExtractSliceOp> { … }; } // namespace void mlir::tensor::populateFoldConstantExtractSlicePatterns( RewritePatternSet &patterns, const ControlConstantExtractSliceFusionFn &controlFn) { … } /// Return the canonical type of the result of an extract_slice op. struct SliceReturnTypeCanonicalizer { … }; /// A canonicalizer wrapper to replace ExtractSliceOps. struct SliceCanonicalizer { … }; void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } // static LogicalResult foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, ShapedType shapedType) { … } /// If we have an ExtractSliceOp consuming an InsertSliceOp with the same /// slice, we can return the InsertSliceOp's source directly. // TODO: This only checks the immediate producer; extend to go up the // insert/extract chain if the slices are disjoint. static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { … } OpFoldResult ExtractSliceOp::fold(FoldAdaptor adaptor) { … } Value mlir::tensor::createCanonicalRankReducingExtractSliceOp( OpBuilder &b, Location loc, Value tensor, RankedTensorType targetType) { … } //===----------------------------------------------------------------------===// // InsertSliceOp //===----------------------------------------------------------------------===// void InsertSliceOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } // Build a InsertSliceOp with mixed static and dynamic entries. void InsertSliceOp::build(OpBuilder &b, OperationState &result, Value source, Value dest, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides, ArrayRef<NamedAttribute> attrs) { … } /// Build an InsertSliceOp with mixed static and dynamic entries packed into a /// Range vector. void InsertSliceOp::build(OpBuilder &b, OperationState &result, Value source, Value dest, ArrayRef<Range> ranges, ArrayRef<NamedAttribute> attrs) { … } // Build a InsertSliceOp with dynamic entries. void InsertSliceOp::build(OpBuilder &b, OperationState &result, Value source, Value dest, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef<NamedAttribute> attrs) { … } /// Rank-reducing type verification for both InsertSliceOp and /// ParallelInsertSliceOp. static SliceVerificationResult verifyInsertSliceOp( RankedTensorType srcType, RankedTensorType dstType, ArrayRef<int64_t> staticOffsets, ArrayRef<int64_t> staticSizes, ArrayRef<int64_t> staticStrides, RankedTensorType *expectedType = nullptr) { … } /// Verifier for InsertSliceOp. LogicalResult InsertSliceOp::verify() { … } /// If we have two consecutive InsertSliceOp writing to the same slice, we /// can mutate the second InsertSliceOp's destination to the first one's. /// /// Example: /// /// ```mlir /// %0 = tensor.insert_slice %slice0 into %input[0, 0] [64, 64] [1, 1] /// %1 = tensor.insert_slice %slice1 into %0[0, 0] [64, 64] [1, 1] /// ``` /// /// folds into: /// /// ```mlir /// %1 = tensor.insert_slice %slice1 into %input[0, 0] [64, 64] [1, 1] /// ``` /// /// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. static LogicalResult foldInsertAfterInsertSlice(InsertSliceOp insertOp) { … } /// Folds round-trip extract/insert slice op pairs. /// Example: /// ```mlir /// %0 = tensor.extract_slice %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] /// %1 = tensor.insert_slice %0 into %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] /// ``` /// can be folded into %val. static Value foldInsertAfterExtractSlice(InsertSliceOp insertOp) { … } OpFoldResult InsertSliceOp::fold(FoldAdaptor) { … } LogicalResult InsertSliceOp::reifyResultShapes( OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } namespace { /// Pattern to rewrite a insert_slice op with constant arguments. /// /// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. template <typename InsertOpTy> class InsertSliceOpConstantArgumentFolder final : public OpRewritePattern<InsertOpTy> { … }; /// Fold tensor_casts with insert_slice operations. If the source or /// destination tensor is a tensor_cast that removes static type information, /// the cast is folded into the insert_slice operation. E.g.: /// /// ```mlir /// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32> /// %2 = tensor.insert_slice %1 into ... : tensor<?x?xf32> into ... /// ``` /// /// folds into: /// /// ```mlir /// %2 = tensor.insert_slice %0 into ... : tensor<8x16xf32> into ... /// ``` /// /// Note: When folding a cast on the destination tensor, the result of the /// insert_slice operation is casted to ensure that the type of the result did /// not change. /// /// This pattern works with both InsertSliceOp and ParallelInsertSliceOp. template <typename InsertOpTy> struct InsertSliceOpCastFolder final : public OpRewritePattern<InsertOpTy> { … }; /// If additional static type information can be deduced from a insert_slice's /// size operands, insert an explicit cast of the op's source operand. This /// enables other canonicalization patterns that are matching for tensor_cast /// ops such as `ForOpTensorCastFolder` in SCF. /// /// Example: /// /// ```mlir /// %r = tensor.insert_slice %0 into %1[...] [64, 64] [1, 1] /// : tensor<?x?xf32> into ... /// ``` /// /// folds into: /// /// ```mlir /// %tmp = tensor.cast %0 : tensor<?x?xf32> to tensor<64x64xf32> /// %r = tensor.insert_slice %tmp into %1[...] [64, 64] [1, 1] /// : tensor<64x64xf32> into ... /// ``` /// /// This patterns works with both InsertSliceOp and ParallelInsertSliceOp. template <typename InsertOpTy> struct InsertSliceOpSourceCastInserter final : public OpRewritePattern<InsertOpTy> { … }; } // namespace llvm::SmallBitVector InsertSliceOp::getDroppedDims() { … } void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder &b, Location loc, Value tensor, Value dest) { … } //===----------------------------------------------------------------------===// // PadOp //===----------------------------------------------------------------------===// void PadOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) { … } // TODO: Replace custom<InferType> directive with AllTypesMatch as soon as it // supports optional types. void printInferType(OpAsmPrinter &printer, Operation *op, Value optOperand, Type typeToInfer, Type typeToInferFrom) { … } ParseResult parseInferType(OpAsmParser &parser, std::optional<OpAsmParser::UnresolvedOperand> optOperand, Type &typeToInfer, Type typeToInferFrom) { … } LogicalResult PadOp::verify() { … } LogicalResult PadOp::verifyRegions() { … } RankedTensorType PadOp::inferResultType(RankedTensorType sourceType, ArrayRef<int64_t> staticLow, ArrayRef<int64_t> staticHigh, ArrayRef<int64_t> resultShape) { … } void PadOp::build(OpBuilder &b, OperationState &result, Type resultType, Value source, ArrayRef<int64_t> staticLow, ArrayRef<int64_t> staticHigh, ValueRange low, ValueRange high, bool nofold, ArrayRef<NamedAttribute> attrs) { … } void PadOp::build(OpBuilder &b, OperationState &result, Type resultType, Value source, ValueRange low, ValueRange high, bool nofold, ArrayRef<NamedAttribute> attrs) { … } void PadOp::build(OpBuilder &b, OperationState &result, Type resultType, Value source, ArrayRef<OpFoldResult> low, ArrayRef<OpFoldResult> high, bool nofold, ArrayRef<NamedAttribute> attrs) { … } void PadOp::build(OpBuilder &b, OperationState &result, Type resultType, Value source, ArrayRef<OpFoldResult> low, ArrayRef<OpFoldResult> high, Value constantPadValue, bool nofold, ArrayRef<NamedAttribute> attrs) { … } llvm::SmallBitVector PadOp::getPaddedDims() { … } namespace { // Folds tensor.pad when padding is static zeros and the attribute // doesn't request otherwise. struct FoldStaticZeroPadding : public OpRewritePattern<PadOp> { … }; // Fold CastOp into PadOp when adding static information. struct FoldSourceTensorCast : public OpRewritePattern<PadOp> { … }; // Fold CastOp using the result of PadOp back into the latter if it adds // static information. struct FoldTargetTensorCast : public OpRewritePattern<PadOp> { … }; /// Fold chains of tensor::ExtractSliceOp, tensor::PadOp pairs that pad /// different dimensions. The pattern applies if the following preconditions /// hold: /// 1) the tensor::ExtractSliceOps are not rank-reducing, /// 2) the tensor::ExtractSliceOps have only unit-strides, /// 3) the tensor::PadOps perform only high-padding, /// 4) the tensor::PadOps have the same constant padding value, /// 5) the tensor::PadOps do not have common padding dimensions, /// 6) one tensor::ExtractSliceOp, tensor::PadOp pair has zero-padding and /// zero-offset for every dimension. /// 7) the tensor::ExtractSliceOp sizes match the source tensor sizes for /// the /// padded source dimensions. /// /// Example: /// /// ```mlir /// %0 = tensor.extract_slice %input[16, 0] [%sz0, 64] [1, 1] /// : tensor<64x64xf32> to tensor<?x64xf32> /// %1 = tensor.pad %0 low[0, 0] high[%pw0, 0] { ... /// } : tensor<?x64xf32> to tensor<8x64xf32> /// %2 = tensor.extract_slice %1[0, 4] [8, %sz1] [1, 1] /// : tensor<8x64xf32> to tensor<8x?xf32> /// %res = tensor.pad %2 nofold low[0, 0] high[0, %pw1] { ... /// } : tensor<8x?xf32> to tensor<8x4xf32> /// ``` /// /// folds into: /// /// ```mlir /// %0 = tensor.extract_slice %input[16, 4] [%sz0, %sz1] [1, 1] /// : tensor<64x64xf32> to tensor<?x?xf32> /// %res = tensor.pad %0 nofold low[0, 0] high[%pw0, %pw1] { ... /// } : tensor<?x?xf32> to tensor<8x4xf32> /// ``` struct FoldOrthogonalPaddings : public OpRewritePattern<PadOp> { … }; struct FoldStaticPadding : public OpRewritePattern<PadOp> { … }; } // namespace void PadOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } /// Return the padding value of the PadOp if it constant. In this context, /// "constant" means an actual constant or "defined outside of the block". /// /// Values are considered constant in three cases: /// - A ConstantLike value. /// - A basic block argument from a different block. /// - A value defined outside of the block. /// /// If the padding value is not constant, an empty Value is returned. Value PadOp::getConstantPaddingValue() { … } OpFoldResult PadOp::fold(FoldAdaptor) { … } //===----------------------------------------------------------------------===// // ParallelInsertSliceOp //===----------------------------------------------------------------------===// OpResult ParallelInsertSliceOp::getTiedOpResult() { … } // Build a ParallelInsertSliceOp with mixed static and dynamic entries. void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result, Value source, Value dest, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides, ArrayRef<NamedAttribute> attrs) { … } /// Build an ParallelInsertSliceOp with mixed static and dynamic entries /// packed into a Range vector. void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result, Value source, Value dest, ArrayRef<Range> ranges, ArrayRef<NamedAttribute> attrs) { … } // Build a ParallelInsertSliceOp with dynamic entries. void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result, Value source, Value dest, ValueRange offsets, ValueRange sizes, ValueRange strides, ArrayRef<NamedAttribute> attrs) { … } LogicalResult ParallelInsertSliceOp::verify() { … } void ParallelInsertSliceOp::getCanonicalizationPatterns( RewritePatternSet &results, MLIRContext *context) { … } llvm::SmallBitVector ParallelInsertSliceOp::getDroppedDims() { … } //===----------------------------------------------------------------------===// // ScatterOp //===----------------------------------------------------------------------===// void ScatterOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } LogicalResult ScatterOp::verify() { … } //===----------------------------------------------------------------------===// // SplatOp //===----------------------------------------------------------------------===// void SplatOp::build(OpBuilder &builder, OperationState &result, Value element, Type aggregateType, ValueRange dynamicSizes) { … } void SplatOp::build(OpBuilder &builder, OperationState &result, Value element, ArrayRef<int64_t> staticShape, ValueRange dynamicSizes) { … } void SplatOp::build(OpBuilder &builder, OperationState &result, Value element, ArrayRef<OpFoldResult> sizes) { … } void SplatOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } LogicalResult SplatOp::verify() { … } LogicalResult SplatOp::reifyResultShapes(OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } OpFoldResult SplatOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // PackOp/UnPackOp Common //===----------------------------------------------------------------------===// template <typename OpTy> static LogicalResult reifyResultShapesImpl(OpTy op, OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } template <typename OpTy> static DenseMap<int64_t, OpFoldResult> getDimAndTileMappingImpl(OpTy op) { … } template <typename OpTy> static SmallVector<OpFoldResult> getMixedTilesImpl(OpTy op) { … } template <typename OpTy> static SmallVector<int64_t> getStaticTilesImpl(OpTy op) { … } /// Returns true if `dimsPos` is invalid. It is invalid when: /// a) It contains duplicate. /// b) At least one dimension is out of bound (`dimPos` is >= 0 and < rank). /// c) The number of elements in `dimsPos` is > than `rank`. static bool isInvalidPackingPosSpecification(ArrayRef<int64_t> dimsPos, size_t rank) { … } /// Returns true if the dimension of `sourceShape` is smaller than the dimension /// of the `limitShape`. static bool areAllInBound(ArrayRef<int64_t> sourceShape, ArrayRef<int64_t> limitShape) { … } template <typename OpTy> static LogicalResult commonVerifierPackAndUnPackOp(OpTy packOrUnPack) { … } namespace { /// Subset of PackOp/UnPackOp fields used to compute the result of applying /// various permutations to the op. // TODO: Add linalg.transpose + pack/unpack folding patterns that just reuse // these. These may or may not become true foldings / canonicalizations // depending on how aggressive we want to be in automatically folding // transposes. struct PackOrUnPackTransposeResult { … }; } // namespace template <typename OpTy> static PackOrUnPackTransposeResult commonPermutationOfPackAndUnPackOp(OpTy packOrUnPackOp, ArrayRef<int64_t> innerPermutation, ArrayRef<int64_t> outerPermutation) { … } //===----------------------------------------------------------------------===// // PackOp //===----------------------------------------------------------------------===// void PackOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) { … } void PackOp::build(OpBuilder &builder, OperationState &state, Value source, Value dest, ArrayRef<int64_t> innerDimsPos, ArrayRef<OpFoldResult> innerTiles, std::optional<Value> paddingValue, ArrayRef<int64_t> outerDimsPerm) { … } LogicalResult PackOp::reifyResultShapes(OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } DenseMap<int64_t, OpFoldResult> PackOp::getDimAndTileMapping() { … } SmallVector<OpFoldResult> PackOp::getMixedTiles() { … } SmallVector<int64_t> PackOp::getStaticTiles() { … } bool PackOp::requirePaddingValue(ArrayRef<int64_t> inputShape, ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outputShape, ArrayRef<int64_t> outerDimsPerm, ArrayRef<OpFoldResult> innerTiles) { … } LogicalResult PackOp::verify() { … } /// Converts OpFoldResults to int64_t shape entries, unconditionally mapping all /// Value's to kDynamic, even if they are arith.constant values. static SmallVector<int64_t> asShapeWithAnyValueAsDynamic(ArrayRef<OpFoldResult> ofrs) { … } /// Helper for PackOp::{getResultShape,inferPackedType}. Returns the shape of /// the packed type. Having a shared helper helps implement these two methods in /// a way that ensures that they agree on which dimensions are dynamic. static SmallVector<int64_t> getPackOpResultTypeShape( ArrayRef<int64_t> sourceShape, ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm) { … } SmallVector<OpFoldResult> PackOp::getResultShape( OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> sourceDims, ArrayRef<OpFoldResult> innerTileSizes, ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm) { … } /// Get the expected packed type based on source type, tile factors, position of /// the inner tiles and permutation of the outer tiled loop. RankedTensorType PackOp::inferPackedType(RankedTensorType sourceType, ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm) { … } Value PackOp::createDestinationTensor(OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> innerTileSizes, ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm) { … } PackOp PackOp::createTransposedClone(OpBuilder &b, Location loc, ArrayRef<int64_t> innerPermutation, ArrayRef<int64_t> outerPermutation) { … } /// Returns true if the tiles and the tiled dims are constant. template <typename OpTy> bool areTilesAndTiledDimsAllConstant(OpTy op) { … } Speculation::Speculatability PackOp::getSpeculatability() { … } // Return true if `inner_dims_pos` and `outer_dims_perm` target the same // dimensions for pack and unpack. static bool hasSameInnerOuterAttribute(PackOp packOp, UnPackOp unPackOp) { … } // Return true if pack and unpack have the same tiles. // Same SSA values or same integer constants. static bool haveSameTiles(PackOp packOp, UnPackOp unPackOp) { … } /// Returns true if the pack op does not need a padding value. static bool paddingIsNotNeeded(PackOp op) { … } /// Returns true if the `srcShape` or `destShape` is different from the one in /// `packOp` and populates each with the inferred static shape. static bool inferStaticShape(PackOp packOp, SmallVectorImpl<int64_t> &srcShape, SmallVectorImpl<int64_t> &destShape) { … } LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) { … } template <typename PackOrUnpackOp> static bool isLikePadUnPad(PackOrUnpackOp packOp, RankedTensorType packedTensorType) { … } bool PackOp::isLikePad() { … } OpFoldResult PackOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // UnPackOp //===----------------------------------------------------------------------===// void UnPackOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } LogicalResult UnPackOp::reifyResultShapes(OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } DenseMap<int64_t, OpFoldResult> UnPackOp::getDimAndTileMapping() { … } SmallVector<OpFoldResult> UnPackOp::getMixedTiles() { … } SmallVector<int64_t> UnPackOp::getStaticTiles() { … } LogicalResult UnPackOp::verify() { … } Speculation::Speculatability UnPackOp::getSpeculatability() { … } void UnPackOp::build(OpBuilder &builder, OperationState &state, Value source, Value dest, ArrayRef<int64_t> innerDimsPos, ArrayRef<OpFoldResult> innerTiles, ArrayRef<int64_t> outerDimsPerm) { … } Value UnPackOp::createDestinationTensor(OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> innerTileSizes, ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm) { … } UnPackOp UnPackOp::createTransposedClone(OpBuilder &b, Location loc, Value transposedSource, ArrayRef<int64_t> innerPermutation, ArrayRef<int64_t> outerPermutation) { … } /// Returns true if the `srcShape` or `destShape` is different from the one in /// `op` and populates each with the inferred static shape. static bool inferStaticShape(UnPackOp op, SmallVectorImpl<int64_t> &srcShape, SmallVectorImpl<int64_t> &destShape) { … } LogicalResult UnPackOp::canonicalize(UnPackOp unPackOp, PatternRewriter &rewriter) { … } bool UnPackOp::isLikeUnPad() { … } OpFoldResult UnPackOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // Common Canonicalizers and Folders. //===----------------------------------------------------------------------===// /// Folds a tensor.cast op into a consuming DestinationStyleOpInterface op if /// the `tensor.cast` has source that is more static than the consuming op. /// /// Example: /// ```mlir /// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32> /// %2 = consumer %1 ... : tensor<?x?xf32> ... /// ``` /// /// folds into: /// /// ```mlir /// %2 = consumer %0 ... : tensor<8x16xf32> ... /// ``` /// TODO: Move the pattern to a proper place, so all other DestinationStyleOp /// can add the pattern to their canonicalizers. struct FoldTensorCastProducerOp : public OpInterfaceRewritePattern<DestinationStyleOpInterface> { … }; //===----------------------------------------------------------------------===// // TensorDialect //===----------------------------------------------------------------------===// void TensorDialect::getCanonicalizationPatterns( RewritePatternSet &results) const { … } //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// #define GET_OP_CLASSES #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"
Codebase Browser
llvm