diff options
Diffstat (limited to 'mlir/lib')
43 files changed, 1803 insertions, 602 deletions
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp index 0b39d1404249..016e59dcb744 100644 --- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp +++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp @@ -128,34 +128,8 @@ AbstractSparseForwardDataFlowAnalysis::visitOperation(Operation *op) { operandLattices.push_back(operandLattice); } - if (auto call = dyn_cast<CallOpInterface>(op)) { - // If the call operation is to an external function, attempt to infer the - // results from the call arguments. - auto callable = - dyn_cast_if_present<CallableOpInterface>(call.resolveCallable()); - if (!getSolverConfig().isInterprocedural() || - (callable && !callable.getCallableRegion())) { - visitExternalCallImpl(call, operandLattices, resultLattices); - return success(); - } - - // Otherwise, the results of a call operation are determined by the - // callgraph. - const auto *predecessors = getOrCreateFor<PredecessorState>( - getProgramPointAfter(op), getProgramPointAfter(call)); - // If not all return sites are known, then conservatively assume we can't - // reason about the data-flow. - if (!predecessors->allPredecessorsKnown()) { - setAllToEntryStates(resultLattices); - return success(); - } - for (Operation *predecessor : predecessors->getKnownPredecessors()) - for (auto &&[operand, resLattice] : - llvm::zip(predecessor->getOperands(), resultLattices)) - join(resLattice, - *getLatticeElementFor(getProgramPointAfter(op), operand)); - return success(); - } + if (auto call = dyn_cast<CallOpInterface>(op)) + return visitCallOperation(call, operandLattices, resultLattices); // Invoke the operation transfer function. return visitOperationImpl(op, operandLattices, resultLattices); @@ -183,24 +157,8 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) { if (block->isEntryBlock()) { // Check if this block is the entry block of a callable region. auto callable = dyn_cast<CallableOpInterface>(block->getParentOp()); - if (callable && callable.getCallableRegion() == block->getParent()) { - const auto *callsites = getOrCreateFor<PredecessorState>( - getProgramPointBefore(block), getProgramPointAfter(callable)); - // If not all callsites are known, conservatively mark all lattices as - // having reached their pessimistic fixpoints. - if (!callsites->allPredecessorsKnown() || - !getSolverConfig().isInterprocedural()) { - return setAllToEntryStates(argLattices); - } - for (Operation *callsite : callsites->getKnownPredecessors()) { - auto call = cast<CallOpInterface>(callsite); - for (auto it : llvm::zip(call.getArgOperands(), argLattices)) - join(std::get<1>(it), - *getLatticeElementFor(getProgramPointBefore(block), - std::get<0>(it))); - } - return; - } + if (callable && callable.getCallableRegion() == block->getParent()) + return visitCallableOperation(callable, argLattices); // Check if the lattices can be determined from region control flow. if (auto branch = dyn_cast<RegionBranchOpInterface>(block->getParentOp())) { @@ -248,6 +206,59 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) { } } +LogicalResult AbstractSparseForwardDataFlowAnalysis::visitCallOperation( + CallOpInterface call, + ArrayRef<const AbstractSparseLattice *> operandLattices, + ArrayRef<AbstractSparseLattice *> resultLattices) { + // If the call operation is to an external function, attempt to infer the + // results from the call arguments. + auto callable = + dyn_cast_if_present<CallableOpInterface>(call.resolveCallable()); + if (!getSolverConfig().isInterprocedural() || + (callable && !callable.getCallableRegion())) { + visitExternalCallImpl(call, operandLattices, resultLattices); + return success(); + } + + // Otherwise, the results of a call operation are determined by the + // callgraph. + const auto *predecessors = getOrCreateFor<PredecessorState>( + getProgramPointAfter(call), getProgramPointAfter(call)); + // If not all return sites are known, then conservatively assume we can't + // reason about the data-flow. + if (!predecessors->allPredecessorsKnown()) { + setAllToEntryStates(resultLattices); + return success(); + } + for (Operation *predecessor : predecessors->getKnownPredecessors()) + for (auto &&[operand, resLattice] : + llvm::zip(predecessor->getOperands(), resultLattices)) + join(resLattice, + *getLatticeElementFor(getProgramPointAfter(call), operand)); + return success(); +} + +void AbstractSparseForwardDataFlowAnalysis::visitCallableOperation( + CallableOpInterface callable, + ArrayRef<AbstractSparseLattice *> argLattices) { + Block *entryBlock = &callable.getCallableRegion()->front(); + const auto *callsites = getOrCreateFor<PredecessorState>( + getProgramPointBefore(entryBlock), getProgramPointAfter(callable)); + // If not all callsites are known, conservatively mark all lattices as + // having reached their pessimistic fixpoints. + if (!callsites->allPredecessorsKnown() || + !getSolverConfig().isInterprocedural()) { + return setAllToEntryStates(argLattices); + } + for (Operation *callsite : callsites->getKnownPredecessors()) { + auto call = cast<CallOpInterface>(callsite); + for (auto it : llvm::zip(call.getArgOperands(), argLattices)) + join(std::get<1>(it), + *getLatticeElementFor(getProgramPointBefore(entryBlock), + std::get<0>(it))); + } +} + void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors( ProgramPoint *point, RegionBranchOpInterface branch, RegionBranchPoint successor, ArrayRef<AbstractSparseLattice *> lattices) { @@ -512,31 +523,34 @@ AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) { if (op->hasTrait<OpTrait::ReturnLike>()) { // Going backwards, the operands of the return are derived from the // results of all CallOps calling this CallableOp. - if (auto callable = dyn_cast<CallableOpInterface>(op->getParentOp())) { - const PredecessorState *callsites = getOrCreateFor<PredecessorState>( - getProgramPointAfter(op), getProgramPointAfter(callable)); - if (callsites->allPredecessorsKnown()) { - for (Operation *call : callsites->getKnownPredecessors()) { - SmallVector<const AbstractSparseLattice *> callResultLattices = - getLatticeElementsFor(getProgramPointAfter(op), - call->getResults()); - for (auto [op, result] : - llvm::zip(operandLattices, callResultLattices)) - meet(op, *result); - } - } else { - // If we don't know all the callers, we can't know where the - // returned values go. Note that, in particular, this will trigger - // for the return ops of any public functions. - setAllToExitStates(operandLattices); - } - return success(); - } + if (auto callable = dyn_cast<CallableOpInterface>(op->getParentOp())) + return visitCallableOperation(op, callable, operandLattices); } return visitOperationImpl(op, operandLattices, resultLattices); } +LogicalResult AbstractSparseBackwardDataFlowAnalysis::visitCallableOperation( + Operation *op, CallableOpInterface callable, + ArrayRef<AbstractSparseLattice *> operandLattices) { + const PredecessorState *callsites = getOrCreateFor<PredecessorState>( + getProgramPointAfter(op), getProgramPointAfter(callable)); + if (callsites->allPredecessorsKnown()) { + for (Operation *call : callsites->getKnownPredecessors()) { + SmallVector<const AbstractSparseLattice *> callResultLattices = + getLatticeElementsFor(getProgramPointAfter(op), call->getResults()); + for (auto [op, result] : llvm::zip(operandLattices, callResultLattices)) + meet(op, *result); + } + } else { + // If we don't know all the callers, we can't know where the + // returned values go. Note that, in particular, this will trigger + // for the return ops of any public functions. + setAllToExitStates(operandLattices); + } + return success(); +} + void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors( RegionBranchOpInterface branch, ArrayRef<AbstractSparseLattice *> operandLattices) { diff --git a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp index dea3d99704d9..01ca5e99a9af 100644 --- a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp +++ b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp @@ -543,14 +543,20 @@ LogicalResult GPUPrintfOpToVPrintfLowering::matchAndRewrite( // the device code, not the host code auto moduleOp = gpuPrintfOp->getParentOfType<gpu::GPUModuleOp>(); + // Create a valid global location removing any metadata attached to the + // location as debug info metadata inside of a function cannot be used outside + // of that function. + Location globalLoc = loc->findInstanceOfOrUnknown<FileLineColLoc>(); + auto vprintfType = LLVM::LLVMFunctionType::get(rewriter.getI32Type(), {ptrType, ptrType}); - LLVM::LLVMFuncOp vprintfDecl = - getOrDefineFunction(moduleOp, loc, rewriter, "vprintf", vprintfType); + LLVM::LLVMFuncOp vprintfDecl = getOrDefineFunction( + moduleOp, globalLoc, rewriter, "vprintf", vprintfType); // Create the global op or find an existing one. - LLVM::GlobalOp global = getOrCreateStringConstant( - rewriter, loc, moduleOp, llvmI8, "printfFormat_", adaptor.getFormat()); + LLVM::GlobalOp global = + getOrCreateStringConstant(rewriter, globalLoc, moduleOp, llvmI8, + "printfFormat_", adaptor.getFormat()); // Get a pointer to the format string's first element Value globalPtr = rewriter.create<LLVM::AddressOfOp>(loc, global); diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp index 78e6ebb523a4..46db5d3fdca3 100644 --- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp @@ -464,27 +464,39 @@ LogicalResult GPUShuffleConversion::matchAndRewrite( template <typename UniformOp, typename NonUniformOp> static Value createGroupReduceOpImpl(OpBuilder &builder, Location loc, - Value arg, bool isGroup, bool isUniform) { + Value arg, bool isGroup, bool isUniform, + std::optional<uint32_t> clusterSize) { Type type = arg.getType(); auto scope = mlir::spirv::ScopeAttr::get(builder.getContext(), isGroup ? spirv::Scope::Workgroup : spirv::Scope::Subgroup); - auto groupOp = spirv::GroupOperationAttr::get(builder.getContext(), - spirv::GroupOperation::Reduce); + auto groupOp = spirv::GroupOperationAttr::get( + builder.getContext(), clusterSize.has_value() + ? spirv::GroupOperation::ClusteredReduce + : spirv::GroupOperation::Reduce); if (isUniform) { return builder.create<UniformOp>(loc, type, scope, groupOp, arg) .getResult(); } - return builder.create<NonUniformOp>(loc, type, scope, groupOp, arg, Value{}) + + Value clusterSizeValue; + if (clusterSize.has_value()) + clusterSizeValue = builder.create<spirv::ConstantOp>( + loc, builder.getI32Type(), + builder.getIntegerAttr(builder.getI32Type(), *clusterSize)); + + return builder + .create<NonUniformOp>(loc, type, scope, groupOp, arg, clusterSizeValue) .getResult(); } -static std::optional<Value> createGroupReduceOp(OpBuilder &builder, - Location loc, Value arg, - gpu::AllReduceOperation opType, - bool isGroup, bool isUniform) { +static std::optional<Value> +createGroupReduceOp(OpBuilder &builder, Location loc, Value arg, + gpu::AllReduceOperation opType, bool isGroup, + bool isUniform, std::optional<uint32_t> clusterSize) { enum class ElemType { Float, Boolean, Integer }; - using FuncT = Value (*)(OpBuilder &, Location, Value, bool, bool); + using FuncT = Value (*)(OpBuilder &, Location, Value, bool, bool, + std::optional<uint32_t>); struct OpHandler { gpu::AllReduceOperation kind; ElemType elemType; @@ -548,7 +560,7 @@ static std::optional<Value> createGroupReduceOp(OpBuilder &builder, for (const OpHandler &handler : handlers) if (handler.kind == opType && elementType == handler.elemType) - return handler.func(builder, loc, arg, isGroup, isUniform); + return handler.func(builder, loc, arg, isGroup, isUniform, clusterSize); return std::nullopt; } @@ -571,7 +583,7 @@ public: auto result = createGroupReduceOp(rewriter, op.getLoc(), adaptor.getValue(), *opType, - /*isGroup*/ true, op.getUniform()); + /*isGroup*/ true, op.getUniform(), std::nullopt); if (!result) return failure(); @@ -589,16 +601,17 @@ public: LogicalResult matchAndRewrite(gpu::SubgroupReduceOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - if (op.getClusterSize()) + if (op.getClusterStride() > 1) { return rewriter.notifyMatchFailure( - op, "lowering for clustered reduce not implemented"); + op, "lowering for cluster stride > 1 is not implemented"); + } if (!isa<spirv::ScalarType>(adaptor.getValue().getType())) return rewriter.notifyMatchFailure(op, "reduction type is not a scalar"); - auto result = createGroupReduceOp(rewriter, op.getLoc(), adaptor.getValue(), - adaptor.getOp(), - /*isGroup=*/false, adaptor.getUniform()); + auto result = createGroupReduceOp( + rewriter, op.getLoc(), adaptor.getValue(), adaptor.getOp(), + /*isGroup=*/false, adaptor.getUniform(), op.getClusterSize()); if (!result) return failure(); diff --git a/mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt b/mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt index 330474a718e3..8e2620029c35 100644 --- a/mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt +++ b/mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt @@ -35,6 +35,7 @@ add_mlir_conversion_library(MLIRVectorToLLVMPass MLIRVectorToLLVM MLIRArmNeonDialect + MLIRArmNeonTransforms MLIRArmSVEDialect MLIRArmSVETransforms MLIRAMXDialect diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp index 0ee6dce9ee94..293e01a5bf4d 100644 --- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp +++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp @@ -14,6 +14,7 @@ #include "mlir/Dialect/AMX/Transforms.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h" +#include "mlir/Dialect/ArmNeon/Transforms.h" #include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h" #include "mlir/Dialect/ArmSVE/Transforms/Transforms.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -82,6 +83,12 @@ void ConvertVectorToLLVMPass::runOnOperation() { populateVectorStepLoweringPatterns(patterns); populateVectorRankReducingFMAPattern(patterns); populateVectorGatherLoweringPatterns(patterns); + if (armI8MM) { + if (armNeon) + arm_neon::populateLowerContractionToSMMLAPatternPatterns(patterns); + if (armSVE) + populateLowerContractionToSVEI8MMPatternPatterns(patterns); + } (void)applyPatternsGreedily(getOperation(), std::move(patterns)); } diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp index fe53d0324936..01cc50014838 100644 --- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp @@ -402,9 +402,7 @@ isVectorizableLoopBodyWithOpCond(AffineForOp loop, return !VectorType::isValidElementType(type); })) return true; - return llvm::any_of(op.getResultTypes(), [](Type type) { - return !VectorType::isValidElementType(type); - }); + return !llvm::all_of(op.getResultTypes(), VectorType::isValidElementType); }); SmallVector<NestedMatch, 8> opsMatched; types.match(forOp, &opsMatched); diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index 2364f8957992..8a708eb29210 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2667,7 +2667,7 @@ AffineParallelOp mlir::affine::getAffineParallelInductionVarOwner(Value val) { if (!ivArg || !ivArg.getOwner()) return nullptr; Operation *containingOp = ivArg.getOwner()->getParentOp(); - auto parallelOp = dyn_cast<AffineParallelOp>(containingOp); + auto parallelOp = dyn_cast_if_present<AffineParallelOp>(containingOp); if (parallelOp && llvm::is_contained(parallelOp.getIVs(), val)) return parallelOp; return nullptr; diff --git a/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractionToSMMLAPattern.cpp b/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractionToSMMLAPattern.cpp index 2a1271dfd6bd..5ce3d2b28aeb 100644 --- a/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractionToSMMLAPattern.cpp +++ b/mlir/lib/Dialect/ArmNeon/Transforms/LowerContractionToSMMLAPattern.cpp @@ -56,6 +56,10 @@ public: // Avoid 0-D vectors and 1-D rhs: if (!lhsType.hasRank() || !rhsType.hasRank() || rhsType.getRank() < 2) return failure(); + // This codegen does not work for scalable vectors. Return failure so this + // pattern is not accidentally chosen over patterns that lower to ArmSVE. + if (lhsType.isScalable() || rhsType.isScalable()) + return failure(); auto dimM = lhsType.getRank() == 1 ? 1 : lhsType.getDimSize(0); auto dimN = rhsType.getDimSize(0); auto dimK = rhsType.getDimSize(1); @@ -238,5 +242,5 @@ public: void mlir::arm_neon::populateLowerContractionToSMMLAPatternPatterns( RewritePatternSet &patterns) { MLIRContext *context = patterns.getContext(); - patterns.add<LowerContractionToSMMLAPattern>(context, /*benefit=*/1); + patterns.add<LowerContractionToSMMLAPattern>(context, /*benefit=*/2); } diff --git a/mlir/lib/Dialect/ArmSVE/CMakeLists.txt b/mlir/lib/Dialect/ArmSVE/CMakeLists.txt index 9f57627c321f..cb1e9d01821a 100644 --- a/mlir/lib/Dialect/ArmSVE/CMakeLists.txt +++ b/mlir/lib/Dialect/ArmSVE/CMakeLists.txt @@ -1,2 +1,3 @@ add_subdirectory(IR) add_subdirectory(Transforms) +add_subdirectory(TransformOps) diff --git a/mlir/lib/Dialect/ArmSVE/TransformOps/ArmSVEVectorTransformOps.cpp b/mlir/lib/Dialect/ArmSVE/TransformOps/ArmSVEVectorTransformOps.cpp new file mode 100644 index 000000000000..b2ca4fc1eaa8 --- /dev/null +++ b/mlir/lib/Dialect/ArmSVE/TransformOps/ArmSVEVectorTransformOps.cpp @@ -0,0 +1,54 @@ +//===- ArmSVEVectorTransformOps.cpp - Implementation transform ops -------===// +// +// 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/ArmSVE/TransformOps/ArmSVEVectorTransformOps.h" + +#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h" +#include "mlir/Dialect/ArmSVE/Transforms/Transforms.h" +#include "mlir/Dialect/Transform/IR/TransformDialect.h" + +using namespace mlir; + +//===----------------------------------------------------------------------===// +// Apply...PatternsOp +//===----------------------------------------------------------------------===// + +void transform::ApplyArmSVELowerContractionPatternsOp::populatePatterns( + RewritePatternSet &patterns) { + mlir::populateLowerContractionToSVEI8MMPatternPatterns(patterns); +} + +//===----------------------------------------------------------------------===// +// Transform op registration +//===----------------------------------------------------------------------===// + +namespace { +class ArmSVEVectorTransformDialectExtension + : public transform::TransformDialectExtension< + ArmSVEVectorTransformDialectExtension> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( + ArmSVEVectorTransformDialectExtension) + + ArmSVEVectorTransformDialectExtension() { + declareGeneratedDialect<arm_sve::ArmSVEDialect>(); + registerTransformOps< +#define GET_OP_LIST +#include "mlir/Dialect/ArmSVE/TransformOps/ArmSVEVectorTransformOps.cpp.inc" + >(); + } +}; +} // namespace + +#define GET_OP_CLASSES +#include "mlir/Dialect/ArmSVE/TransformOps/ArmSVEVectorTransformOps.cpp.inc" + +void mlir::arm_sve::registerTransformDialectExtension( + DialectRegistry ®istry) { + registry.addExtensions<ArmSVEVectorTransformDialectExtension>(); +} diff --git a/mlir/lib/Dialect/ArmSVE/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/ArmSVE/TransformOps/CMakeLists.txt new file mode 100644 index 000000000000..235fe8f5ca0c --- /dev/null +++ b/mlir/lib/Dialect/ArmSVE/TransformOps/CMakeLists.txt @@ -0,0 +1,18 @@ +add_mlir_dialect_library(MLIRArmSVEVectorTransformOps + ArmSVEVectorTransformOps.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSVE/TransformOps + + DEPENDS + MLIRArmSVEVectorTransformOpsIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRLLVMCommonConversion + MLIRLLVMDialect + MLIRVectorDialect + MLIRTransformDialect + MLIRArmSVEDialect + MLIRArmSVETransforms + ) diff --git a/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt b/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt index a70c489a51fe..65f98b44b1b6 100644 --- a/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt @@ -1,6 +1,7 @@ add_mlir_dialect_library(MLIRArmSVETransforms LegalizeForLLVMExport.cpp LegalizeVectorStorage.cpp + LowerContractionToSVEI8MMPattern.cpp DEPENDS MLIRArmSVEConversionsIncGen diff --git a/mlir/lib/Dialect/ArmSVE/Transforms/LowerContractionToSVEI8MMPattern.cpp b/mlir/lib/Dialect/ArmSVE/Transforms/LowerContractionToSVEI8MMPattern.cpp new file mode 100644 index 000000000000..b1233c5c06eb --- /dev/null +++ b/mlir/lib/Dialect/ArmSVE/Transforms/LowerContractionToSVEI8MMPattern.cpp @@ -0,0 +1,365 @@ +//===- LowerContractionToSMMLAPattern.cpp - Contract to SMMLA ---*- C++ -*-===// +// +// 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 lowering patterns from vector.contract to operations +// that map to instructions from the SVE FEAT_I8MM extension. +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h" +#include "mlir/Dialect/ArmSVE/Transforms/Transforms.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Utils/IndexingUtils.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "mlir/Dialect/UB/IR/UBOps.h" + +#define DEBUG_TYPE "lower-contract-to-arm-sve-i8mm" + +using namespace mlir; + +namespace { +// Get the operand of a `vector.contract`. This function is intended to abstract +// away from the particular way a value is extended before feeding it into the +// `vector.contract` - via zero-extend or an explicit or implicit sign-extend +// (for implicit sign-extension see `vector.contract` documentation). +// +// The template parameter `Op` indicates the extension operation (explicit or +// implicit) for which we are checking. +// +// Return success only for extensions from `i8` to `i32`. +template <typename Op> +std::optional<Value> getExtOperand(Value v, Type i8Ty, Type i32Ty) { + + static_assert(llvm::is_one_of<Op, arith::ExtSIOp, arith::ExtUIOp>::value, + "Must be instantiated with either sign- or zero- extension op"); + + // If the operand is not defined by an explicit extend operation of the + // accepted operation type allow for an implicit sign-extension. + auto extOp = dyn_cast_or_null<Op>(v.getDefiningOp()); + if (!extOp) { + if constexpr (std::is_same<Op, arith::ExtSIOp>::value) { + auto vTy = cast<VectorType>(v.getType()); + if (vTy.getElementType() != i8Ty) + return {}; + return v; + } + return {}; + } + + // If the operand is defined by an explicit extend operation of the accepted + // operation type, check it's extended from `i8` to `i32`. + auto inOp = extOp.getIn(); + auto inTy = dyn_cast<VectorType>(inOp.getType()); + if (!inTy || inTy.getElementType() != i8Ty) + return {}; + + auto outTy = dyn_cast<VectorType>(extOp.getType()); + if (!outTy || outTy.getElementType() != i32Ty) + return {}; + + return inOp; +} + +// Designate the operation (resp. instruction) used to do sub-tile matrix +// multiplications. +enum class MMLA { + Signed, // smmla + Unsigned, // ummla + Mixed, // usmmla + MixedSwapped // usmmla with LHS and RHS swapped +}; + +// Create the matrix mulitply and accumulate operation according to `op`. +Value createMMLA(PatternRewriter &rewriter, MMLA op, Location loc, + mlir::VectorType accType, Value acc, Value lhs, Value rhs) { + switch (op) { + case MMLA::Signed: + return rewriter.create<arm_sve::SmmlaOp>(loc, accType, acc, lhs, rhs); + case MMLA::Unsigned: + return rewriter.create<arm_sve::UmmlaOp>(loc, accType, acc, lhs, rhs); + case MMLA::Mixed: + return rewriter.create<arm_sve::UsmmlaOp>(loc, accType, acc, lhs, rhs); + case MMLA::MixedSwapped: + // The accumulator comes transposed and the result will be transposed + // later, so all we have to do here is swap the operands. + return rewriter.create<arm_sve::UsmmlaOp>(loc, accType, acc, rhs, lhs); + } +} + +/// Lower a contraction operation that performs a matrix multiplication +/// of two 8-bit integer matrix tiles with logical dimensions <Mx8> and <8x[N]> +/// for the left-hand side and the right-hand side, respectively, +/// yielding a <Mx[N]> 32-bit integer result. +/// +/// The operands' shapes are such that the operands can be evenly split into +/// sub-tiles with dimensions as expected by the targeted FEAT_I8MM +/// instructions. The intent is that M and N are chosen (by higher level +/// transforms) in such a way as to maximise register usage. The main use case +/// we envision as of now is MMT4D, thus the RHS operand is expected +/// pre-transposed. +/// +/// The matrix multiplication is performed by unrolling the usual tiled matrix +/// multiplication algorithm using sub-tiles with dimensions <2x8> for the LHS, +/// <8x[2]> for the RHS, and <2x[2]> for the result and the input accumulator. +/// +/// One way to illustrate the operation is as follows: +/// +/// RHS<8x[N]>: <8x[2]> <8x[2]> ... <8x[2]> +/// +----------------------------- +/// LHS<Mx8>: <2x8> | <2x[2]> <2x[2]> ... <2x[2]> +/// <2x8> | <2x[2]> <2x[2]> ... <2x[2]> +/// ... | ... ... ... ... +/// <2x8> | <2x[2]> <2x[2]> ... <2x[2]> +/// +/// The RHS operand is unpacked into N/2 values, each representing a sequence of +/// VSCALE number of sub-tiles with dimensions <8x2>. +/// The LHS operand is initially unpacked into M/2 values, each representing a +/// sub-tile with dimensions <2x8>, and then each such sub-tile is replicated +/// VSCALE times. +/// Multiplying thus replicated LHS sub-tile by the corresponding RHS sub-tile +/// correctly computes an entire result sub-tile. +class LowerContractionToSVEI8MMPattern + : public OpRewritePattern<vector::ContractionOp> { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(vector::ContractionOp op, + PatternRewriter &rewriter) const override { + + Location loc = op.getLoc(); + mlir::VectorType lhsType = op.getLhsType(); + mlir::VectorType rhsType = op.getRhsType(); + + // Check the rank the types so we can safely examine their dimensions. + if (lhsType.getRank() != 2 || rhsType.getRank() != 2) + return rewriter.notifyMatchFailure(op, "non-matching operand shape"); + + auto M = lhsType.getDimSize(0); + auto N = rhsType.getDimSize(0); + auto K = rhsType.getDimSize(1); + + // Check the operands have the expected shape: + // * for LHS: fixed vector MxK + // * for RHS: scalable vector [N]xK + // * K == 8 + // * M and N even and at least 2 + if (lhsType.isScalable() || !rhsType.getScalableDims()[0] || + rhsType.getScalableDims()[1] || lhsType.getDimSize(1) != K || K != 8 || + M < 2 || M % 2 != 0 || N < 2 || N % 2 != 0 || + !rhsType.getScalableDims()[0]) + return rewriter.notifyMatchFailure(op, "non-matching operand shape"); + + // Check permutation maps. For now only accept + // lhs: (d0, d1, d2) -> (d0, d2) + // rhs: (d0, d1, d2) -> (d1, d2) + // acc: (d0, d1, d2) -> (d0, d1) + // This corresponds to matrix multiplication with transposed RHS. + if (op.getIndexingMapsArray()[0] != + AffineMap::getMultiDimMapWithTargets(3, ArrayRef{0u, 2u}, + op.getContext()) || + op.getIndexingMapsArray()[1] != + AffineMap::getMultiDimMapWithTargets(3, ArrayRef{1u, 2u}, + op.getContext()) || + op.getIndexingMapsArray()[2] != + AffineMap::getMultiDimMapWithTargets(3, ArrayRef{0u, 1u}, + op.getContext())) + return rewriter.notifyMatchFailure(op, "non-matching permutation maps"); + + // Check iterator types for matrix multiplication. + auto itTypes = op.getIteratorTypesArray(); + if (itTypes.size() != 3 || itTypes[0] != vector::IteratorType::parallel || + itTypes[1] != vector::IteratorType::parallel || + itTypes[2] != vector::IteratorType::reduction) + return rewriter.notifyMatchFailure( + op, "iterator types do not correspond to matrix multiplication"); + + // Check the combining kind is addition. + if (op.getKind() != vector::CombiningKind::ADD) + return rewriter.notifyMatchFailure(op, + "combining kind is not an addition"); + + // Check the output is a vector of i32 elements. + auto outTy = dyn_cast<VectorType>(op.getResultType()); + if (!outTy || outTy.getElementType() != rewriter.getI32Type()) + return rewriter.notifyMatchFailure(op, + "output type is not a vector of i32"); + + // Check inputs are sign-/zero- extensions from i8 to i32. Get the values + // before the extension. All four signed/unsigned combinations for input + // operands are supported, but they are lowered to different operations. + // Determine which is the appropriate operation to lower to. + MMLA mmlaOp = MMLA::Signed; + auto maybeLhs = getExtOperand<arith::ExtSIOp>( + op.getLhs(), rewriter.getI8Type(), rewriter.getI32Type()); + if (!maybeLhs) { + mmlaOp = MMLA::Unsigned; + maybeLhs = getExtOperand<arith::ExtUIOp>( + op.getLhs(), rewriter.getI8Type(), rewriter.getI32Type()); + } + if (!maybeLhs) + return rewriter.notifyMatchFailure( + op, "LHS is not a sign- or zero- extended i8"); + + auto maybeRhs = getExtOperand<arith::ExtSIOp>( + op.getRhs(), rewriter.getI8Type(), rewriter.getI32Type()); + if (maybeRhs) { + if (mmlaOp == MMLA::Unsigned) + mmlaOp = MMLA::Mixed; + } else { + if (mmlaOp == MMLA::Signed) + mmlaOp = MMLA::MixedSwapped; + maybeRhs = getExtOperand<arith::ExtUIOp>( + op.getRhs(), rewriter.getI8Type(), rewriter.getI32Type()); + } + if (!maybeRhs) + return rewriter.notifyMatchFailure( + op, "RHS is not a sign- or zero- extended i8"); + + // One-dimensional vector types for arm_sve.*mmla + auto nxv16i8 = VectorType::get(/*shape=*/16, rewriter.getI8Type(), + /*scalableDims=*/{true}); + auto nxv4i32 = VectorType::get(/*shape=*/4, rewriter.getI32Type(), + /*scalableDims=*/{true}); + + // Extract LHS sub-tiles with logicall shape <2x8>. + SmallVector<Value> lhsTile; + for (int64_t i = 0; i < M; i += 2) { + // Extract two consecutive rows of the LHS tile. + auto r0 = rewriter.create<vector::ExtractOp>(loc, *maybeLhs, + ArrayRef<int64_t>{i}); + auto r1 = rewriter.create<vector::ExtractOp>(loc, *maybeLhs, + ArrayRef<int64_t>{i + 1}); + // Concatenate to obtain a 16 x i8 flattened sub-tile. + auto t = rewriter.create<vector::ShuffleOp>( + loc, r0, r1, + llvm::ArrayRef<int64_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, + 14, 15}); + // Turn it into a scalable vector. + auto s = rewriter.create<vector::ScalableInsertOp>( + loc, t, rewriter.create<ub::PoisonOp>(loc, nxv16i8), 0); + // Replicate the sub-tile VSCALE times to fill the entire vector. + auto r = rewriter.create<arm_sve::DupQLaneOp>(loc, s, 0); + lhsTile.push_back(r); + } + + // "Flatten" the RHS tile from <[N]x8> to <[8*N]>. + auto rhs = rewriter.create<vector::ShapeCastOp>( + maybeRhs->getLoc(), + VectorType::get(/*shape=*/8 * N, rewriter.getI8Type(), + /*scalableDims=*/{true}), + *maybeRhs); + + // Extract the RHS sub-tiles with logical shape <8x[2]>. + SmallVector<Value> rhsTile; + for (int64_t j = 0; j < N; j += 2) + rhsTile.push_back( + rewriter.create<vector::ScalableExtractOp>(loc, nxv16i8, rhs, j * 8)); + + // Handy types for packing/unpacking of the accumulator tile. + auto accRowTy = VectorType::get(/*shape=*/N, rewriter.getI32Type(), + /*scalableDims=*/{true}); + auto accRowX2Ty = VectorType::get(/*shape=*/2 * N, rewriter.getI32Type(), + /*scalableDims=*/{true}); + auto accRow64Ty = VectorType::get(/*shape=*/N / 2, rewriter.getI64Type(), + /*scalableDims=*/{true}); + auto accRowX264Ty = VectorType::get(/*shape=*/N, rewriter.getI64Type(), + /*scalableDims=*/{true}); + + // Extract and pack the ACC sub-tiles. + SmallVector<Value> accTile; + for (int64_t i = 0; i < M; i += 2) { + // Extract two consecutive rows of the accumulator tile. + auto r0 = rewriter.create<vector::ExtractOp>(loc, op.getAcc(), + ArrayRef<int64_t>{i}); + auto r1 = rewriter.create<vector::ExtractOp>(loc, op.getAcc(), + ArrayRef<int64_t>{i + 1}); + Value accTileVec; + if (mmlaOp == MMLA::MixedSwapped) { + // We need to swap the positions of the LHS and RHS (since we don't have + // a signed * unsigned operation), but then each individual 2x2 tile of + // the acumulator and (later) the result need to be transposed. + accTileVec = rewriter.create<vector::InterleaveOp>(loc, r0, r1); + } else { + // Bitcast them to 64-bit elements, so subsequent + // interleave/deinterleave work on pairs of 32-bit numbers. + auto r0I64 = rewriter.create<vector::BitCastOp>(loc, accRow64Ty, r0); + auto r1I64 = rewriter.create<vector::BitCastOp>(loc, accRow64Ty, r1); + + // Interleave the rows, effectively flattening each 2x2 tile into 4 + // consecutive elements. + auto intrI64 = rewriter.create<vector::InterleaveOp>(loc, r0I64, r1I64); + + // Bitcast back to 32-bit elements. + accTileVec = + rewriter.create<vector::BitCastOp>(loc, accRowX2Ty, intrI64); + } + // Extract ACC sub-tiles. + for (int64_t j = 0; j < N; j += 2) + accTile.push_back(rewriter.create<vector::ScalableExtractOp>( + loc, nxv4i32, accTileVec, j * 2)); + } + + // Emit sub-tile matrix multiplications. + SmallVector<Value> outTile; + for (int64_t i = 0; i < M / 2; ++i) + for (int64_t j = 0; j < N / 2; ++j) { + Value mmla = createMMLA(rewriter, mmlaOp, loc, nxv4i32, + accTile[i * N / 2 + j], lhsTile[i], rhsTile[j]); + outTile.push_back(mmla); + } + + // Unpack the OUT sub-tiles and insert into the result. + Value result = rewriter.create<ub::PoisonOp>(loc, op.getResultType()); + for (int64_t i = 0; i < M / 2; ++i) { + // Collect a number of sub-tiles in a row. + Value row = rewriter.create<ub::PoisonOp>(loc, accRowX2Ty); + for (int64_t j = 0; j < N / 2; ++j) + row = rewriter.create<vector::ScalableInsertOp>( + loc, outTile[i * N / 2 + j], row, j * 4); + + // Unpack the row to obtain two rows of the output. If we have the out + // sub-tiles transposed we obtain two consecutive output rows by + // separating even and odd elements, i.e. a simple deinterleave. + // Otherwise, the interleave is by pairs. + Value out0, out1; + if (mmlaOp == MMLA::MixedSwapped) { + auto tmp = rewriter.create<vector::DeinterleaveOp>(loc, row); + out0 = tmp.getRes1(); + out1 = tmp.getRes2(); + } else { + // Deinterleave by pairs. + auto row64 = rewriter.create<vector::BitCastOp>(loc, accRowX264Ty, row); + auto deintr64 = rewriter.create<vector::DeinterleaveOp>(loc, row64); + + // Bitcast back into 32-bit elements and insert into the result. + out0 = rewriter.create<vector::BitCastOp>(loc, accRowTy, + deintr64.getRes1()); + out1 = rewriter.create<vector::BitCastOp>(loc, accRowTy, + deintr64.getRes2()); + } + result = rewriter.create<vector::InsertOp>(loc, out0, result, i * 2); + result = rewriter.create<vector::InsertOp>(loc, out1, result, i * 2 + 1); + } + + rewriter.replaceOp(op, result); + return success(); + } +}; + +} // namespace + +void mlir::populateLowerContractionToSVEI8MMPatternPatterns( + RewritePatternSet &patterns) { + MLIRContext *context = patterns.getContext(); + patterns.add<LowerContractionToSVEI8MMPattern>(context, /*benefit=*/2); +} diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp index eed7a56fff8a..ca914df8b789 100644 --- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp @@ -94,7 +94,9 @@ void Ownership::combine(Ownership other) { *this = getCombined(other); } // DeallocationState //===----------------------------------------------------------------------===// -DeallocationState::DeallocationState(Operation *op) : liveness(op) {} +DeallocationState::DeallocationState(Operation *op, + SymbolTableCollection &symbolTables) + : symbolTable(symbolTables), liveness(op) {} void DeallocationState::updateOwnership(Value memref, Ownership ownership, Block *block) { diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp index fc6424a25ac7..d1d106220a38 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp @@ -341,15 +341,25 @@ static LogicalResult getFuncOpsOrderedByCalls( // Iteratively remove function operations that do not call any of the // functions remaining in the callCounter map and add them to ordered list. - while (!numberCallOpsContainedInFuncOp.empty()) { - auto it = llvm::find_if(numberCallOpsContainedInFuncOp, - [](auto entry) { return entry.getSecond() == 0; }); - if (it == numberCallOpsContainedInFuncOp.end()) - break; - orderedFuncOps.push_back(it->getFirst()); - for (auto callee : calledBy[it->getFirst()]) - numberCallOpsContainedInFuncOp[callee]--; - numberCallOpsContainedInFuncOp.erase(it); + SmallVector<func::FuncOp> worklist; + + for (const auto &entry : numberCallOpsContainedInFuncOp) { + if (entry.second == 0) + worklist.push_back(entry.first); + } + + while (!worklist.empty()) { + func::FuncOp func = worklist.pop_back_val(); + orderedFuncOps.push_back(func); + + for (func::FuncOp caller : calledBy[func]) { + auto &count = numberCallOpsContainedInFuncOp[caller]; + + if (--count == 0) + worklist.push_back(caller); + } + + numberCallOpsContainedInFuncOp.erase(func); } // Put all other functions in the list of remaining functions. These are diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp index c5b2f3020712..1eeafc4df8cf 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp @@ -166,8 +166,9 @@ namespace { /// program have a corresponding de-allocation. class BufferDeallocation { public: - BufferDeallocation(Operation *op, DeallocationOptions options) - : state(op), options(options) {} + BufferDeallocation(Operation *op, DeallocationOptions options, + SymbolTableCollection &symbolTables) + : state(op, symbolTables), options(options) {} /// Performs the actual placement/creation of all dealloc operations. LogicalResult deallocate(FunctionOpInterface op); @@ -1027,11 +1028,13 @@ struct OwnershipBasedBufferDeallocationPass DeallocationOptions options; options.privateFuncDynamicOwnership = privateFuncDynamicOwnership; + mlir::SymbolTableCollection symbolTables; + auto status = getOperation()->walk([&](func::FuncOp func) { if (func.isExternal()) return WalkResult::skip(); - if (failed(deallocateBuffersOwnershipBased(func, options))) + if (failed(deallocateBuffersOwnershipBased(func, options, symbolTables))) return WalkResult::interrupt(); return WalkResult::advance(); @@ -1047,11 +1050,11 @@ struct OwnershipBasedBufferDeallocationPass // Implement bufferization API //===----------------------------------------------------------------------===// -LogicalResult -bufferization::deallocateBuffersOwnershipBased(FunctionOpInterface op, - DeallocationOptions options) { +LogicalResult bufferization::deallocateBuffersOwnershipBased( + FunctionOpInterface op, DeallocationOptions options, + SymbolTableCollection &symbolTables) { // Gather all required allocation nodes and prepare the deallocation phase. - BufferDeallocation deallocation(op, options); + BufferDeallocation deallocation(op, options, symbolTables); // Place all required temporary clone and dealloc nodes. return deallocation.deallocate(op); diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index 648b6b087e59..aaf6b0593c2e 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -1205,6 +1205,42 @@ LogicalResult NVVM::VoteSyncOp::verify() { return success(); } +LogicalResult NVVM::PrefetchOp::verify() { + using MemSpace = NVVM::NVVMMemorySpace; + using CacheLevel = NVVM::PrefetchCacheLevel; + + unsigned addressSpace = + llvm::cast<LLVM::LLVMPointerType>(getAddr().getType()).getAddressSpace(); + std::optional<NVVM::CacheEvictionPriority> evictPriority = getEvictPriority(); + + if (getUniform()) { + if (getCacheLevel() != CacheLevel::L1) + return emitOpError("unsupported cache level, the only supported uniform " + "cache level is L1"); + + if (addressSpace != MemSpace::kGenericMemorySpace) + return emitOpError( + "prefetch to uniform cache requires a generic pointer"); + } + + if (evictPriority) { + if (getCacheLevel() != CacheLevel::L2) + return emitOpError( + "cache eviction priority supported only for cache level L2"); + + if (addressSpace != MemSpace::kGlobalMemorySpace) + return emitOpError("cache eviction priority requires a global pointer"); + + if (*evictPriority != NVVM::CacheEvictionPriority::EvictNormal && + *evictPriority != NVVM::CacheEvictionPriority::EvictLast) + return emitOpError( + "unsupported cache eviction priority, only evict_last and " + "evict_normal are supported"); + } + + return success(); +} + /// Packs the given `field` into the `result`. /// The `result` is 64-bits and each `field` can be 32-bits or narrower. static llvm::Value * @@ -1712,6 +1748,70 @@ NVVM::IDArgPair DotAccumulate4WayOp::getIntrinsicIDAndArgs( return {ids[type], args}; } +NVVM::IDArgPair DotAccumulate2WayOp::getIntrinsicIDAndArgs( + Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) { + auto curOp = cast<NVVM::DotAccumulate2WayOp>(op); + + llvm::SmallVector<llvm::Value *> args; + args.push_back(getAsPackedI32(mt.lookupValue(curOp.getA()), builder)); + args.push_back(getAsPackedI32(mt.lookupValue(curOp.getB()), builder)); + args.push_back(builder.getInt1(curOp.getBHi())); + args.push_back(mt.lookupValue(curOp.getC())); + + bool isASigned = curOp.getAType() == NVVM::DotAccumulateType::SIGNED; + bool isBSigned = curOp.getBType() == NVVM::DotAccumulateType::SIGNED; + unsigned type = (isASigned << 1) | isBSigned; + const llvm::Intrinsic::ID ids[] = { + llvm::Intrinsic::nvvm_idp2a_u_u, + llvm::Intrinsic::nvvm_idp2a_u_s, + llvm::Intrinsic::nvvm_idp2a_s_u, + llvm::Intrinsic::nvvm_idp2a_s_s, + }; + return {ids[type], args}; +} + +llvm::Intrinsic::ID PrefetchOp::getIntrinsicID(NVVM::PrefetchOp &op) { + using MemSpace = NVVM::NVVMMemorySpace; + using CacheLevel = NVVM::PrefetchCacheLevel; + + NVVM::PrefetchCacheLevel cacheLevel = op.getCacheLevel(); + std::optional<NVVM::CacheEvictionPriority> evictPriority = + op.getEvictPriority(); + unsigned addressSpace = + llvm::cast<LLVM::LLVMPointerType>(op.getAddr().getType()) + .getAddressSpace(); + + if (op.getUniform() && cacheLevel == CacheLevel::L1) + return llvm::Intrinsic::nvvm_prefetchu_L1; + + if (evictPriority && cacheLevel == CacheLevel::L2) { + switch (*evictPriority) { + case NVVM::CacheEvictionPriority::EvictLast: + return llvm::Intrinsic::nvvm_prefetch_global_L2_evict_last; + case NVVM::CacheEvictionPriority::EvictNormal: + return llvm::Intrinsic::nvvm_prefetch_global_L2_evict_normal; + default: + llvm_unreachable("Invalid cache eviction priority"); + } + } + + switch (addressSpace) { + case MemSpace::kGenericMemorySpace: + return cacheLevel == CacheLevel::L1 ? llvm::Intrinsic::nvvm_prefetch_L1 + : llvm::Intrinsic::nvvm_prefetch_L2; + case MemSpace::kGlobalMemorySpace: + return cacheLevel == CacheLevel::L1 + ? llvm::Intrinsic::nvvm_prefetch_global_L1 + : llvm::Intrinsic::nvvm_prefetch_global_L2; + case MemSpace::kLocalMemorySpace: + return cacheLevel == CacheLevel::L1 + ? llvm::Intrinsic::nvvm_prefetch_local_L1 + : llvm::Intrinsic::nvvm_prefetch_local_L2; + default: + llvm_unreachable("Invalid pointer address space"); + } +} + //===----------------------------------------------------------------------===// // NVVMDialect initialization, type parsing, and registration. //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp index dd8ef9608a82..7188987e5e93 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp @@ -876,11 +876,8 @@ public: return failure(); } // Currently only support static inner tile sizes. - if (llvm::any_of(packOp.getStaticTiles(), [](int64_t size) { - return ShapedType::isDynamic(size); - })) { + if (llvm::any_of(packOp.getStaticTiles(), ShapedType::isDynamic)) return failure(); - } // User controlled propagation function. if (!controlFn(&packOp.getSourceMutable())) @@ -1002,11 +999,8 @@ public: return failure(); } // Currently only support static inner tile sizes. - if (llvm::any_of(unPackOp.getStaticTiles(), [](int64_t size) { - return ShapedType::isDynamic(size); - })) { + if (llvm::any_of(unPackOp.getStaticTiles(), ShapedType::isDynamic)) return failure(); - } Operation *consumerOp = *result.user_begin(); return TypeSwitch<Operation *, LogicalResult>(consumerOp) diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp index 8718c57b9e86..615d1f66414b 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -227,8 +227,7 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter, // 1. Filter out NYI cases. auto packedTensorType = cast<RankedTensorType>(packOp->getResultTypes().front()); - if (llvm::any_of(packOp.getStaticInnerTiles(), - [](int64_t size) { return ShapedType::isDynamic(size); })) { + if (llvm::any_of(packOp.getStaticInnerTiles(), ShapedType::isDynamic)) { return rewriter.notifyMatchFailure( packOp, "non-static shape NYI, needs a more powerful tensor.expand_shape op"); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp index c5b62227777a..ff28bd7c4834 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -820,9 +820,8 @@ tensorExtractVectorizationPrecondition(Operation *op, bool vectorizeNDExtract) { return failure(); } - if (llvm::any_of(extractOp->getResultTypes(), [](Type type) { - return !VectorType::isValidElementType(type); - })) { + if (!llvm::all_of(extractOp->getResultTypes(), + VectorType::isValidElementType)) { return failure(); } @@ -1506,101 +1505,192 @@ static SmallVector<int64_t> getTiledPackShape(linalg::PackOp packOp, return applyPermutation(destShape, linalg::getPackInverseDestPerm(packOp)); } +/// Determines whether a mask for xfer_write is trivially "all true" +/// +/// Given all the inputs required to generate a mask (mask sizes and shapes), +/// and an xfer_write operation (write indices and the destination tensor +/// shape), determines whether the corresponding mask would be trivially +/// foldable (i.e., trivially "all true"). +/// +/// Use this method to avoid generating spurious masks and relaying on +/// vectorization post-processing to remove them. +/// +/// Pre-conditions for a mask to be trivially foldable: +/// * All involved shapes (mask + destination tensor) are static. +/// * All write indices are constant. +/// * All mask sizes are constant (including `arith.constant`). +/// +/// If the pre-conditions are met, the method checks for each destination +/// dimension `d`: +/// (1) destDimSize[rankDiff + d] <= maskShape[d] +/// (2) destDimSize[rankDiff + d] <= writeIndex[d] + maskSize[d] +/// +/// rankDiff = rank(dest) - rank(mask). +/// +/// This method takes a conservative view: it may return false even if the mask +/// is technically foldable. +/// +/// EXAMPLE 1 (trivially foldable, all shapes match, mask sizes match the shape +/// of the dest tensor): +/// %c0 = arith.constant 0 : index +/// %mask = vector.create_mask 5, 1 +/// vector.mask %mask { +/// vector.transfer_write %vecToStore_1, %dest{[%c0, %c0] +/// {in_bounds = [true, true]} +/// : vector<5x1xi32>, tensor<5x1xi32> +/// } +/// +/// EXAMPLE 2 (not trivially foldable - vector shape exceeds the tensor shape, +/// mask is required to avoid out-of-bounds write): +/// %c0 = arith.constant 0 : index +/// %mask = vector.create_mask 5, 1 +/// vector.mask %mask { +/// vector.transfer_write %vecToStore_2, %dest[%c0, %c0] +/// {in_bounds = [true, true]} +/// : vector<8x1xi32>, tensor<5x1xi32> +/// } +/// +/// TODO: Re-use in createReadOrMaskedRead +static bool isMaskTriviallyFoldable(SmallVector<OpFoldResult> &maskSizes, + SmallVector<Value> &writeIdxs, + ArrayRef<int64_t> destShape, + ArrayRef<int64_t> maskShape) { + // Masking is unavoidable in the case of dynamic tensors. + if (ShapedType::isDynamicShape(destShape)) + return false; + + // Collect all constant mask sizes. + SmallVector<int64_t, 4> cstMaskSizes; + for (auto [i, dimSize] : llvm::enumerate(maskSizes)) { + if (auto intSize = getConstantIntValue(dimSize)) { + cstMaskSizes.push_back(*intSize); + } + } + + // If any of the mask sizes is non-constant, bail out. + if (cstMaskSizes.size() != maskShape.size()) + return false; + + // Collect all constant write indices. + SmallVector<int64_t, 4> cstWriteIdxs; + for (auto [i, idx] : llvm::enumerate(writeIdxs)) { + APSInt intVal; + if (matchPattern(idx, m_ConstantInt(&intVal))) { + cstWriteIdxs.push_back(intVal.getSExtValue()); + } + } + + // If any of the write indices is non-constant, bail out. + if (cstWriteIdxs.size() != destShape.size()) + return false; + + // Go over all destination dims and check (1) and (2). Take into account that: + // * The number of mask sizes will match the rank of the vector to store. + // This could be lower than the rank of the destination tensor. + // * Mask sizes could be larger than the corresponding mask shape (hence + // `clamp`). + // TODO: The 2nd item should be rejected by the verifier. + int64_t rankDiff = destShape.size() - cstMaskSizes.size(); + for (auto [i, idx] : llvm::enumerate(cstMaskSizes)) { + if (/*(1)*/ maskShape[i] > destShape[rankDiff + i] || + /*(2)*/ destShape[rankDiff + i] < + (std::clamp(cstMaskSizes[i], int64_t(0), maskShape[i]) + + cstWriteIdxs[i])) + return false; + } + + return true; +} + /// Creates an optionally masked TransferWriteOp /// /// Generates the following operation: -/// %res = vector.transfer_write %vectorToStore into %dest +/// %res = vector.transfer_write %vecToStore into %dest /// -/// If the leading N dimensions of the destination tensor do not match -/// `inputVecSizesForLeadingDims` (N = rank(inputVecSizesForLeadingDims)), -/// masking is applied to ensure correctness: +/// If shape(vecToStore) != shape(dest), masking is used to ensure correctness: /// -/// %mask = vector.create_mask(%destShape) +/// %mask = vector.create_mask(%destShape) : %vecToStoreShape /// %res = vector.mask %mask { -/// vector.transfer_write %vectorToStore into %dest +/// vector.transfer_write %vecToStore into %dest /// } /// +/// The mask shape is identical to `vecToStore` (with the element type == +/// i1), and the mask values are based on the shape of the `dest` tensor. +/// /// If `useInBoundsInsteadOfMasking` is set to `true`, the `in_bounds` attribute /// is used instead of masking: /// -/// %write = vector.transfer_write %vectorToStore into %dest +/// %write = vector.transfer_write %vecToStore into %dest /// in_bounds_flags = (...) /// %res = vector.transfer_write %input into %dest /// {in_bounds = in_bounds_flags} /// -/// NOTE: All write offsets are set to 0. -/// TODO: Allow specyfying write offsets. -/// NOTE: When N < rank(input), the missing vector sizes are effectively -/// extracted from the trailing sizes of `destSizes`. This means those sizes -/// must be static. -/// TODO: Support cases where an arbitrary dim is dynamic - this will require -/// specifying all the vector sizes. +/// Finally, `writeIndices` specifies the offsets to use. If empty, all indices +/// are set to 0. static Operation * -createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value vectorToStore, - Value dest, - ArrayRef<int64_t> inputVecSizesForLeadingDims, +createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value vecToStore, + Value dest, SmallVector<Value> writeIndices = {}, bool useInBoundsInsteadOfMasking = false) { ShapedType destType = cast<ShapedType>(dest.getType()); - assert(cast<VectorType>(vectorToStore.getType()).getRank() == - static_cast<int64_t>(destType.getRank()) && - "Rank mismatch!"); - (void)destType; + int64_t destRank = destType.getRank(); + auto destShape = destType.getShape(); - int64_t rank = cast<ShapedType>(dest.getType()).getRank(); - auto destShape = cast<ShapedType>(dest.getType()).getShape(); + VectorType vecToStoreType = cast<VectorType>(vecToStore.getType()); + int64_t vecToStoreRank = vecToStoreType.getRank(); + auto vecToStoreShape = vecToStoreType.getShape(); // Compute the in_bounds attribute - SmallVector<bool> inBoundsVal(rank, true); + SmallVector<bool> inBoundsVal(vecToStoreRank, true); if (useInBoundsInsteadOfMasking) { - // In this case, assume that all the required vector sizes have been - // provided. - assert(inputVecSizesForLeadingDims.size() == - static_cast<size_t>(destType.getRank()) && - "Insufficient number of input vector sizes!"); // Update the inBounds attribute. - for (unsigned i = 0; i < rank; i++) - inBoundsVal[i] = (destShape[i] == inputVecSizesForLeadingDims[i]) && - !ShapedType::isDynamic(destShape[i]); + // FIXME: This computation is too weak - it ignores the write indices. + for (unsigned i = 0; i < vecToStoreRank; i++) + inBoundsVal[i] = + (destShape[destRank - vecToStoreRank + i] >= vecToStoreShape[i]) && + !ShapedType::isDynamic(destShape[destRank - vecToStoreRank + i]); + } + + // If missing, initialize the write indices to 0. + assert(writeIndices.empty() || + writeIndices.size() == static_cast<size_t>(destRank) && + "Invalid number of write indices!"); + if (writeIndices.empty()) { + auto zero = builder.create<arith::ConstantIndexOp>(loc, 0); + writeIndices.assign(destRank, zero); } // Generate the xfer_write Op - auto zero = builder.create<arith::ConstantIndexOp>(loc, 0); - Operation *write = builder.create<vector::TransferWriteOp>( - loc, - /*vector=*/vectorToStore, - /*source=*/dest, - /*indices=*/SmallVector<Value>(rank, zero), - /*inBounds=*/inBoundsVal); - assert(llvm::none_of( - destShape.drop_front(inputVecSizesForLeadingDims.size()), - [](int64_t size) { return size == ShapedType::kDynamic; }) && - "Only dims aligned with inputVecSizesForLeadingDims may be dynamic"); + Operation *write = + builder.create<vector::TransferWriteOp>(loc, + /*vector=*/vecToStore, + /*source=*/dest, + /*indices=*/writeIndices, + /*inBounds=*/inBoundsVal); // If masking is disabled, exit. if (useInBoundsInsteadOfMasking) return write; - // Check if masking is needed. - bool needMaskForWrite = - !llvm::equal(inputVecSizesForLeadingDims, - destShape.take_front(inputVecSizesForLeadingDims.size())); + // Check if masking is needed. If not, exit. + if (llvm::equal(vecToStoreShape, destShape.take_back(vecToStoreRank))) + return write; - // If masking is needed, generate the mask and mask the operation. - if (needMaskForWrite) { - SmallVector<int64_t> writeMaskShape; - writeMaskShape.append(inputVecSizesForLeadingDims.begin(), - inputVecSizesForLeadingDims.end()); - writeMaskShape.append(destShape.begin() + - inputVecSizesForLeadingDims.size(), - destShape.end()); - auto writeMaskType = VectorType::get(writeMaskShape, builder.getI1Type()); - Value maskForWrite = builder.create<vector::CreateMaskOp>( - loc, writeMaskType, tensor::getMixedSizes(builder, loc, dest)); - write = mlir::vector::maskOperation(builder, write, maskForWrite); - } + // Compute the mask and mask the write Op. + auto writeMaskType = VectorType::get(vecToStoreShape, builder.getI1Type()); - return write; + SmallVector<OpFoldResult> destSizes = + tensor::getMixedSizes(builder, loc, dest); + SmallVector<OpFoldResult> maskSizes(destSizes.end() - vecToStoreRank, + destSizes.end()); + + if (isMaskTriviallyFoldable(maskSizes, writeIndices, destShape, + vecToStoreShape)) + return write; + + Value maskForWrite = + builder.createOrFold<vector::CreateMaskOp>(loc, writeMaskType, maskSizes); + return mlir::vector::maskOperation(builder, write, maskForWrite); } /// Vectorize linalg::PackOp with (1) static inner_tiles (2) constant @@ -1701,9 +1791,7 @@ vectorizeAsTensorPackOp(RewriterBase &rewriter, linalg::PackOp packOp, loc, reifiedReturnShapes[0], transposeOp.getResult().getType().getElementType()); Operation *write = - createWriteOrMaskedWrite(rewriter, loc, transposeOp.getResult(), dest, - /*inputVecSizesForLeadingDims=*/inputVectorSizes, - /*useInBoundsInsteadOfMasking=*/false); + createWriteOrMaskedWrite(rewriter, loc, transposeOp.getResult(), dest); newResults.push_back(write->getResult(0)); return success(); } @@ -1839,10 +1927,9 @@ vectorizeAsTensorUnpackOp(RewriterBase &rewriter, linalg::UnPackOp unpackOp, Value dest = rewriter.create<tensor::EmptyOp>( loc, reifiedRetShapes[0], shapeCastOp.getResult().getType().getElementType()); - Operation *write = - createWriteOrMaskedWrite(rewriter, loc, shapeCastOp.getResult(), dest, - /*inputVecSizesForLeadingDims=*/writeVectorSizes, - useInBoundsInsteadOfMasking); + Operation *write = createWriteOrMaskedWrite( + rewriter, loc, shapeCastOp.getResult(), dest, + /*writeIndices=*/{}, useInBoundsInsteadOfMasking); newResults.push_back(write->getResult(0)); return success(); } @@ -1874,10 +1961,7 @@ vectorizeAsTensorPadOp(RewriterBase &rewriter, tensor::PadOp padOp, // Create Xfer write Op Value dest = rewriter.create<tensor::EmptyOp>( loc, reifiedReturnShapes[0], padOp.getResultType().getElementType()); - Operation *write = - createWriteOrMaskedWrite(rewriter, loc, maskedRead, dest, - /*inputVecSizesForLeadingDims=*/inputVectorSizes, - /*useInBoundsInsteadOfMasking=*/false); + Operation *write = createWriteOrMaskedWrite(rewriter, loc, maskedRead, dest); newResults.push_back(write->getResult(0)); return success(); } @@ -2163,14 +2247,12 @@ static LogicalResult vectorizeLinalgOpPrecondition( })) { continue; } - if (llvm::any_of(innerOp.getOperandTypes(), [](Type type) { - return !VectorType::isValidElementType(type); - })) { + if (!llvm::all_of(innerOp.getOperandTypes(), + VectorType::isValidElementType)) { return failure(); } - if (llvm::any_of(innerOp.getResultTypes(), [](Type type) { - return !VectorType::isValidElementType(type); - })) { + if (!llvm::all_of(innerOp.getResultTypes(), + VectorType::isValidElementType)) { return failure(); } } @@ -2883,22 +2965,14 @@ vectorizeAsInsertSliceOp(RewriterBase &rewriter, tensor::InsertSliceOp sliceOp, sliceOp.getLoc(), elemType, rewriter.getZeroAttr(elemType)); } - // 2. Get the vector shape and in-bounds attributes + // 2. Get the vector shape SmallVector<int64_t> vecShape; - SmallVector<bool> readInBounds; - SmallVector<bool> writeInBounds; size_t rankDiff = resultType.getRank() - sourceType.getRank(); for (int64_t i = 0, end = sourceType.getRank(); i < end; ++i) { if (!inputVectorSizes.empty()) { vecShape.push_back(inputVectorSizes[i]); - readInBounds.push_back(false); - writeInBounds.push_back(false); } else if (!sourceType.isDynamicDim(i)) { vecShape.push_back(sourceType.getDimSize(i)); - // Source shape is statically known: Neither read nor write are - // out-of-bounds. - readInBounds.push_back(true); - writeInBounds.push_back(true); } else if (!resultType.isDynamicDim(i)) { // Source shape is not statically known, but result shape is. // Vectorize with size of result shape. This may be larger than the @@ -2906,69 +2980,30 @@ vectorizeAsInsertSliceOp(RewriterBase &rewriter, tensor::InsertSliceOp sliceOp, // FIXME: Using rankDiff implies that the source tensor is inserted at // the end of the destination tensor. However, that's not required. vecShape.push_back(resultType.getDimSize(rankDiff + i)); - // Read may be out-of-bounds because the result size could be larger - // than the source size. - readInBounds.push_back(false); - // Write will be in-bounds provided that the corresponding write idx is 0. - // To keep this logic simple, conservatively mark as out-of-bounds. - writeInBounds.push_back(false); } else { // Neither source nor result dim of padOp is static. Cannot vectorize // the copy. - // TODO: Add support for masking return failure(); } } auto vecType = VectorType::get(vecShape, sourceType.getElementType()); // 3. Generate TransferReadOp + TransferWriteOp - ReifiedRankedShapedTypeDims reifiedSrcSizes; - Value maskOp; - - // If vector sizes are user provided, make sure to mask. First, generate the - // mask. - if (!inputVectorSizes.empty()) { - auto *srcDefOp = source.getDefiningOp(); - if (!srcDefOp) { - LDBG("Unable to get the defining Op of " << sliceOp); - return failure(); - } - - LogicalResult status = - cast<ReifyRankedShapedTypeOpInterface>(srcDefOp).reifyResultShapes( - rewriter, reifiedSrcSizes); - if (status.failed()) { - LDBG("Unable to reify result shapes of " << srcDefOp); - return failure(); - } - - // Create the mask - auto readMaskType = VectorType::get(inputVectorSizes, rewriter.getI1Type()); - maskOp = rewriter.create<vector::CreateMaskOp>( - sliceOp.getLoc(), readMaskType, reifiedSrcSizes[0]); - } + auto loc = sliceOp.getLoc(); + // Create read SmallVector<Value> readIndices( - vecType.getRank(), - rewriter.create<arith::ConstantIndexOp>(sliceOp.getLoc(), 0)); - Operation *read = rewriter.create<vector::TransferReadOp>( - sliceOp.getLoc(), vecType, source, readIndices, padValue, - ArrayRef<bool>{readInBounds}); - - if (maskOp) { - read = mlir::vector::maskOperation(rewriter, read, maskOp); - } - - auto writeIndices = getValueOrCreateConstantIndexOp( - rewriter, sliceOp.getLoc(), sliceOp.getMixedOffsets()); - - Operation *write = rewriter.create<vector::TransferWriteOp>( - sliceOp.getLoc(), read->getResult(0), sliceOp.getDest(), writeIndices, - ArrayRef<bool>{writeInBounds}); - - if (maskOp) { - write = mlir::vector::maskOperation(rewriter, write, maskOp); - } + vecType.getRank(), rewriter.create<arith::ConstantIndexOp>(loc, 0)); + Value read = mlir::vector::createReadOrMaskedRead( + rewriter, loc, source, vecType.getShape(), padValue, + /*useInBoundsInsteadOfMasking=*/inputVectorSizes.empty()); + + // Create write + auto writeIndices = + getValueOrCreateConstantIndexOp(rewriter, loc, sliceOp.getMixedOffsets()); + Operation *write = + createWriteOrMaskedWrite(rewriter, loc, read, sliceOp.getDest(), + writeIndices, inputVectorSizes.empty()); // 4. Finalize newResults.push_back(write->getResult(0)); diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp index 219619981629..34ae83b25c39 100644 --- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp +++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp @@ -799,8 +799,7 @@ LogicalResult SparseTensorEncodingAttr::verify( "before singleton level"; auto *curCOOEnd = std::find_if_not(it, lvlTypes.end(), isSingletonLT); - if (!std::all_of(it, curCOOEnd, - [](LevelType i) { return isSingletonLT(i); })) + if (!std::all_of(it, curCOOEnd, isSingletonLT)) return emitError() << "expected all singleton lvlTypes " "following a singleton level"; // We can potentially support mixed SoA/AoS singleton levels. @@ -833,8 +832,7 @@ LogicalResult SparseTensorEncodingAttr::verify( it != std::end(lvlTypes)) { if (it != lvlTypes.end() - 1) return emitError() << "expected n_out_of_m to be the last level type"; - if (!std::all_of(lvlTypes.begin(), it, - [](LevelType i) { return isDenseLT(i); })) + if (!std::all_of(lvlTypes.begin(), it, isDenseLT)) return emitError() << "expected all dense lvlTypes " "before a n_out_of_m level"; if (dimToLvl && (dimToLvl.getNumDims() != dimToLvl.getNumResults())) { diff --git a/mlir/lib/Dialect/Tensor/Extensions/AllExtensions.cpp b/mlir/lib/Dialect/Tensor/Extensions/AllExtensions.cpp index 93e1a2021857..0421a6c0ff80 100644 --- a/mlir/lib/Dialect/Tensor/Extensions/AllExtensions.cpp +++ b/mlir/lib/Dialect/Tensor/Extensions/AllExtensions.cpp @@ -13,4 +13,4 @@ using namespace mlir; void mlir::tensor::registerAllExtensions(DialectRegistry ®istry) { registerShardingInterfaceExternalModels(registry); -}
\ No newline at end of file +} diff --git a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt index 5425615dac39..66b73f9bead5 100644 --- a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt @@ -30,6 +30,7 @@ add_mlir_dialect_library(MLIRTensorDialect MLIRInferIntRangeInterface MLIRInferTypeOpInterface MLIRParallelCombiningOpInterface + MLIRRuntimeVerifiableOpInterface MLIRShapedOpInterfaces MLIRSideEffectInterfaces MLIRSupport diff --git a/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp index e7d8f52d309c..5b65e47bc937 100644 --- a/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp @@ -12,6 +12,7 @@ #include "mlir/Dialect/Complex/IR/Complex.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h" +#include "mlir/Interfaces/RuntimeVerifiableOpInterface.h" #include "mlir/Interfaces/SubsetOpInterface.h" #include "mlir/Transforms/InliningUtils.h" @@ -58,6 +59,8 @@ void TensorDialect::initialize() { InsertSliceOp, ReshapeOp>(); declarePromisedInterfaces<ReifyRankedShapedTypeOpInterface, ExpandShapeOp, CollapseShapeOp, PadOp>(); + declarePromisedInterfaces<RuntimeVerifiableOpInterface, CastOp, DimOp, + ExtractOp, InsertOp, ExtractSliceOp>(); declarePromisedInterfaces<SubsetOpInterface, ExtractSliceOp, InsertSliceOp, ParallelInsertSliceOp>(); declarePromisedInterfaces<SubsetInsertionOpInterface, InsertSliceOp, diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp index 12e8b257ce9f..04242cad9ecb 100644 --- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -1061,8 +1061,7 @@ void DimOp::getCanonicalizationPatterns(RewritePatternSet &results, void EmptyOp::build(OpBuilder &builder, OperationState &result, ArrayRef<int64_t> staticShape, Type elementType, Attribute encoding) { - assert(all_of(staticShape, - [](int64_t sz) { return !ShapedType::isDynamic(sz); }) && + assert(none_of(staticShape, ShapedType::isDynamic) && "expected only static sizes"); build(builder, result, staticShape, elementType, ValueRange{}, encoding); } @@ -1624,76 +1623,6 @@ OpFoldResult GatherOp::fold(FoldAdaptor adaptor) { // InsertOp //===----------------------------------------------------------------------===// -namespace { - -/// Pattern to fold an insert op of a constant destination and scalar to a new -/// constant. -/// -/// Example: -/// ``` -/// %0 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32> -/// %c0 = arith.constant 0 : index -/// %c4_f32 = arith.constant 4.0 : f32 -/// %1 = tensor.insert %c4_f32 into %0[%c0] : tensor<4xf32> -/// ``` -/// is rewritten into: -/// ``` -/// %1 = arith.constant dense<[4.0, 2.0, 3.0, 4.0]> : tensor<4xf32> -/// ``` -class InsertOpConstantFold final : public OpRewritePattern<InsertOp> { -public: - using OpRewritePattern<InsertOp>::OpRewritePattern; - - LogicalResult matchAndRewrite(InsertOp insertOp, - PatternRewriter &rewriter) const override { - // Requires a ranked tensor type. - auto destType = - llvm::dyn_cast<RankedTensorType>(insertOp.getDest().getType()); - if (!destType) - return failure(); - - // Pattern requires constant indices - SmallVector<uint64_t, 8> indices; - for (OpFoldResult indice : getAsOpFoldResult(insertOp.getIndices())) { - auto indiceAttr = dyn_cast<Attribute>(indice); - if (!indiceAttr) - return failure(); - indices.push_back(llvm::cast<IntegerAttr>(indiceAttr).getInt()); - } - - // Requires a constant scalar to insert - OpFoldResult scalar = getAsOpFoldResult(insertOp.getScalar()); - Attribute scalarAttr = dyn_cast<Attribute>(scalar); - if (!scalarAttr) - return failure(); - - if (auto constantOp = dyn_cast_or_null<arith::ConstantOp>( - insertOp.getDest().getDefiningOp())) { - if (auto sourceAttr = - llvm::dyn_cast<ElementsAttr>(constantOp.getValue())) { - // Update the attribute at the inserted index. - auto sourceValues = sourceAttr.getValues<Attribute>(); - auto flattenedIndex = sourceAttr.getFlattenedIndex(indices); - std::vector<Attribute> updatedValues; - updatedValues.reserve(sourceAttr.getNumElements()); - for (unsigned i = 0; i < sourceAttr.getNumElements(); ++i) { - updatedValues.push_back(i == flattenedIndex ? scalarAttr - : sourceValues[i]); - } - rewriter.replaceOpWithNewOp<arith::ConstantOp>( - insertOp, sourceAttr.getType(), - DenseElementsAttr::get(cast<ShapedType>(sourceAttr.getType()), - updatedValues)); - return success(); - } - } - - return failure(); - } -}; - -} // namespace - void InsertOp::getAsmResultNames( function_ref<void(Value, StringRef)> setNameFn) { setNameFn(getResult(), "inserted"); @@ -1717,11 +1646,6 @@ OpFoldResult InsertOp::fold(FoldAdaptor adaptor) { return {}; } -void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results, - MLIRContext *context) { - results.add<InsertOpConstantFold>(context); -} - //===----------------------------------------------------------------------===// // GenerateOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt index 7880d1c5a0c5..99e1c4fec846 100644 --- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt @@ -8,6 +8,7 @@ add_mlir_dialect_library(MLIRTensorTransforms MergeConsecutiveInsertExtractSlicePatterns.cpp ReshapePatterns.cpp RewriteAsConstant.cpp + RuntimeOpVerification.cpp SwapExtractSliceWithProducerPatterns.cpp SubsetInsertionOpInterfaceImpl.cpp @@ -28,6 +29,7 @@ add_mlir_dialect_library(MLIRTensorTransforms MLIRLinalgDialect MLIRMemRefDialect MLIRPass + MLIRRuntimeVerifiableOpInterface MLIRSCFDialect MLIRSubsetOpInterface MLIRTensorDialect diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp index 2b229d60c691..657624b817af 100644 --- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp +++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp @@ -185,13 +185,11 @@ struct BubbleUpExpandThroughParallelCollapse ArrayRef<int64_t> collapsedStaticShapes = staticSourceSize.slice( collapseReassociation.front(), collapseReassociation.size()); int64_t numCollapsedDynamic = - llvm::count_if(collapsedStaticShapes, - [](int64_t d) { return ShapedType::isDynamic(d); }); + llvm::count_if(collapsedStaticShapes, ShapedType::isDynamic); ArrayRef<int64_t> expandedStaticShapes = staticResultSize.slice( expandReassociation.front(), expandReassociation.size()); int64_t numExpandedDynamic = - llvm::count_if(expandedStaticShapes, - [](int64_t d) { return ShapedType::isDynamic(d); }); + llvm::count_if(expandedStaticShapes, ShapedType::isDynamic); if (numCollapsedDynamic > 1 || numExpandedDynamic > 1 || collapsedStaticShapes != expandedStaticShapes) { return failure(); diff --git a/mlir/lib/Dialect/Tensor/Transforms/RuntimeOpVerification.cpp b/mlir/lib/Dialect/Tensor/Transforms/RuntimeOpVerification.cpp new file mode 100644 index 000000000000..6138821ee8c6 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/Transforms/RuntimeOpVerification.cpp @@ -0,0 +1,208 @@ +//===- RuntimeOpVerification.cpp - Op Verification ------------------------===// +// +// 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/Tensor/Transforms/RuntimeOpVerification.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Utils/IndexingUtils.h" +#include "mlir/Interfaces/RuntimeVerifiableOpInterface.h" + +using namespace mlir; + +namespace mlir { +namespace tensor { +namespace { +/// Generate a runtime check for lb <= value < ub. +Value generateInBoundsCheck(OpBuilder &builder, Location loc, Value value, + Value lb, Value ub) { + Value inBounds1 = builder.createOrFold<arith::CmpIOp>( + loc, arith::CmpIPredicate::sge, value, lb); + Value inBounds2 = builder.createOrFold<arith::CmpIOp>( + loc, arith::CmpIPredicate::slt, value, ub); + Value inBounds = + builder.createOrFold<arith::AndIOp>(loc, inBounds1, inBounds2); + return inBounds; +} + +struct CastOpInterface + : public RuntimeVerifiableOpInterface::ExternalModel<CastOpInterface, + CastOp> { + void generateRuntimeVerification(Operation *op, OpBuilder &builder, + Location loc) const { + auto castOp = cast<CastOp>(op); + auto srcType = cast<TensorType>(castOp.getSource().getType()); + + // Nothing to check if the result is an unranked tensor. + auto resultType = dyn_cast<RankedTensorType>(castOp.getType()); + if (!resultType) + return; + + if (isa<UnrankedTensorType>(srcType)) { + // Check rank. + Value srcRank = builder.create<RankOp>(loc, castOp.getSource()); + Value resultRank = + builder.create<arith::ConstantIndexOp>(loc, resultType.getRank()); + Value isSameRank = builder.create<arith::CmpIOp>( + loc, arith::CmpIPredicate::eq, srcRank, resultRank); + builder.create<cf::AssertOp>( + loc, isSameRank, + RuntimeVerifiableOpInterface::generateErrorMessage(op, + "rank mismatch")); + } + + // Check dimension sizes. + for (const auto &it : llvm::enumerate(resultType.getShape())) { + // Static dim size -> static/dynamic dim size does not need verification. + if (auto rankedSrcType = dyn_cast<RankedTensorType>(srcType)) + if (!rankedSrcType.isDynamicDim(it.index())) + continue; + + // Static/dynamic dim size -> dynamic dim size does not need verification. + if (resultType.isDynamicDim(it.index())) + continue; + + Value srcDimSz = + builder.create<DimOp>(loc, castOp.getSource(), it.index()); + Value resultDimSz = + builder.create<arith::ConstantIndexOp>(loc, it.value()); + Value isSameSz = builder.create<arith::CmpIOp>( + loc, arith::CmpIPredicate::eq, srcDimSz, resultDimSz); + builder.create<cf::AssertOp>( + loc, isSameSz, + RuntimeVerifiableOpInterface::generateErrorMessage( + op, "size mismatch of dim " + std::to_string(it.index()))); + } + } +}; + +struct DimOpInterface + : public RuntimeVerifiableOpInterface::ExternalModel<DimOpInterface, + DimOp> { + void generateRuntimeVerification(Operation *op, OpBuilder &builder, + Location loc) const { + auto dimOp = cast<DimOp>(op); + Value rank = builder.create<RankOp>(loc, dimOp.getSource()); + Value zero = builder.create<arith::ConstantIndexOp>(loc, 0); + builder.create<cf::AssertOp>( + loc, generateInBoundsCheck(builder, loc, dimOp.getIndex(), zero, rank), + RuntimeVerifiableOpInterface::generateErrorMessage( + op, "index is out of bounds")); + } +}; + +/// Verifies that the indices on extract/insert ops are in-bounds of the +/// tensor's index space: 0 <= index#i < dim#i +template <typename OpTy> +struct ExtractInsertOpInterface + : public RuntimeVerifiableOpInterface::ExternalModel< + ExtractInsertOpInterface<OpTy>, OpTy> { + void generateRuntimeVerification(Operation *op, OpBuilder &builder, + Location loc) const { + auto extractInsertOp = cast<OpTy>(op); + + Value tensor; + if constexpr (std::is_same_v<OpTy, ExtractOp>) { + tensor = extractInsertOp.getTensor(); + } else if constexpr (std::is_same_v<OpTy, InsertOp>) { + tensor = extractInsertOp.getDest(); + } else { + llvm_unreachable("invalid op"); + } + auto tensorType = cast<RankedTensorType>(tensor.getType()); + auto rank = tensorType.getRank(); + if (rank == 0) { + // Nothing to check for 0-d tensors. + return; + } + + auto indices = extractInsertOp.getIndices(); + auto zero = builder.create<arith::ConstantIndexOp>(loc, 0); + Value assertCond; + for (auto i : llvm::seq<int64_t>(0, rank)) { + Value dimOp = builder.createOrFold<tensor::DimOp>(loc, tensor, i); + Value inBounds = + generateInBoundsCheck(builder, loc, indices[i], zero, dimOp); + assertCond = + i > 0 ? builder.createOrFold<arith::AndIOp>(loc, assertCond, inBounds) + : inBounds; + } + builder.create<cf::AssertOp>( + loc, assertCond, + RuntimeVerifiableOpInterface::generateErrorMessage( + op, "out-of-bounds access")); + } +}; + +struct ExtractSliceOpInterface + : public RuntimeVerifiableOpInterface::ExternalModel< + ExtractSliceOpInterface, ExtractSliceOp> { + void generateRuntimeVerification(Operation *op, OpBuilder &builder, + Location loc) const { + auto extractSliceOp = cast<ExtractSliceOp>(op); + RankedTensorType sourceType = extractSliceOp.getSource().getType(); + + // For each dimension, assert that: + // 0 <= offset < dim_size + // 0 <= offset + (size - 1) * stride < dim_size + Value zero = builder.create<arith::ConstantIndexOp>(loc, 0); + Value one = builder.create<arith::ConstantIndexOp>(loc, 1); + for (int64_t i = 0, e = sourceType.getRank(); i < e; ++i) { + Value offset = getValueOrCreateConstantIndexOp( + builder, loc, extractSliceOp.getMixedOffsets()[i]); + Value size = getValueOrCreateConstantIndexOp( + builder, loc, extractSliceOp.getMixedSizes()[i]); + Value stride = getValueOrCreateConstantIndexOp( + builder, loc, extractSliceOp.getMixedStrides()[i]); + + // Verify that offset is in-bounds. + Value dimSize = builder.createOrFold<tensor::DimOp>( + loc, extractSliceOp.getSource(), i); + Value offsetInBounds = + generateInBoundsCheck(builder, loc, offset, zero, dimSize); + builder.create<cf::AssertOp>( + loc, offsetInBounds, + RuntimeVerifiableOpInterface::generateErrorMessage( + op, "offset " + std::to_string(i) + " is out-of-bounds")); + + // Verify that slice does not run out-of-bounds. + Value sizeMinusOne = builder.create<arith::SubIOp>(loc, size, one); + Value sizeMinusOneTimesStride = + builder.create<arith::MulIOp>(loc, sizeMinusOne, stride); + Value lastPos = + builder.create<arith::AddIOp>(loc, offset, sizeMinusOneTimesStride); + Value lastPosInBounds = + generateInBoundsCheck(builder, loc, lastPos, zero, dimSize); + builder.create<cf::AssertOp>( + loc, lastPosInBounds, + RuntimeVerifiableOpInterface::generateErrorMessage( + op, "extract_slice runs out-of-bounds along dimension " + + std::to_string(i))); + } + } +}; +} // namespace +} // namespace tensor +} // namespace mlir + +void mlir::tensor::registerRuntimeVerifiableOpInterfaceExternalModels( + DialectRegistry ®istry) { + registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) { + CastOp::attachInterface<CastOpInterface>(*ctx); + DimOp::attachInterface<DimOpInterface>(*ctx); + ExtractOp::attachInterface<ExtractInsertOpInterface<ExtractOp>>(*ctx); + ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx); + InsertOp::attachInterface<ExtractInsertOpInterface<InsertOp>>(*ctx); + + // Load additional dialects of which ops may get created. + ctx->loadDialect<arith::ArithDialect, cf::ControlFlowDialect>(); + }); +} diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp index a22e6b7aa979..a32e4ccbed59 100644 --- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp +++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp @@ -1857,7 +1857,8 @@ LogicalResult tosa::MulOp::inferReturnTypeComponents( } LogicalResult tosa::MulOp::verify() { - auto resElemType = getElementTypeOrSelf(getOutput()); + const Value output = getOutput(); + auto resElemType = getElementTypeOrSelf(output); // Verify if the element type among operands and result match tosa // specification. @@ -1897,59 +1898,39 @@ LogicalResult tosa::MulOp::verify() { // Verify the op has same ranks for all main operands (excludes extra operands // such as shift of mul op, so this is the only difference with the built-in // `SameOperandsAndResultRank` trait) and results types, if known. - - // delegate function that returns true if type is a shaped type with known - // rank - auto hasRank = [](const Type type) { - if (auto shaped_type = dyn_cast<ShapedType>(type)) - return shaped_type.hasRank(); - - return false; - }; - - auto rankedOperandTypes = - llvm::to_vector(llvm::make_filter_range(getOperandTypes(), hasRank)); - - auto rankedResultTypes = - llvm::make_filter_range(getOperation()->getResultTypes(), hasRank); - - // If all operands and results are unranked, then no further verification. - if (rankedOperandTypes.empty() && rankedResultTypes.empty()) + TypeRange operandTypes = getOperandTypes(); + ShapedType aType = cast<ShapedType>(operandTypes[0]); + ShapedType bType = cast<ShapedType>(operandTypes[1]); + + const bool aHasRank = aType.hasRank(); + const bool bHasRank = bType.hasRank(); + if (aHasRank && bHasRank) { + const int64_t aRank = aType.getRank(); + const int64_t bRank = bType.getRank(); + if (aRank != bRank) + return emitOpError("a and b operands don't have matching ranks, got ") + << aRank << " and " << bRank; + + // check for broadcast compatible shapes + SmallVector<int64_t> resultShape; + if (!mlir::OpTrait::util::getBroadcastedShape( + aType.getShape(), bType.getShape(), resultShape)) + return emitOpError("a and b operands don't have broadcast-compatible " + "shapes, got ") + << aType << " and " << bType; + } + + ShapedType resultType = cast<ShapedType>(output.getType()); + if (!resultType.hasRank()) return success(); - // delegate function that returns rank of shaped type with known rank - auto getRank = [](const Type type) { - return cast<ShapedType>(type).getRank(); - }; - - auto rank = !rankedOperandTypes.empty() ? getRank(*rankedOperandTypes.begin()) - : getRank(*rankedResultTypes.begin()); - - for (size_t i = 0; i < 2; ++i) { - if (rank != getRank(rankedOperandTypes[i])) { - return emitOpError("operands don't have matching ranks"); - } - } - - for (const auto type : rankedResultTypes) { - if (rank != getRank(type)) { - return emitOpError("result type has different rank than operands"); - } - } - - // check for broadcast compatible shapes in first two operands (ignoring - // shift) - - // delegate function that returns shape of shaped type - auto getShape = [](const Type type) { - return mlir::cast<ShapedType>(type).getShape(); - }; - SmallVector<int64_t> resultShape; - if (!mlir::OpTrait::util::getBroadcastedShape(getShape(rankedOperandTypes[0]), - getShape(rankedOperandTypes[1]), - resultShape)) { - return emitOpError("operands don't have broadcast-compatible shapes"); - } + const int64_t resultRank = resultType.getRank(); + if (aHasRank && resultRank != aType.getRank()) + return emitOpError("result type has different rank than a, got ") + << resultRank << " vs " << aType.getRank(); + if (bHasRank && resultRank != bType.getRank()) + return emitOpError("result type has different rank than b, got ") + << resultRank << " vs " << bType.getRank(); return success(); } @@ -2692,6 +2673,73 @@ LogicalResult tosa::ScatterOp::verify() { .failed()) { return failure(); } + + const ShapeAdaptor valuesInShape(getValuesIn().getType()); + const ShapeAdaptor indicesShape(getIndices().getType()); + const ShapeAdaptor inputShape(getInput().getType()); + const ShapeAdaptor outputShape(getValuesOut().getType()); + + int64_t N = ShapedType::kDynamic; + int64_t K = ShapedType::kDynamic; + int64_t W = ShapedType::kDynamic; + int64_t C = ShapedType::kDynamic; + if (valuesInShape.hasRank()) { + N = valuesInShape.getDimSize(0); + K = valuesInShape.getDimSize(1); + C = valuesInShape.getDimSize(2); + } + if (indicesShape.hasRank()) { + const int64_t indicesN = indicesShape.getDimSize(0); + W = indicesShape.getDimSize(1); + if (N == ShapedType::kDynamic) + N = indicesN; + else if (indicesN != ShapedType::kDynamic && N != indicesN) + return emitOpError() << "requires indices dimension 0 to have size " << N + << ", got " << indicesN; + } + if (inputShape.hasRank()) { + const int64_t inputN = inputShape.getDimSize(0); + const int64_t inputW = inputShape.getDimSize(1); + const int64_t inputC = inputShape.getDimSize(2); + if (N == ShapedType::kDynamic) + N = inputN; + else if (inputN != ShapedType::kDynamic && N != inputN) + return emitOpError() << "requires input dimension 0 to have size " << N + << ", got " << inputN; + if (W == ShapedType::kDynamic) + W = inputW; + else if (inputW != ShapedType::kDynamic && W != inputW) + return emitOpError() << "requires input dimension 1 to have size " << W + << ", got " << inputW; + + if (C == ShapedType::kDynamic) + C = inputC; + else if (inputC != ShapedType::kDynamic && C != inputC) + return emitOpError() << "requires input dimension 2 to have size " << C + << ", got " << inputC; + } + if (outputShape.hasRank()) { + const int64_t outputN = outputShape.getDimSize(0); + const int64_t outputK = outputShape.getDimSize(1); + const int64_t outputC = outputShape.getDimSize(2); + if (N != ShapedType::kDynamic && outputN != ShapedType::kDynamic && + N != outputN) + return emitOpError() << "requires values_out dimension 0 to have size " + << N << ", got " << outputN; + if (K == ShapedType::kDynamic) + K = outputK; + else if (outputK != ShapedType::kDynamic && K != outputK) + return emitOpError() << "requires values_out dimension 1 to have size " + << K << ", got " << outputK; + if (C != ShapedType::kDynamic && outputC != ShapedType::kDynamic && + C != outputC) + return emitOpError() << "requires values_out dimension 2 to have size " + << C << ", got " << outputC; + } + if (K != ShapedType::kDynamic && W != ShapedType::kDynamic && !(K >= W)) + return emitOpError() << "requires dimensions K >= W, got K=" << K + << " and W=" << W; + return success(); } diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp index 9aa0051070bd..29ec9f8db261 100644 --- a/mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp +++ b/mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp @@ -344,4 +344,4 @@ public: validateSameOperandsAndResultRankTrait(func.getBody()); } }; -} // namespace
\ No newline at end of file +} // namespace diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp index 673743f22249..a0f9518e3d12 100644 --- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp +++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp @@ -53,6 +53,12 @@ using namespace mlir; +static ParseResult parseApplyRegisteredPassOptions( + OpAsmParser &parser, ArrayAttr &options, + SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dynamicOptions); +static void printApplyRegisteredPassOptions(OpAsmPrinter &printer, + Operation *op, ArrayAttr options, + ValueRange dynamicOptions); static ParseResult parseSequenceOpOperands( OpAsmParser &parser, std::optional<OpAsmParser::UnresolvedOperand> &root, Type &rootType, @@ -766,17 +772,53 @@ void transform::ApplyLoopInvariantCodeMotionOp::getEffects( // ApplyRegisteredPassOp //===----------------------------------------------------------------------===// -DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne( - transform::TransformRewriter &rewriter, Operation *target, - ApplyToEachResultList &results, transform::TransformState &state) { - // Make sure that this transform is not applied to itself. Modifying the - // transform IR while it is being interpreted is generally dangerous. Even - // more so when applying passes because they may perform a wide range of IR - // modifications. - DiagnosedSilenceableFailure payloadCheck = - ensurePayloadIsSeparateFromTransform(*this, target); - if (!payloadCheck.succeeded()) - return payloadCheck; +void transform::ApplyRegisteredPassOp::getEffects( + SmallVectorImpl<MemoryEffects::EffectInstance> &effects) { + consumesHandle(getTargetMutable(), effects); + onlyReadsHandle(getDynamicOptionsMutable(), effects); + producesHandle(getOperation()->getOpResults(), effects); + modifiesPayload(effects); +} + +DiagnosedSilenceableFailure +transform::ApplyRegisteredPassOp::apply(transform::TransformRewriter &rewriter, + transform::TransformResults &results, + transform::TransformState &state) { + // Obtain a single options-string from options passed statically as + // string attributes as well as "dynamically" through params. + std::string options; + OperandRange dynamicOptions = getDynamicOptions(); + size_t dynamicOptionsIdx = 0; + for (auto [idx, optionAttr] : llvm::enumerate(getOptions())) { + if (idx > 0) + options += " "; // Interleave options seperator. + + if (auto strAttr = dyn_cast<StringAttr>(optionAttr)) { + options += strAttr.getValue(); + } else if (isa<UnitAttr>(optionAttr)) { + assert(dynamicOptionsIdx < dynamicOptions.size() && + "number of dynamic option markers (UnitAttr) in options ArrayAttr " + "should be the same as the number of options passed as params"); + ArrayRef<Attribute> dynamicOption = + state.getParams(dynamicOptions[dynamicOptionsIdx++]); + if (dynamicOption.size() != 1) + return emitSilenceableError() << "options passed as a param must have " + "a single value associated, param " + << dynamicOptionsIdx - 1 << " associates " + << dynamicOption.size(); + + if (auto dynamicOptionStr = dyn_cast<StringAttr>(dynamicOption[0])) { + options += dynamicOptionStr.getValue(); + } else { + return emitSilenceableError() + << "options passed as a param must be a string, got " + << dynamicOption[0]; + } + } else { + llvm_unreachable( + "expected options element to be either StringAttr or UnitAttr"); + } + } // Get pass or pass pipeline from registry. const PassRegistryEntry *info = PassPipelineInfo::lookup(getPassName()); @@ -786,9 +828,9 @@ DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne( return emitDefiniteFailure() << "unknown pass or pass pipeline: " << getPassName(); - // Create pass manager and run the pass or pass pipeline. + // Create pass manager and add the pass or pass pipeline. PassManager pm(getContext()); - if (failed(info->addToPipeline(pm, getOptions(), [&](const Twine &msg) { + if (failed(info->addToPipeline(pm, options, [&](const Twine &msg) { emitError(msg); return failure(); }))) { @@ -796,16 +838,114 @@ DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne( << "failed to add pass or pass pipeline to pipeline: " << getPassName(); } - if (failed(pm.run(target))) { - auto diag = emitSilenceableError() << "pass pipeline failed"; - diag.attachNote(target->getLoc()) << "target op"; - return diag; + + auto targets = SmallVector<Operation *>(state.getPayloadOps(getTarget())); + for (Operation *target : targets) { + // Make sure that this transform is not applied to itself. Modifying the + // transform IR while it is being interpreted is generally dangerous. Even + // more so when applying passes because they may perform a wide range of IR + // modifications. + DiagnosedSilenceableFailure payloadCheck = + ensurePayloadIsSeparateFromTransform(*this, target); + if (!payloadCheck.succeeded()) + return payloadCheck; + + // Run the pass or pass pipeline on the current target operation. + if (failed(pm.run(target))) { + auto diag = emitSilenceableError() << "pass pipeline failed"; + diag.attachNote(target->getLoc()) << "target op"; + return diag; + } } - results.push_back(target); + // The applied pass will have directly modified the payload IR(s). + results.set(llvm::cast<OpResult>(getResult()), targets); return DiagnosedSilenceableFailure::success(); } +static ParseResult parseApplyRegisteredPassOptions( + OpAsmParser &parser, ArrayAttr &options, + SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dynamicOptions) { + auto dynamicOptionMarker = UnitAttr::get(parser.getContext()); + SmallVector<Attribute> optionsArray; + + auto parseOperandOrString = [&]() -> OptionalParseResult { + OpAsmParser::UnresolvedOperand operand; + OptionalParseResult parsedOperand = parser.parseOptionalOperand(operand); + if (parsedOperand.has_value()) { + if (failed(parsedOperand.value())) + return failure(); + + dynamicOptions.push_back(operand); + optionsArray.push_back( + dynamicOptionMarker); // Placeholder for knowing where to + // inject the dynamic option-as-param. + return success(); + } + + StringAttr stringAttr; + OptionalParseResult parsedStringAttr = + parser.parseOptionalAttribute(stringAttr); + if (parsedStringAttr.has_value()) { + if (failed(parsedStringAttr.value())) + return failure(); + optionsArray.push_back(stringAttr); + return success(); + } + + return std::nullopt; + }; + + OptionalParseResult parsedOptionsElement = parseOperandOrString(); + while (parsedOptionsElement.has_value()) { + if (failed(parsedOptionsElement.value())) + return failure(); + parsedOptionsElement = parseOperandOrString(); + } + + if (optionsArray.empty()) { + return parser.emitError(parser.getCurrentLocation()) + << "expected at least one option (either a string or a param)"; + } + options = parser.getBuilder().getArrayAttr(optionsArray); + return success(); +} + +static void printApplyRegisteredPassOptions(OpAsmPrinter &printer, + Operation *op, ArrayAttr options, + ValueRange dynamicOptions) { + size_t currentDynamicOptionIdx = 0; + for (auto [idx, optionAttr] : llvm::enumerate(options)) { + if (idx > 0) + printer << " "; // Interleave options separator. + + if (isa<UnitAttr>(optionAttr)) + printer.printOperand(dynamicOptions[currentDynamicOptionIdx++]); + else if (auto strAttr = dyn_cast<StringAttr>(optionAttr)) + printer.printAttribute(strAttr); + else + llvm_unreachable("each option should be either a StringAttr or UnitAttr"); + } +} + +LogicalResult transform::ApplyRegisteredPassOp::verify() { + size_t numUnitsInOptions = 0; + for (Attribute optionsElement : getOptions()) { + if (isa<UnitAttr>(optionsElement)) + numUnitsInOptions++; + else if (!isa<StringAttr>(optionsElement)) + return emitOpError() << "expected each option to be either a StringAttr " + << "or a UnitAttr, got " << optionsElement; + } + + if (getDynamicOptions().size() != numUnitsInOptions) + return emitOpError() + << "expected the same number of options passed as params as " + << "UnitAttr elements in options ArrayAttr"; + + return success(); +} + //===----------------------------------------------------------------------===// // CastOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp index fcfb401fd986..3179b4f97540 100644 --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -2522,6 +2522,10 @@ void BroadcastOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges, setResultRanges(getResult(), argRanges.front()); } +std::optional<SmallVector<int64_t, 4>> BroadcastOp::getShapeForUnroll() { + return llvm::to_vector<4>(getResultVectorType().getShape()); +} + /// Return the dimensions of the result vector that were formerly ones in the /// source tensor and thus correspond to "dim-1" broadcasting. static llvm::SetVector<int64_t> diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorShapeCast.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorShapeCast.cpp index 23324a007377..39c16fab21c4 100644 --- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorShapeCast.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorShapeCast.cpp @@ -21,177 +21,298 @@ #include "mlir/IR/Location.h" #include "mlir/IR/PatternMatch.h" #include "mlir/IR/TypeUtilities.h" +#include <numeric> #define DEBUG_TYPE "vector-shape-cast-lowering" using namespace mlir; -using namespace mlir::vector; - -/// Increments n-D `indices` by `step` starting from the innermost dimension. -static void incIdx(SmallVectorImpl<int64_t> &indices, VectorType vecType, - int step = 1) { - for (int dim : llvm::reverse(llvm::seq<int>(0, indices.size()))) { - assert(indices[dim] < vecType.getDimSize(dim) && - "Indices are out of bound"); - indices[dim] += step; - if (indices[dim] < vecType.getDimSize(dim)) - break; - indices[dim] = 0; - step = 1; +/// Perform the inplace update +/// rhs <- lhs + rhs +/// +/// where `rhs` is a number expressed in mixed base `base` with most signficant +/// dimensions on the left. For example if `rhs` is {a,b,c} and `base` is +/// {5,3,2} then `rhs` has value a*3*2 + b*2 + c. +/// +/// Some examples where `base` is {5,3,2}: +/// rhs = {0,0,0}, lhs = 1 --> rhs = {0,0,1} +/// rhs = {0,0,1}, lhs = 1 --> rhs = {0,1,0} +/// rhs = {0,0,0}, lhs = 25 --> rhs = {4,0,1} +/// +/// Invalid: +/// rhs = {0,0,2}, lhs = 1 : rhs not in base {5,3,2} +/// +/// Overflows not handled correctly: +/// rhs = {4,2,1}, lhs = 2 --> rhs = {0,0,0} (not {0,0,1}) +static void inplaceAdd(int64_t lhs, ArrayRef<int64_t> base, + MutableArrayRef<int64_t> rhs) { + + // For dimensions in [numIndices - 1, ..., 3, 2, 1, 0]: + for (int dim : llvm::reverse(llvm::seq<int>(0, rhs.size()))) { + int64_t dimBase = base[dim]; + assert(rhs[dim] < dimBase && "rhs not in base"); + + int64_t incremented = rhs[dim] + lhs; + + // If the incremented value excedes the dimension base, we must spill to the + // next most significant dimension and repeat (we might need to spill to + // more significant dimensions multiple times). + lhs = incremented / dimBase; + rhs[dim] = incremented % dimBase; + if (lhs == 0) + break; } } namespace { -/// ShapeOp n-D -> 1-D downcast serves the purpose of flattening N-D to 1-D -/// vectors progressively. This iterates over the n-1 major dimensions of the -/// n-D vector and performs rewrites into: -/// vector.extract from n-D + vector.insert_strided_slice offset into 1-D -class ShapeCastOpNDDownCastRewritePattern - : public OpRewritePattern<vector::ShapeCastOp> { -public: - using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(vector::ShapeCastOp op, - PatternRewriter &rewriter) const override { - auto sourceVectorType = op.getSourceVectorType(); - auto resultVectorType = op.getResultVectorType(); - if (sourceVectorType.isScalable() || resultVectorType.isScalable()) - return failure(); +/// shape_cast is converted to a sequence of extract, extract_strided_slice, +/// insert_strided_slice, and insert operations. The running example will be: +/// +/// %0 = vector.shape_cast %arg0 : +/// vector<2x2x3x4x7x11xi8> to vector<8x6x7x11xi8> +/// +/// In this example the source and result shapes share a common suffix of 7x11. +/// This means we can always decompose the shape_cast into extract, insert, and +/// their strided equivalents, on vectors with shape suffix 7x11. +/// +/// The greatest common divisor (gcd) of the first dimension preceding the +/// common suffix is gcd(4,6) = 2. The algorithm implemented here will operate +/// on vectors with shapes that are `multiples` of (what we define as) the +/// 'atomic shape', 2x7x11. The atomic shape is `gcd` x `common-suffix`. +/// +/// vector<2x2x3x4x7x11xi8> to +/// vector<8x6x7x11xi8> +/// | |||| +/// | ++++------------> common suffix of 7x11 +/// +-----------------> gcd(4,6) is 2 | | +/// | | | +/// v v v +/// atomic shape <----- 2x7x11 +/// +/// +/// +/// The decomposition implemented in this pattern consists of a sequence of +/// repeated steps: +/// +/// (1) Extract vectors from the suffix of the source. +/// In our example this is 2x2x3x4x7x11 -> 4x7x11. +/// +/// (2) Do extract_strided_slice down to the atomic shape. +/// In our example this is 4x7x11 -> 2x7x11. +/// +/// (3) Do insert_strided_slice to the suffix of the result. +/// In our example this is 2x7x11 -> 6x7x11. +/// +/// (4) insert these vectors into the result vector. +/// In our example this is 6x7x11 -> 8x6x7x11. +/// +/// These steps occur with different periods. In this example +/// (1) occurs 12 times, +/// (2) and (3) occur 24 times, and +/// (4) occurs 8 times. +/// +/// Two special cases are handled independently in this pattern +/// (i) A shape_cast that just does leading 1 insertion/removal +/// (ii) A shape_cast where the gcd is 1. +/// +/// These 2 cases can have more compact IR generated by not using the generic +/// algorithm described above. +/// +class ShapeCastOpRewritePattern : public OpRewritePattern<vector::ShapeCastOp> { + + // Case (i) of description. + // Assumes source and result shapes are identical up to some leading ones. + static LogicalResult leadingOnesLowering(vector::ShapeCastOp shapeCast, + PatternRewriter &rewriter) { + + const Location loc = shapeCast.getLoc(); + const VectorType sourceType = shapeCast.getSourceVectorType(); + const VectorType resultType = shapeCast.getResultVectorType(); + + const int64_t sourceRank = sourceType.getRank(); + const int64_t resultRank = resultType.getRank(); + const int64_t delta = sourceRank - resultRank; + const int64_t sourceLeading = delta > 0 ? delta : 0; + const int64_t resultLeading = delta > 0 ? 0 : -delta; + + const Value source = shapeCast.getSource(); + const Value poison = rewriter.create<ub::PoisonOp>(loc, resultType); + const Value extracted = rewriter.create<vector::ExtractOp>( + loc, source, SmallVector<int64_t>(sourceLeading, 0)); + const Value result = rewriter.create<vector::InsertOp>( + loc, extracted, poison, SmallVector<int64_t>(resultLeading, 0)); + + rewriter.replaceOp(shapeCast, result); + return success(); + } - int64_t srcRank = sourceVectorType.getRank(); - int64_t resRank = resultVectorType.getRank(); - if (srcRank < 2 || resRank != 1) - return failure(); + // Case (ii) of description. + // Assumes a shape_cast where the suffix shape of the source starting at + // `sourceDim` and the suffix shape of the result starting at `resultDim` are + // identical. + static LogicalResult noStridedSliceLowering(vector::ShapeCastOp shapeCast, + int64_t sourceDim, + int64_t resultDim, + PatternRewriter &rewriter) { - // Compute the number of 1-D vector elements involved in the reshape. - int64_t numElts = 1; - for (int64_t dim = 0; dim < srcRank - 1; ++dim) - numElts *= sourceVectorType.getDimSize(dim); + const Location loc = shapeCast.getLoc(); - auto loc = op.getLoc(); - SmallVector<int64_t> srcIdx(srcRank - 1, 0); - SmallVector<int64_t> resIdx(resRank, 0); - int64_t extractSize = sourceVectorType.getShape().back(); - Value result = rewriter.create<ub::PoisonOp>(loc, resultVectorType); + const Value source = shapeCast.getSource(); + const ArrayRef<int64_t> sourceShape = + shapeCast.getSourceVectorType().getShape(); - // Compute the indices of each 1-D vector element of the source extraction - // and destination slice insertion and generate such instructions. - for (int64_t i = 0; i < numElts; ++i) { - if (i != 0) { - incIdx(srcIdx, sourceVectorType, /*step=*/1); - incIdx(resIdx, resultVectorType, /*step=*/extractSize); - } + const VectorType resultType = shapeCast.getResultVectorType(); + const ArrayRef<int64_t> resultShape = resultType.getShape(); - Value extract = - rewriter.create<vector::ExtractOp>(loc, op.getSource(), srcIdx); - result = rewriter.create<vector::InsertStridedSliceOp>( - loc, extract, result, - /*offsets=*/resIdx, /*strides=*/1); - } + const int64_t nSlices = + std::accumulate(sourceShape.begin(), sourceShape.begin() + sourceDim, 1, + std::multiplies<int64_t>()); - rewriter.replaceOp(op, result); - return success(); - } -}; + SmallVector<int64_t> extractIndex(sourceDim, 0); + SmallVector<int64_t> insertIndex(resultDim, 0); + Value result = rewriter.create<ub::PoisonOp>(loc, resultType); -/// ShapeOp 1-D -> n-D upcast serves the purpose of unflattening n-D from 1-D -/// vectors progressively. This iterates over the n-1 major dimension of the n-D -/// vector and performs rewrites into: -/// vector.extract_strided_slice from 1-D + vector.insert into n-D -/// Note that 1-D extract_strided_slice are lowered to efficient vector.shuffle. -class ShapeCastOpNDUpCastRewritePattern - : public OpRewritePattern<vector::ShapeCastOp> { -public: - using OpRewritePattern::OpRewritePattern; + for (int i = 0; i < nSlices; ++i) { + Value extracted = + rewriter.create<vector::ExtractOp>(loc, source, extractIndex); - LogicalResult matchAndRewrite(vector::ShapeCastOp op, - PatternRewriter &rewriter) const override { - auto sourceVectorType = op.getSourceVectorType(); - auto resultVectorType = op.getResultVectorType(); - if (sourceVectorType.isScalable() || resultVectorType.isScalable()) - return failure(); - - int64_t srcRank = sourceVectorType.getRank(); - int64_t resRank = resultVectorType.getRank(); - if (srcRank != 1 || resRank < 2) - return failure(); - - // Compute the number of 1-D vector elements involved in the reshape. - int64_t numElts = 1; - for (int64_t dim = 0; dim < resRank - 1; ++dim) - numElts *= resultVectorType.getDimSize(dim); - - // Compute the indices of each 1-D vector element of the source slice - // extraction and destination insertion and generate such instructions. - auto loc = op.getLoc(); - SmallVector<int64_t> srcIdx(srcRank, 0); - SmallVector<int64_t> resIdx(resRank - 1, 0); - int64_t extractSize = resultVectorType.getShape().back(); - Value result = rewriter.create<ub::PoisonOp>(loc, resultVectorType); - for (int64_t i = 0; i < numElts; ++i) { - if (i != 0) { - incIdx(srcIdx, sourceVectorType, /*step=*/extractSize); - incIdx(resIdx, resultVectorType, /*step=*/1); - } + result = rewriter.create<vector::InsertOp>(loc, extracted, result, + insertIndex); - Value extract = rewriter.create<vector::ExtractStridedSliceOp>( - loc, op.getSource(), /*offsets=*/srcIdx, /*sizes=*/extractSize, - /*strides=*/1); - result = rewriter.create<vector::InsertOp>(loc, extract, result, resIdx); + inplaceAdd(1, sourceShape.take_front(sourceDim), extractIndex); + inplaceAdd(1, resultShape.take_front(resultDim), insertIndex); } - rewriter.replaceOp(op, result); + rewriter.replaceOp(shapeCast, result); return success(); } -}; -// We typically should not lower general shape cast operations into data -// movement instructions, since the assumption is that these casts are -// optimized away during progressive lowering. For completeness, however, -// we fall back to a reference implementation that moves all elements -// into the right place if we get here. -class ShapeCastOpRewritePattern : public OpRewritePattern<vector::ShapeCastOp> { public: using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(vector::ShapeCastOp op, PatternRewriter &rewriter) const override { Location loc = op.getLoc(); - auto sourceVectorType = op.getSourceVectorType(); - auto resultVectorType = op.getResultVectorType(); + VectorType sourceType = op.getSourceVectorType(); + VectorType resultType = op.getResultVectorType(); + + if (sourceType.isScalable() || resultType.isScalable()) + return rewriter.notifyMatchFailure( + op, + "shape_cast where vectors are scalable not handled by this pattern"); + + const ArrayRef<int64_t> sourceShape = sourceType.getShape(); + const ArrayRef<int64_t> resultShape = resultType.getShape(); + const int64_t sourceRank = sourceType.getRank(); + const int64_t resultRank = resultType.getRank(); + const int64_t numElms = sourceType.getNumElements(); + const Value source = op.getSource(); + + // Set the first dimension (starting at the end) in the source and result + // respectively where the dimension sizes differ. Using the running example: + // + // dimensions: [0 1 2 3 4 5 ] [0 1 2 3 ] + // shapes: (2,2,3,4,7,11) -> (8,6,7,11) + // ^ ^ + // | | + // sourceSuffixStartDim is 3 | + // | + // resultSuffixStartDim is 1 + int64_t sourceSuffixStartDim = sourceRank - 1; + int64_t resultSuffixStartDim = resultRank - 1; + while (sourceSuffixStartDim >= 0 && resultSuffixStartDim >= 0 && + (sourceType.getDimSize(sourceSuffixStartDim) == + resultType.getDimSize(resultSuffixStartDim))) { + --sourceSuffixStartDim; + --resultSuffixStartDim; + } - if (sourceVectorType.isScalable() || resultVectorType.isScalable()) - return failure(); - - // Special case for n-D / 1-D lowerings with better implementations. - int64_t srcRank = sourceVectorType.getRank(); - int64_t resRank = resultVectorType.getRank(); - if ((srcRank > 1 && resRank == 1) || (srcRank == 1 && resRank > 1)) - return failure(); - - // Generic ShapeCast lowering path goes all the way down to unrolled scalar - // extract/insert chains. - int64_t numElts = 1; - for (int64_t r = 0; r < srcRank; r++) - numElts *= sourceVectorType.getDimSize(r); - // Replace with data movement operations: - // x[0,0,0] = y[0,0] - // x[0,0,1] = y[0,1] - // x[0,1,0] = y[0,2] - // etc., incrementing the two index vectors "row-major" - // within the source and result shape. - SmallVector<int64_t> srcIdx(srcRank, 0); - SmallVector<int64_t> resIdx(resRank, 0); - Value result = rewriter.create<ub::PoisonOp>(loc, resultVectorType); - for (int64_t i = 0; i < numElts; i++) { - if (i != 0) { - incIdx(srcIdx, sourceVectorType); - incIdx(resIdx, resultVectorType); + // This is the case (i) where there are just some leading ones to contend + // with in the source or result. It can be handled with a single + // extract/insert pair. + if (resultSuffixStartDim < 0 || sourceSuffixStartDim < 0) + return leadingOnesLowering(op, rewriter); + + const int64_t sourceSuffixStartDimSize = + sourceType.getDimSize(sourceSuffixStartDim); + const int64_t resultSuffixStartDimSize = + resultType.getDimSize(resultSuffixStartDim); + const int64_t greatestCommonDivisor = + std::gcd(sourceSuffixStartDimSize, resultSuffixStartDimSize); + const int64_t stridedSliceRank = sourceRank - sourceSuffixStartDim; + const size_t extractPeriod = + sourceSuffixStartDimSize / greatestCommonDivisor; + const size_t insertPeriod = + resultSuffixStartDimSize / greatestCommonDivisor; + + SmallVector<int64_t> atomicShape(sourceShape.begin() + sourceSuffixStartDim, + sourceShape.end()); + atomicShape[0] = greatestCommonDivisor; + + const int64_t numAtomicElms = std::accumulate( + atomicShape.begin(), atomicShape.end(), 1, std::multiplies<int64_t>()); + const size_t nAtomicSlices = numElms / numAtomicElms; + + // This is the case (ii) where the strided dimension size is 1. More compact + // IR is generated in this case if we just extract and insert the elements + // directly. In other words, we don't use extract_strided_slice and + // insert_strided_slice. + if (greatestCommonDivisor == 1) + return noStridedSliceLowering(op, sourceSuffixStartDim + 1, + resultSuffixStartDim + 1, rewriter); + + // The insert_strided_slice result's type + const ArrayRef<int64_t> insertStridedShape = + resultShape.drop_front(resultSuffixStartDim); + const VectorType insertStridedType = + VectorType::get(insertStridedShape, resultType.getElementType()); + + SmallVector<int64_t> extractIndex(sourceSuffixStartDim, 0); + SmallVector<int64_t> insertIndex(resultSuffixStartDim, 0); + SmallVector<int64_t> extractOffsets(stridedSliceRank, 0); + SmallVector<int64_t> insertOffsets(stridedSliceRank, 0); + const SmallVector<int64_t> sizes(stridedSliceRank, 1); + + Value extracted = {}; + Value extractedStrided = {}; + Value insertedSlice = {}; + Value result = rewriter.create<ub::PoisonOp>(loc, resultType); + const Value partResult = + rewriter.create<ub::PoisonOp>(loc, insertStridedType); + + for (size_t i = 0; i < nAtomicSlices; ++i) { + + const size_t extractStridedPhase = i % extractPeriod; + const size_t insertStridedPhase = i % insertPeriod; + + // vector.extract + if (extractStridedPhase == 0) { + extracted = + rewriter.create<vector::ExtractOp>(loc, source, extractIndex); + inplaceAdd(1, sourceShape.take_front(sourceSuffixStartDim), + extractIndex); } - Value extract = - rewriter.create<vector::ExtractOp>(loc, op.getSource(), srcIdx); - result = rewriter.create<vector::InsertOp>(loc, extract, result, resIdx); + // vector.extract_strided_slice + extractOffsets[0] = extractStridedPhase * greatestCommonDivisor; + extractedStrided = rewriter.create<vector::ExtractStridedSliceOp>( + loc, extracted, extractOffsets, atomicShape, sizes); + + // vector.insert_strided_slice + if (insertStridedPhase == 0) { + insertedSlice = partResult; + } + insertOffsets[0] = insertStridedPhase * greatestCommonDivisor; + insertedSlice = rewriter.create<vector::InsertStridedSliceOp>( + loc, extractedStrided, insertedSlice, insertOffsets, sizes); + + // vector.insert + if (insertStridedPhase + 1 == insertPeriod) { + result = rewriter.create<vector::InsertOp>(loc, insertedSlice, result, + insertIndex); + inplaceAdd(1, resultType.getShape().take_front(resultSuffixStartDim), + insertIndex); + } } rewriter.replaceOp(op, result); return success(); @@ -252,7 +373,8 @@ public: // from >= 2-D scalable vectors or scalable vectors of fixed vectors. if (!isTrailingDimScalable(sourceVectorType) || !isTrailingDimScalable(resultVectorType)) { - return failure(); + return rewriter.notifyMatchFailure( + op, "trailing dims are not scalable, not handled by this pattern"); } // The sizes of the trailing dimension of the source and result vectors, the @@ -329,8 +451,8 @@ public: // 4. Increment the insert/extract indices, stepping by minExtractionSize // for the trailing dimensions. - incIdx(srcIdx, sourceVectorType, /*step=*/minExtractionSize); - incIdx(resIdx, resultVectorType, /*step=*/minExtractionSize); + inplaceAdd(minExtractionSize, sourceVectorType.getShape(), srcIdx); + inplaceAdd(minExtractionSize, resultVectorType.getShape(), resIdx); } rewriter.replaceOp(op, result); @@ -347,8 +469,6 @@ public: void mlir::vector::populateVectorShapeCastLoweringPatterns( RewritePatternSet &patterns, PatternBenefit benefit) { - patterns.add<ShapeCastOpNDDownCastRewritePattern, - ShapeCastOpNDUpCastRewritePattern, ShapeCastOpRewritePattern, - ScalableShapeCastOpRewritePattern>(patterns.getContext(), - benefit); + patterns.add<ShapeCastOpRewritePattern, ScalableShapeCastOpRewritePattern>( + patterns.getContext(), benefit); } diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp index 1cc477d9dca9..fc443ab0d138 100644 --- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp @@ -631,14 +631,78 @@ private: vector::UnrollVectorOptions options; }; +struct UnrollBroadcastPattern : public OpRewritePattern<vector::BroadcastOp> { + UnrollBroadcastPattern(MLIRContext *context, + const vector::UnrollVectorOptions &options, + PatternBenefit benefit = 1) + : OpRewritePattern<vector::BroadcastOp>(context, benefit), + options(options) {} + + LogicalResult matchAndRewrite(vector::BroadcastOp broadcastOp, + PatternRewriter &rewriter) const override { + auto targetShape = getTargetShape(options, broadcastOp); + if (!targetShape) + return failure(); + + Location loc = broadcastOp.getLoc(); + VectorType srcType = dyn_cast<VectorType>(broadcastOp.getSourceType()); + VectorType resType = broadcastOp.getResultVectorType(); + VectorType targetType = + resType.cloneWith(*targetShape, resType.getElementType()); + Value result = rewriter.create<arith::ConstantOp>( + loc, resType, rewriter.getZeroAttr(resType)); + + SmallVector<int64_t> originalShape = *broadcastOp.getShapeForUnroll(); + SmallVector<int64_t> strides(originalShape.size(), 1); + + for (SmallVector<int64_t> offsets : + StaticTileOffsetRange(originalShape, *targetShape)) { + Value newSrc; + if (!srcType) { + // Scalar to vector broadcast. + newSrc = broadcastOp.getSource(); + } else { + // Vector to vector broadcast. + int64_t rank = srcType.getRank(); + SmallVector<int64_t> srcOffsets(offsets.end() - rank, offsets.end()); + SmallVector<int64_t> srcShape(targetShape->end() - rank, + targetShape->end()); + SmallVector<int64_t> srcStrides(strides.end() - rank, strides.end()); + // adjust the offset and shape for src if the corresponding dim is 1. + for (int64_t i = 0; i < rank; ++i) { + if (srcType.getDimSize(i) == 1) { + srcOffsets[i] = 0; + srcShape[i] = 1; + } + } + newSrc = rewriter.createOrFold<vector::ExtractStridedSliceOp>( + loc, broadcastOp.getSource(), srcOffsets, srcShape, srcStrides); + } + + Operation *newOp = cloneOpWithOperandsAndTypes(rewriter, loc, broadcastOp, + newSrc, targetType); + + result = rewriter.createOrFold<vector::InsertStridedSliceOp>( + loc, newOp->getResult(0), result, offsets, strides); + } + + rewriter.replaceOp(broadcastOp, result); + return success(); + } + +private: + vector::UnrollVectorOptions options; +}; + } // namespace void mlir::vector::populateVectorUnrollPatterns( RewritePatternSet &patterns, const UnrollVectorOptions &options, PatternBenefit benefit) { - patterns.add<UnrollTransferReadPattern, UnrollTransferWritePattern, - UnrollContractionPattern, UnrollElementwisePattern, - UnrollReductionPattern, UnrollMultiReductionPattern, - UnrollTransposePattern, UnrollGatherPattern>( - patterns.getContext(), options, benefit); + patterns + .add<UnrollTransferReadPattern, UnrollTransferWritePattern, + UnrollContractionPattern, UnrollElementwisePattern, + UnrollReductionPattern, UnrollMultiReductionPattern, + UnrollTransposePattern, UnrollGatherPattern, UnrollBroadcastPattern>( + patterns.getContext(), options, benefit); } diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp index d5dd6f2027be..590d244daef4 100644 --- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp +++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp @@ -337,15 +337,16 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc, auto sourceShape = sourceShapedType.getShape(); assert(sourceShape.size() == inputVectorSizes.size() && "expected same ranks."); - auto maskType = VectorType::get(inputVectorSizes, builder.getI1Type()); auto vectorType = VectorType::get(inputVectorSizes, padValue.getType()); assert(padValue.getType() == sourceShapedType.getElementType() && "expected same pad element type to match source element type"); int64_t readRank = inputVectorSizes.size(); auto zero = builder.create<arith::ConstantIndexOp>(loc, 0); SmallVector<bool> inBoundsVal(readRank, true); + if (useInBoundsInsteadOfMasking) { // Update the inBounds attribute. + // FIXME: This computation is too weak - it ignores the read indices. for (unsigned i = 0; i < readRank; i++) inBoundsVal[i] = (sourceShape[i] == inputVectorSizes[i]) && !ShapedType::isDynamic(sourceShape[i]); @@ -362,6 +363,8 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc, return transferReadOp; SmallVector<OpFoldResult> mixedSourceDims = tensor::getMixedSizes(builder, loc, source); + + auto maskType = VectorType::get(inputVectorSizes, builder.getI1Type()); Value mask = builder.create<vector::CreateMaskOp>(loc, maskType, mixedSourceDims); return mlir::vector::maskOperation(builder, transferReadOp, mask) diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index 39f5c003320d..f704fbfbe8ff 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -315,7 +315,8 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, // process and dynamically linked libraries. auto objectLinkingLayerCreator = [&](ExecutionSession &session) { auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>( - session, [sectionMemoryMapper = options.sectionMemoryMapper]() { + session, [sectionMemoryMapper = + options.sectionMemoryMapper](const MemoryBuffer &) { return std::make_unique<SectionMemoryManager>(sectionMemoryMapper); }); diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp index 7c9e6c89d4d8..d3e9aebc4dc2 100644 --- a/mlir/lib/IR/OperationSupport.cpp +++ b/mlir/lib/IR/OperationSupport.cpp @@ -680,9 +680,13 @@ llvm::hash_code OperationEquivalence::computeHash( // - Operation Name // - Attributes // - Result Types + DictionaryAttr dictAttrs; + if (!(flags & Flags::IgnoreDiscardableAttrs)) + dictAttrs = op->getRawDictionaryAttrs(); llvm::hash_code hash = - llvm::hash_combine(op->getName(), op->getRawDictionaryAttrs(), - op->getResultTypes(), op->hashProperties()); + llvm::hash_combine(op->getName(), dictAttrs, op->getResultTypes()); + if (!(flags & Flags::IgnoreProperties)) + hash = llvm::hash_combine(hash, op->hashProperties()); // - Location if required if (!(flags & Flags::IgnoreLocations)) @@ -836,14 +840,19 @@ OperationEquivalence::isRegionEquivalentTo(Region *lhs, Region *rhs, return true; // 1. Compare the operation properties. + if (!(flags & IgnoreDiscardableAttrs) && + lhs->getRawDictionaryAttrs() != rhs->getRawDictionaryAttrs()) + return false; + if (lhs->getName() != rhs->getName() || - lhs->getRawDictionaryAttrs() != rhs->getRawDictionaryAttrs() || lhs->getNumRegions() != rhs->getNumRegions() || lhs->getNumSuccessors() != rhs->getNumSuccessors() || lhs->getNumOperands() != rhs->getNumOperands() || - lhs->getNumResults() != rhs->getNumResults() || - !lhs->getName().compareOpProperties(lhs->getPropertiesStorage(), - rhs->getPropertiesStorage())) + lhs->getNumResults() != rhs->getNumResults()) + return false; + if (!(flags & IgnoreProperties) && + !(lhs->getName().compareOpProperties(lhs->getPropertiesStorage(), + rhs->getPropertiesStorage()))) return false; if (!(flags & IgnoreLocations) && lhs->getLoc() != rhs->getLoc()) return false; diff --git a/mlir/lib/Interfaces/ViewLikeInterface.cpp b/mlir/lib/Interfaces/ViewLikeInterface.cpp index 70dd7b4aec88..d70cfebbe53e 100644 --- a/mlir/lib/Interfaces/ViewLikeInterface.cpp +++ b/mlir/lib/Interfaces/ViewLikeInterface.cpp @@ -27,9 +27,7 @@ LogicalResult mlir::verifyListOfOperandsOrIntegers(Operation *op, return op->emitError("expected ") << numElements << " " << name << " values, got " << staticVals.size(); unsigned expectedNumDynamicEntries = - llvm::count_if(staticVals, [](int64_t staticVal) { - return ShapedType::isDynamic(staticVal); - }); + llvm::count_if(staticVals, ShapedType::isDynamic); if (values.size() != expectedNumDynamicEntries) return op->emitError("expected ") << expectedNumDynamicEntries << " dynamic " << name << " values"; @@ -270,5 +268,5 @@ bool mlir::detail::sameOffsetsSizesAndStrides( unsigned mlir::detail::getNumDynamicEntriesUpToIdx(ArrayRef<int64_t> staticVals, unsigned idx) { return std::count_if(staticVals.begin(), staticVals.begin() + idx, - [&](int64_t val) { return ShapedType::isDynamic(val); }); + ShapedType::isDynamic); } diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp index e0e9b5f54042..d1371d2b508f 100644 --- a/mlir/lib/Pass/Pass.cpp +++ b/mlir/lib/Pass/Pass.cpp @@ -18,6 +18,7 @@ #include "mlir/IR/Threading.h" #include "mlir/IR/Verifier.h" #include "mlir/Support/FileUtilities.h" +#include "mlir/Support/IndentedOstream.h" #include "llvm/ADT/Hashing.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/ScopeExit.h" @@ -80,14 +81,19 @@ void Pass::copyOptionValuesFrom(const Pass *other) { } /// Prints out the pass in the textual representation of pipelines. If this is -/// an adaptor pass, print its pass managers. -void Pass::printAsTextualPipeline(raw_ostream &os) { +/// an adaptor pass, print its pass managers. When `pretty` is true, the +/// printed pipeline is formatted for readability. +void Pass::printAsTextualPipeline(raw_ostream &os, bool pretty) { // Special case for adaptors to print its pass managers. if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(this)) { llvm::interleave( adaptor->getPassManagers(), - [&](OpPassManager &pm) { pm.printAsTextualPipeline(os); }, - [&] { os << ","; }); + [&](OpPassManager &pm) { pm.printAsTextualPipeline(os, pretty); }, + [&] { + os << ","; + if (pretty) + os << "\n"; + }); return; } // Otherwise, print the pass argument followed by its options. If the pass @@ -390,27 +396,51 @@ StringRef OpPassManager::getOpAnchorName() const { } /// Prints out the passes of the pass manager as the textual representation -/// of pipelines. +/// of pipelines. When `pretty` is true, the printed pipeline is formatted for +/// readability. void printAsTextualPipeline( - raw_ostream &os, StringRef anchorName, - const llvm::iterator_range<OpPassManager::pass_iterator> &passes) { + raw_indented_ostream &os, StringRef anchorName, + const llvm::iterator_range<OpPassManager::pass_iterator> &passes, + bool pretty = false) { os << anchorName << "("; + if (pretty) { + os << "\n"; + os.indent(); + } llvm::interleave( - passes, [&](mlir::Pass &pass) { pass.printAsTextualPipeline(os); }, - [&]() { os << ","; }); + passes, + [&](mlir::Pass &pass) { pass.printAsTextualPipeline(os, pretty); }, + [&]() { + os << ","; + if (pretty) + os << "\n"; + }); + if (pretty) { + os << "\n"; + os.unindent(); + } os << ")"; } -void OpPassManager::printAsTextualPipeline(raw_ostream &os) const { +void printAsTextualPipeline( + raw_ostream &os, StringRef anchorName, + const llvm::iterator_range<OpPassManager::pass_iterator> &passes, + bool pretty) { + raw_indented_ostream indentedOS(os); + printAsTextualPipeline(indentedOS, anchorName, passes, pretty); +} +void OpPassManager::printAsTextualPipeline(raw_ostream &os, bool pretty) const { StringRef anchorName = getOpAnchorName(); + raw_indented_ostream indentedOS(os); ::printAsTextualPipeline( - os, anchorName, + indentedOS, anchorName, {MutableArrayRef<std::unique_ptr<Pass>>{impl->passes}.begin(), - MutableArrayRef<std::unique_ptr<Pass>>{impl->passes}.end()}); + MutableArrayRef<std::unique_ptr<Pass>>{impl->passes}.end()}, + pretty); } void OpPassManager::dump() { llvm::errs() << "Pass Manager with " << impl->passes.size() << " passes:\n"; - printAsTextualPipeline(llvm::errs()); + printAsTextualPipeline(llvm::errs(), /*pretty=*/true); llvm::errs() << "\n"; } @@ -466,7 +496,6 @@ llvm::hash_code OpPassManager::hash() { return hashCode; } - //===----------------------------------------------------------------------===// // OpToOpPassAdaptor //===----------------------------------------------------------------------===// @@ -871,7 +900,8 @@ LogicalResult PassManager::run(Operation *op) { // Initialize all of the passes within the pass manager with a new generation. llvm::hash_code newInitKey = context->getRegistryHash(); llvm::hash_code pipelineKey = hash(); - if (newInitKey != initializationKey || pipelineKey != pipelineInitializationKey) { + if (newInitKey != initializationKey || + pipelineKey != pipelineInitializationKey) { if (failed(initialize(context, impl->initializationGeneration + 1))) return failure(); initializationKey = newInitKey; diff --git a/mlir/lib/Pass/PassCrashRecovery.cpp b/mlir/lib/Pass/PassCrashRecovery.cpp index 8c6d865cb31d..b048ff946239 100644 --- a/mlir/lib/Pass/PassCrashRecovery.cpp +++ b/mlir/lib/Pass/PassCrashRecovery.cpp @@ -443,7 +443,8 @@ makeReproducerStreamFactory(StringRef outputFile) { void printAsTextualPipeline( raw_ostream &os, StringRef anchorName, - const llvm::iterator_range<OpPassManager::pass_iterator> &passes); + const llvm::iterator_range<OpPassManager::pass_iterator> &passes, + bool pretty = false); std::string mlir::makeReproducer( StringRef anchorName, diff --git a/mlir/lib/Reducer/OptReductionPass.cpp b/mlir/lib/Reducer/OptReductionPass.cpp index fbbe75d30388..7eb29a02702a 100644 --- a/mlir/lib/Reducer/OptReductionPass.cpp +++ b/mlir/lib/Reducer/OptReductionPass.cpp @@ -88,4 +88,4 @@ void OptReductionPass::runOnOperation() { moduleVariant->destroy(); LLVM_DEBUG(llvm::dbgs() << "Pass Complete\n\n"); -}
\ No newline at end of file +} diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp index 4a6605e9c74c..ec1a9c909c94 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp @@ -5737,7 +5737,13 @@ convertHostOrTargetOperation(Operation *op, llvm::IRBuilderBase &builder, llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterIP = ompBuilder->createBarrier(builder.saveIP(), llvm::omp::OMPD_barrier); - return handleError(afterIP, *op); + LogicalResult res = handleError(afterIP, *op); + if (res.succeeded()) { + // If the barrier generated a cancellation check, the insertion + // point might now need to be changed to a new continuation block + builder.restoreIP(*afterIP); + } + return res; }) .Case([&](omp::TaskyieldOp op) { if (failed(checkImplementationStatus(*op))) diff --git a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp index a21d691ae514..3957dbc0db98 100644 --- a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp +++ b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp @@ -1062,12 +1062,30 @@ LogicalResult spirv::Deserializer::processCooperativeMatrixTypeKHR( << operands[2]; } - unsigned rows = getConstantInt(operands[3]).getInt(); - unsigned columns = getConstantInt(operands[4]).getInt(); + IntegerAttr rowsAttr = getConstantInt(operands[3]); + IntegerAttr columnsAttr = getConstantInt(operands[4]); + IntegerAttr useAttr = getConstantInt(operands[5]); + + if (!rowsAttr) + return emitError(unknownLoc, "OpTypeCooperativeMatrixKHR `Rows` references " + "undefined constant <id> ") + << operands[3]; + + if (!columnsAttr) + return emitError(unknownLoc, "OpTypeCooperativeMatrixKHR `Columns` " + "references undefined constant <id> ") + << operands[4]; + + if (!useAttr) + return emitError(unknownLoc, "OpTypeCooperativeMatrixKHR `Use` references " + "undefined constant <id> ") + << operands[5]; + + unsigned rows = rowsAttr.getInt(); + unsigned columns = columnsAttr.getInt(); std::optional<spirv::CooperativeMatrixUseKHR> use = - spirv::symbolizeCooperativeMatrixUseKHR( - getConstantInt(operands[5]).getInt()); + spirv::symbolizeCooperativeMatrixUseKHR(useAttr.getInt()); if (!use) { return emitError( unknownLoc, |