//===- LinalgOps.cpp - Implementation of the linalg operations ------------===// // // 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 // //===----------------------------------------------------------------------===// // // This file implements the Linalg operations. // //===----------------------------------------------------------------------===// #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/AsmParser/AsmParser.h" #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/Math/IR/Math.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/SparseTensor/IR/SparseTensor.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/AffineExprVisitor.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/OpImplementation.h" #include "mlir/IR/OperationSupport.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallSet.h" #include "llvm/ADT/StringSet.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/FormatVariadic.h" #include "llvm/Support/MathExtras.h" #include "llvm/Support/raw_ostream.h" #include <optional> usingnamespacemlir; usingnamespacemlir::linalg; /// Return a `memref.dim` or `tensor.dim` for the shape of `v` at `dim`. static OpFoldResult getDimValue(OpBuilder &builder, Location loc, Value v, int64_t dim) { … } /// Returns a memref.subview or a tensor.extract_slice based on the type of the /// `source`. static Operation *getSlice(OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) { … } //===----------------------------------------------------------------------===// // Helper functions //===----------------------------------------------------------------------===// Value linalg::createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) { … } OpFoldResult linalg::createFoldedDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) { … } //===----------------------------------------------------------------------===// // Support for named Linalg ops defined in ods-gen. //===----------------------------------------------------------------------===// RegionBuilderFn; /// Fills the region of a structured operation using the provided /// `regionBuilder`. The method is used by both named structured ops created by /// ods-gen and by manually defined C++ ops. It is called by both builders and /// parsers and creates a block with arguments corresponding to the elemental /// types of `inputTypes` and `outputTypes`. All output types are asserted to be /// ShapedType. static void fillStructuredOpRegion(OpBuilder &opBuilder, Region ®ion, TypeRange inputTypes, TypeRange outputTypes, ArrayRef<NamedAttribute> attrs, RegionBuilderFn regionBuilder) { … } /// Creates a structured operation given `inputs`, `outputs`, and `attributes`. /// The result types are derived automatically if `resultTensorTypes` is none. /// The body of the operation is filled using `regionBuilder`. All ods-gen /// created structured operations use the method to implement their builders. static void buildStructuredOp(OpBuilder &b, OperationState &state, std::optional<TypeRange> resultTensorTypes, ValueRange inputs, ValueRange outputs, ArrayRef<NamedAttribute> attributes, RegionBuilderFn regionBuilder) { … } /// Common parsing used for both named structured ops created by ods-gen and by /// manually defined C++ ops. Does not handle regions. static ParseResult parseCommonStructuredOpParts(OpAsmParser &parser, OperationState &result, SmallVectorImpl<Type> &inputTypes, SmallVectorImpl<Type> &outputTypes, bool addOperandSegmentSizes = true) { … } static void printCommonStructuredOpParts(OpAsmPrinter &p, ValueRange inputs, ValueRange outputs) { … } //===----------------------------------------------------------------------===// // Specific parsing and printing for named structured ops created by ods-gen. //===----------------------------------------------------------------------===// static ParseResult parseNamedStructuredOpRegion( OpAsmParser &parser, Region ®ion, unsigned numRegionArgs, TypeRange inputTypes, TypeRange outputTypes, ArrayRef<NamedAttribute> attrs, RegionBuilderFn regionBuilder) { … } static ParseResult parseNamedStructuredOpResults(OpAsmParser &parser, SmallVectorImpl<Type> &resultTypes) { … } static ParseResult parseNamedStructuredOp(OpAsmParser &parser, OperationState &result, unsigned numRegionArgs, RegionBuilderFn regionBuilder) { … } static void printNamedStructuredOpResults(OpAsmPrinter &p, TypeRange resultTypes) { … } static void printNamedStructuredOp(OpAsmPrinter &p, Operation *op, ValueRange inputs, ValueRange outputs) { … } //===----------------------------------------------------------------------===// // Region builder helper. // TODO: Move this to a utility library. // The public methods on this class are referenced directly from generated code. // Helper build the unary, binary, and type conversion functions defined by the // DSL. See LinalgNamedStructuredOps.yamlgen.cpp.inc for the code that uses this // class. // // Implementations of the math functions must be polymorphic over numeric types, // internally performing necessary casts. If the function application makes no // sense, then the only recourse is to assert and return nullptr. This can be // extended later if it becomes possible to fail construction of the region. The // invariant should be enforced at a higher level. // // TODO: These helpers are currently type polymorphic over the class of integer // and floating point types, but they will not internally cast within bit // widths of a class (mixed precision such as i8->i32) or across classes // (i.e. mixed float and integer). Many such combinations are ambiguous or need // to be handled with care and work is being considered to extend the op // language to make such cases explicit. In the mean-time, violating this will // fail verification, which is deemed acceptable. //===----------------------------------------------------------------------===// namespace { class RegionBuilderHelper { … }; } // namespace //===----------------------------------------------------------------------===// // CopyOp //===----------------------------------------------------------------------===// namespace { struct EraseSelfCopy : OpRewritePattern<CopyOp> { … }; } // namespace void CopyOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // FillOp //===----------------------------------------------------------------------===// namespace { /// Fold linalg.fill -> tensor.expand/collapse_shape chain. /// /// For such op chains, we can create new linalg.fill ops with the result /// type of the tensor.expand/collapse_shape op. template <typename TensorReshapeOp> struct FoldFillWithTensorReshape : OpRewritePattern<TensorReshapeOp> { … }; /// Fold tensor.pad(linalg.fill) into linalg.fill if the padding value and the /// filling value are the same. struct FoldFillWithPad final : public OpRewritePattern<tensor::PadOp> { … }; /// Fold tensor.insert_slice(tensor.pad(<input>), linalg.fill) into /// tensor.insert_slice(<input>, linalg.fill) if the padding value and the /// filling value are the same. struct FoldInsertPadIntoFill : public OpRewritePattern<tensor::InsertSliceOp> { … }; /// Fold tensor.extract(linalg.fill(<input>)) into <input> struct FoldFillWithTensorExtract : public OpRewritePattern<tensor::ExtractOp> { … }; /// Folds pack(fill) into a single fill op if /// 1. The pack op does not have padding value, or /// 2. The filled value and padding value are the same. static FailureOr<FillOp> foldFillPackIntoFillOp(RewriterBase &rewriter, tensor::PackOp packOp) { … } /// Wrapper pattern that applies foldFillPackIntoFillOp method. struct FoldFillWithPack : public OpRewritePattern<tensor::PackOp> { … }; /// Fold fill with copy. struct FoldFillWithCopy : OpRewritePattern<linalg::CopyOp> { … }; /// Fold fill with transpose. struct FoldFillWithTranspose : OpRewritePattern<linalg::TransposeOp> { … }; /// Fold a concat with all elements being fills of the same value /// into a fill of the concat result shape. struct FoldConcatsOfFill : public OpRewritePattern<tensor::ConcatOp> { … }; } // namespace void FillOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // GenericOp //===----------------------------------------------------------------------===// static void buildGenericRegion( OpBuilder &builder, Location loc, Region ®ion, ValueRange inputs, ValueRange outputs, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild) { … } void GenericOp::getAsmBlockArgumentNames(Region ®ion, OpAsmSetValueNameFn setNameFn) { … } void GenericOp::build( OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes, ValueRange inputs, ValueRange outputs, ArrayAttr indexingMaps, ArrayAttr iteratorTypes, StringAttr doc, StringAttr libraryCall, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, ArrayRef<NamedAttribute> attributes) { … } void GenericOp::build( OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes, ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps, ArrayRef<utils::IteratorType> iteratorTypes, StringRef doc, StringRef libraryCall, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, ArrayRef<NamedAttribute> attributes) { … } void GenericOp::build( OpBuilder &builder, OperationState &result, ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps, ArrayRef<utils::IteratorType> iteratorTypes, StringRef doc, StringRef libraryCall, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, ArrayRef<NamedAttribute> attributes) { … } void GenericOp::build( OpBuilder &builder, OperationState &result, ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps, ArrayRef<utils::IteratorType> iteratorTypes, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, ArrayRef<NamedAttribute> attributes) { … } void GenericOp::build( OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes, ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps, ArrayRef<utils::IteratorType> iteratorTypes, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, ArrayRef<NamedAttribute> attributes) { … } void GenericOp::print(OpAsmPrinter &p) { … } ParseResult GenericOp::parse(OpAsmParser &parser, OperationState &result) { … } static void getGenericEffectsImpl( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects, LinalgOp linalgOp) { … } void GenericOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } static Speculation::Speculatability getGenericSpeculatabilityImpl(LinalgOp linalgOp) { … } Speculation::Speculatability GenericOp::getSpeculatability() { … } LogicalResult GenericOp::verify() { … } namespace { /// Remove any linalg operation (on tensors) that are just copying /// the values from inputs to the results. Requirements are /// 1) All iterator types are parallel /// 2) The body contains just a yield operation with the yielded values being /// the arguments corresponding to the operands. template <typename OpTy> struct EraseIdentityLinalgOp : public OpRewritePattern<OpTy> { … }; } // namespace void GenericOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } LogicalResult GenericOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) { … } //===----------------------------------------------------------------------===// // MapOp //===----------------------------------------------------------------------===// static ParseResult parseDstStyleOp( OpAsmParser &parser, OperationState &result, function_ref<ParseResult(OpAsmParser &, NamedAttrList &)> parseAttrsFn = nullptr) { … } void MapOp::getAsmBlockArgumentNames(Region ®ion, OpAsmSetValueNameFn setNameFn) { … } void MapOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) { … } void MapOp::build( OpBuilder &builder, OperationState &result, ValueRange inputs, Value init, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, ArrayRef<NamedAttribute> attributes) { … } static void addBodyWithPayloadOp(OpAsmParser &parser, OperationState &result, const OperationName &payloadOpName, const NamedAttrList &payloadOpAttrs, ArrayRef<Value> operands, bool initFirst = false) { … } ParseResult MapOp::parse(OpAsmParser &parser, OperationState &result) { … } // Retrieve the operation from the body, if it is the only one (except // yield) and if it gets the same amount of arguments as the body does. // If initFirst flag is enabled, we check that init takes the first position in // operands of payload. static Operation *findPayloadOp(Block *body, bool initFirst = false) { … } void printShortForm(OpAsmPrinter &p, Operation *payloadOp) { … } void MapOp::print(OpAsmPrinter &p) { … } LogicalResult MapOp::verify() { … } SmallVector<utils::IteratorType> MapOp::getIteratorTypesArray() { … } ArrayAttr MapOp::getIndexingMaps() { … } void MapOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } Speculation::Speculatability MapOp::getSpeculatability() { … } //===----------------------------------------------------------------------===// // ReduceOp //===----------------------------------------------------------------------===// void ReduceOp::getAsmBlockArgumentNames(Region ®ion, OpAsmSetValueNameFn setNameFn) { … } void ReduceOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } void ReduceOp::build( OpBuilder &builder, OperationState &result, ValueRange inputs, ValueRange inits, ArrayRef<int64_t> dimensions, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, ArrayRef<NamedAttribute> attributes) { … } SmallVector<utils::IteratorType> ReduceOp::getIteratorTypesArray() { … } ArrayAttr ReduceOp::getIndexingMaps() { … } void ReduceOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } Speculation::Speculatability ReduceOp::getSpeculatability() { … } static ParseResult parseDenseI64ArrayAttr(OpAsmParser &parser, NamedAttrList &attributes, StringRef attributeName) { … } ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) { … } static void printDenseI64ArrayAttr(OpAsmPrinter &p, StringRef attributeName, ArrayRef<int64_t> attributeValue) { … } void ReduceOp::print(OpAsmPrinter &p) { … } LogicalResult ReduceOp::verify() { … } //===----------------------------------------------------------------------===// // TransposeOp //===----------------------------------------------------------------------===// static void buildIdentityRegion(OpBuilder &builder, Location loc, Region ®ion, ValueRange inputs, ValueRange outputs) { … } void TransposeOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &result, Value input, Value init, DenseI64ArrayAttr permutation, ArrayRef<NamedAttribute> attributes) { … } void TransposeOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &result, Value input, Value init, ArrayRef<int64_t> permutation, ArrayRef<NamedAttribute> attributes) { … } ParseResult TransposeOp::parse(OpAsmParser &parser, OperationState &result) { … } void TransposeOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } void TransposeOp::print(OpAsmPrinter &p) { … } LogicalResult TransposeOp::verify() { … } SmallVector<utils::IteratorType> TransposeOp::getIteratorTypesArray() { … } ArrayAttr TransposeOp::getIndexingMaps() { … } void TransposeOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } Speculation::Speculatability TransposeOp::getSpeculatability() { … } LogicalResult TransposeOp::fold(FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &result) { … } /// Fold transpose with transpose. struct FoldTransposeWithTranspose : OpRewritePattern<linalg::TransposeOp> { … }; /// This pattern canonicalize transpose by swapping the order of /// broadcast and transpose: /// transpose(broadcast(input)) -> broadcast(transpose(input)) struct SwapTransposeWithBroadcast : OpRewritePattern<linalg::TransposeOp> { … }; void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // BroadcastOp //===----------------------------------------------------------------------===// void BroadcastOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &result, Value input, Value init, DenseI64ArrayAttr dimensions, ArrayRef<NamedAttribute> attributes) { … } void BroadcastOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &result, Value input, Value init, ArrayRef<int64_t> dimensions, ArrayRef<NamedAttribute> attributes) { … } ParseResult BroadcastOp::parse(OpAsmParser &parser, OperationState &result) { … } void BroadcastOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { … } void BroadcastOp::print(OpAsmPrinter &p) { … } LogicalResult BroadcastOp::verify() { … } SmallVector<utils::IteratorType> BroadcastOp::getIteratorTypesArray() { … } ArrayAttr BroadcastOp::getIndexingMaps() { … } void BroadcastOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } Speculation::Speculatability BroadcastOp::getSpeculatability() { … } void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // YieldOp //===----------------------------------------------------------------------===// void linalg::YieldOp::print(OpAsmPrinter &p) { … } ParseResult YieldOp::parse(OpAsmParser &parser, OperationState &result) { … } // Check the operand number and types must match the element types of the // LinalgOp interface's shaped operands. static LogicalResult verifyYield(linalg::YieldOp op, LinalgOp linalgOp) { … } LogicalResult linalg::YieldOp::verify() { … } //===----------------------------------------------------------------------===// // IndexOp //===----------------------------------------------------------------------===// LogicalResult IndexOp::verify() { … } /////// Operations corresponding to library calls defined with Tablegen //////// #include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.cpp.inc" #define GET_OP_CLASSES #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc" #define GET_OP_CLASSES #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc" AffineMap mlir::linalg::extractOrIdentityMap(std::optional<AffineMap> maybeMap, unsigned rank, MLIRContext *context) { … } SmallVector<AffineExpr, 4> mlir::linalg::makeAffineDimExprs(unsigned num, unsigned &startIdx, MLIRContext *context) { … } SmallVector<AffineExpr, 4> mlir::linalg::concat(ArrayRef<AffineExpr> a, ArrayRef<AffineExpr> b) { … } static LogicalResult appendMangledType(llvm::raw_string_ostream &ss, Type t) { … } std::string mlir::linalg::generateLibraryCallName(Operation *op) { … } //===----------------------------------------------------------------------===// // Canonicalizers and Folders. //===----------------------------------------------------------------------===// namespace { struct EraseDeadLinalgOp : public OpInterfaceRewritePattern<LinalgOp> { … }; /// Fold LinalgOps with `tensor.cast` consumer if the `tensor.cast` has /// result that is more static than the linalg op. struct FoldTensorCastConsumerOp : public OpRewritePattern<tensor::CastOp> { … }; /// For each of the operand in `operands` this function maps the static sizes of /// dimensions to their affine dim expressions. static void populateMap(LinalgOp linalgOp, MutableArrayRef<OpOperand> operands, llvm::DenseMap<AffineExpr, int64_t> &affineExprToSize) { … } /// Creates new operand w.r.t 'opOperand' of `linalgOp` with static sizes /// mapped in `affineExprToSize`. New operands are created in `newOperands` and /// their result types is stored in `resultTypes`. If `opOperand` requires no /// change then `changeNeeded` is false and same operand is added in the /// `newOperands` list. static void createNewOperandWithStaticSizes( Location loc, PatternRewriter &rewriter, OpOperand *opOperand, llvm::DenseMap<AffineExpr, int64_t> &affineExprToSize, LinalgOp linalgOp, SmallVector<Value> &newOperands, SmallVector<Type> &resultTypes, bool &changeNeeded) { … } /// Static shapes for the operands can be inferred if any one of the operands /// have a static shape. This can be done by referring to the affine dim /// expressions for the operand. struct InferStaticShapeOfOperands : public OpInterfaceRewritePattern<LinalgOp> { … }; } // namespace // All named ops canonicalizers and folders are auto-generated in the // .cpp.inc. //===----------------------------------------------------------------------===// // SoftmaxOp //===----------------------------------------------------------------------===// LogicalResult SoftmaxOp::verify() { … } SmallVector<Range> SoftmaxOp::getIterationDomain(OpBuilder &builder) { … } SmallVector<utils::IteratorType> SoftmaxOp::getLoopIteratorTypes() { … } FailureOr<TilingResult> SoftmaxOp::getTiledImplementation(OpBuilder &builder, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes) { … } LogicalResult SoftmaxOp::getResultTilePosition( OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets, SmallVector<OpFoldResult> &resultSizes) { … } // cast(dynamic) -> static. LogicalResult SoftmaxOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) { … } LogicalResult SoftmaxOp::reifyResultShapes(OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { … } void SoftmaxOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } // Helper functions for softmax decomposition. // @{ // Helper function to produce the iterator types (reduction or parallel) and // affine maps for the iterators used in the decomposition of softmax. // This method creates: // If allParallel == true: // - iterator type: {parallel, ..., parallel} // - affine maps: // -- identity with inputRank dimensions. // -- (d0, ..., dN) -> (d0, ..., d_dim-1, d_dim+1, ..., dN), // where N == inputRank. // // If allParallel == false: // - iterator type at dim(i) == parallel for i != \p dim and // dim(dim) == reduction. // - affine map: // -- identity with inputRank dimensions. // -- (d0, ..., dN) -> (d0, ..., d_dim-1, d_dim+1, ..., dN), // where N == inputRank. static std::tuple<SmallVector<utils::IteratorType>, SmallVector<AffineMap>> computeIteratorTypesAndIndexingMaps(OpBuilder &builder, int64_t inputRank, int64_t dim, bool allParallel = false) { … } // Helper function to produce a linalg.generic that computes a reduction on // dimension \p dim with the operation type \p T. template <typename T> static Value reduce(OpBuilder &builder, Location loc, Value input, Value output, int64_t dim) { … } /// Produce a linalg generic that computes the second step of the softmax /// decomposition: res = exp(input - max), where \p max is the max of \p input /// on dimension \p dim. static Value buildSubAndExpOp(OpBuilder &builder, Location loc, Value input, Value max, Value output, int64_t dim) { … } /// Produce a linalg generic that computes the final step of the softmax /// decomposition. /// \returns linalg.generic ins(\p numerator, \p denominator) outs(\p output) { /// yield n / d /// } static Value buildDivOp(OpBuilder &builder, Location loc, Value numerator, Value denominator, Value output, int64_t dim) { … } // @} End helper functions for softmax decomposition. /// Given an N-dimensional tensor x, this method converts /// softmax(x) to the following sequence of operations: /// /// 1. Compute the max of x along dimension d. This results /// in a N-1 dimensional tensor m. /// m = max(x, dim = d) /// /// 2. Subtract a broadcasted m from x and exponentiate. This results in /// a N dimensional tensor z. /// z = exp(x - m) /// /// 3. Compute the sum of z along dimension d. This results in /// a N-1 dimensional tensor l. /// l = sum(z, dim = d) /// /// 4. Divide z and l. This gives the N-dimensional softmax. /// softmax = z / l /// FailureOr<SmallVector<Value>> SoftmaxOp::decomposeOperation(OpBuilder &b) { … } //===----------------------------------------------------------------------===// // WinogradFilterTransformOp //===----------------------------------------------------------------------===// LogicalResult WinogradFilterTransformOp::verify() { … } SmallVector<Range> WinogradFilterTransformOp::getIterationDomain(OpBuilder &builder) { … } SmallVector<utils::IteratorType> WinogradFilterTransformOp::getLoopIteratorTypes() { … } LogicalResult WinogradFilterTransformOp::getResultTilePosition( OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets, SmallVector<OpFoldResult> &resultSizes) { … } /// Implement tiling for winograd_filter_transform /// The input of winograd_filter_transform is (F, KH, KW, C). /// The output of winograd_filter_transform is (alphaH, alphaW, C, F) /// Users can specify the tile sizes of F and C. /// `offsets` are the values for the offsets of F, KH, KW, C for one tile. /// `sizes` are the values for the sizes of F, KH, KW, C for one tile. FailureOr<TilingResult> WinogradFilterTransformOp::getTiledImplementation( OpBuilder &builder, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes) { … } //===----------------------------------------------------------------------===// // WinogradInputTransformOp //===----------------------------------------------------------------------===// LogicalResult WinogradInputTransformOp::verify() { … } SmallVector<Range> WinogradInputTransformOp::getIterationDomain(OpBuilder &builder) { … } SmallVector<utils::IteratorType> WinogradInputTransformOp::getLoopIteratorTypes() { … } LogicalResult WinogradInputTransformOp::getResultTilePosition( OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets, SmallVector<OpFoldResult> &resultSizes) { … } /// Implement tiling for winograd_input_transform /// The input of winograd_input_transform is (N, H, W, C). /// The output of winograd_input_transform is (alphaH, alphaW, tileH, tileW, N, /// C) Users can specify the tile sizes of tileH, tileW, N, and C. `offsets` are /// the values for the offsets of tileH, tileW, N, C for one tile. `sizes` are /// the values for the sizes of tileH, tileW, N, C for one tile. FailureOr<TilingResult> WinogradInputTransformOp::getTiledImplementation(OpBuilder &builder, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes) { … } //===----------------------------------------------------------------------===// // WinogradOutputTransformOp //===----------------------------------------------------------------------===// LogicalResult WinogradOutputTransformOp::verify() { … } SmallVector<Range> WinogradOutputTransformOp::getIterationDomain(OpBuilder &builder) { … } SmallVector<utils::IteratorType> WinogradOutputTransformOp::getLoopIteratorTypes() { … } LogicalResult WinogradOutputTransformOp::getResultTilePosition( OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets, SmallVector<OpFoldResult> &resultSizes) { … } /// Implement tiling for winograd_output_transform /// The input of winograd_output_transform is (alphaH, alphaW, tileH, tileW, N, /// F). The output of winograd_output_transform is (N, H, W, F) Users can /// specify the tile sizes of tileH, tileW, N, and F. `offsets` are the values /// for the offsets of tileH, tileW, N, F for one tile. `sizes` are the values /// for the sizes of tileH, tileW, N, F for one tile. FailureOr<TilingResult> WinogradOutputTransformOp::getTiledImplementation( OpBuilder &builder, ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes) { … } //===----------------------------------------------------------------------===// // LinalgDialect //===----------------------------------------------------------------------===// void LinalgDialect::getCanonicalizationPatterns( RewritePatternSet &results) const { … } Operation *LinalgDialect::materializeConstant(OpBuilder &builder, Attribute value, Type type, Location loc) { … }
Codebase Browser
llvm