NFC: improve naming and doc for tiled and distributed loop info (#7558) The logic behind rediscovering the loop tiling and distribution information is quite dense. This makes at least the API clearer.
diff --git a/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp index 6b2ad8e..06b12bc 100644 --- a/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp +++ b/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -130,7 +130,7 @@ } static SmallVector<int64_t> getDefaultWorkloadPerWorkgroup( - ArrayRef<TiledLoopInfo> tiledLoops, + ArrayRef<LoopTilingAndDistributionInfo> tiledLoops, ArrayRef<int64_t> nativeVectorSizeInElements) { if (tiledLoops.empty()) { return {}; @@ -138,7 +138,7 @@ assert(tiledLoops.size() == nativeVectorSizeInElements.size()); unsigned maxDim = 0; for (auto tiledLoop : tiledLoops) { - maxDim = std::max<unsigned>(tiledLoop.distributionDim, maxDim); + maxDim = std::max<unsigned>(tiledLoop.processorDistributionDim, maxDim); } SmallVector<int64_t> workloadPerWorkgroup(maxDim + 1, 1); SmallVector<int64_t> numWorkgroupsPerDim(maxDim + 1, 1); @@ -149,9 +149,9 @@ auto ceilFn = [](int64_t a, int64_t b) { return (a + b - 1) / b; }; for (auto tiledLoop : enumerate(tiledLoops)) { - Optional<int64_t> lb = getStaticValue(tiledLoop.value().lb); - Optional<int64_t> ub = getStaticValue(tiledLoop.value().ub); - unsigned dim = tiledLoop.value().distributionDim; + Optional<int64_t> lb = getStaticValue(tiledLoop.value().untiledLowerBound); + Optional<int64_t> ub = getStaticValue(tiledLoop.value().untiledUpperBound); + unsigned dim = tiledLoop.value().processorDistributionDim; if (!lb || !ub) { workloadPerWorkgroup[dim] = defaultWorkgroupTileSize; workload[dim] = ShapedType::kDynamicSize; @@ -209,7 +209,7 @@ /// Sets the default launch configuration to use for a tiled + distributed /// dispatch region based on the `tiledLoops` found. static LogicalResult setDefaultLaunchConfig( - FuncOp entryPointFn, ArrayRef<TiledLoopInfo> tiledLoops) { + FuncOp entryPointFn, ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) { unsigned typeWidthInBytes = getReferenceTypeLengthInBytes(entryPointFn); SmallVector<int64_t> nativeVectorSizeInElements(tiledLoops.size(), 1); if (!tiledLoops.empty()) { @@ -232,9 +232,9 @@ /// Sets the lowering configuration for dispatch region with root op that /// implements the contraction operation interface. -static LogicalResult setRootConfig(FuncOp entryPointFn, - linalg::ContractionOpInterface contractionOp, - ArrayRef<TiledLoopInfo> tiledLoops) { +static LogicalResult setRootConfig( + FuncOp entryPointFn, linalg::ContractionOpInterface contractionOp, + ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) { if (getLoweringConfig(contractionOp)) return success(); auto lhsShapedType = contractionOp.lhs().getType().cast<ShapedType>(); @@ -286,15 +286,17 @@ return maxSize; }; for (unsigned i = tiledLoops.size() - 2; i < tiledLoops.size(); ++i) { - if (!tiledLoops[i].lb.is<Attribute>() || - !tiledLoops[i].ub.is<Attribute>()) { + if (!tiledLoops[i].untiledLowerBound.is<Attribute>() || + !tiledLoops[i].untiledUpperBound.is<Attribute>()) { continue; } - int64_t lb = tiledLoops[i].lb.get<Attribute>().cast<IntegerAttr>().getInt(); - int64_t ub = tiledLoops[i].ub.get<Attribute>().cast<IntegerAttr>().getInt(); - workloadPerWorkgroup[tiledLoops.size() - 1 - i] = - getTileSize(lb, ub, workloadPerWorkgroup[tiledLoops.size() - 1 - i], - vectorSizeVals[i]); + auto lb = + tiledLoops[i].untiledLowerBound.get<Attribute>().cast<IntegerAttr>(); + auto ub = + tiledLoops[i].untiledUpperBound.get<Attribute>().cast<IntegerAttr>(); + workloadPerWorkgroup[tiledLoops.size() - 1 - i] = getTileSize( + lb.getInt(), ub.getInt(), + workloadPerWorkgroup[tiledLoops.size() - 1 - i], vectorSizeVals[i]); } setTranslationInfo( entryPointFn, @@ -328,8 +330,9 @@ /// Sets the lowering configuration for dispatch region for linalg.mmt4d root /// op -static LogicalResult setRootConfig(FuncOp entryPointFn, linalg::Mmt4DOp mmt4dOp, - ArrayRef<TiledLoopInfo> tiledLoops) { +static LogicalResult setRootConfig( + FuncOp entryPointFn, linalg::Mmt4DOp mmt4dOp, + ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) { // TODO(ataei): These are hand tuned for some performance benchmarks for // now, we want to adapt the same strategy as matmul that dynamically sets // tile size. @@ -379,8 +382,9 @@ /// Sets the lowering configuration for dispatch region for linalg_ext.fft /// root op. -static LogicalResult setRootConfig(FuncOp entryPointFn, linalg_ext::FftOp fftOp, - ArrayRef<TiledLoopInfo> tiledLoops) { +static LogicalResult setRootConfig( + FuncOp entryPointFn, linalg_ext::FftOp fftOp, + ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) { auto partitionedLoops = getPartitionedLoops(fftOp); unsigned maxDepth = partitionedLoops.back() + 1; SmallVector<int64_t> workgroupTileSizes(maxDepth, defaultWorkgroupTileSize); @@ -415,9 +419,9 @@ /// Finds the root operation in the given list of linalg operations and sets /// its configuration. Returns error for multiple root operations. -static LogicalResult setRootConfig(FuncOp entryPointFn, - ArrayRef<Operation *> computeOps, - ArrayRef<TiledLoopInfo> tiledLoops) { +static LogicalResult setRootConfig( + FuncOp entryPointFn, ArrayRef<Operation *> computeOps, + ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) { Operation *rootOp = nullptr; for (auto computeOp : computeOps) { auto setRootConfigFn = [&](Operation *op) -> LogicalResult { @@ -445,7 +449,7 @@ /// Sets the translation information to use for a dispatch region. static LogicalResult setTranslationInfoAndRootConfig( FuncOp entryPointFn, ArrayRef<Operation *> computeOps, - ArrayRef<TiledLoopInfo> tiledLoops) { + ArrayRef<LoopTilingAndDistributionInfo> tiledLoops) { // First check if the operations have a preset pipeline. for (auto computeOp : computeOps) { if (IREE::Codegen::CompilationInfoAttr compilationInfo = @@ -486,7 +490,7 @@ if (!entryPointOp) continue; if (getTranslationInfo(entryPointOp)) continue; SmallVector<Operation *> computeOps; - SmallVector<TiledLoopInfo> tiledLoops; + SmallVector<LoopTilingAndDistributionInfo> tiledLoops; // If there are no linalg ops, not using Linalg based lowering. if (failed(getComputeOps(funcOp, computeOps, tiledLoops))) {
diff --git a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp index 0d1ae4e..d6bda1d 100644 --- a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp +++ b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -263,7 +263,7 @@ if (!entryPointOp) continue; if (getTranslationInfo(entryPointOp)) continue; SmallVector<Operation *> computeOps; - SmallVector<TiledLoopInfo> tiledLoops; + SmallVector<LoopTilingAndDistributionInfo> tiledLoops; if (failed(getComputeOps(funcOp, computeOps, tiledLoops))) { return funcOp.emitOpError("failed to get compute ops"); }
diff --git a/iree/compiler/Codegen/LLVMGPU/LLVMGPURemoveTrivialLoops.cpp b/iree/compiler/Codegen/LLVMGPU/LLVMGPURemoveTrivialLoops.cpp index df0e622..5837ff4 100644 --- a/iree/compiler/Codegen/LLVMGPU/LLVMGPURemoveTrivialLoops.cpp +++ b/iree/compiler/Codegen/LLVMGPU/LLVMGPURemoveTrivialLoops.cpp
@@ -60,20 +60,27 @@ /// of element per workgroups. static SmallVector<int64_t> getNumWorkgroup( FuncOp funcOp, IREE::HAL::ExecutableEntryPointOp entryPointOp) { - SmallVector<TiledLoopInfo> tiledLoopInfo = getTiledLoopInfo(funcOp); + SmallVector<LoopTilingAndDistributionInfo> tiledLoopInfo = + getTiledAndDistributedLoopInfo(funcOp); SmallVector<int64_t> workloadSize(tiledLoopInfo.size()); - for (TiledLoopInfo &tileInfo : tiledLoopInfo) { - if (tileInfo.distributionDim >= workloadSize.size()) + for (LoopTilingAndDistributionInfo &tileInfo : tiledLoopInfo) { + if (tileInfo.processorDistributionDim >= workloadSize.size()) return SmallVector<int64_t>(); - if (!tileInfo.lb.is<Attribute>() || !tileInfo.ub.is<Attribute>() || - !tileInfo.step.is<Attribute>()) { + if (!tileInfo.untiledLowerBound.is<Attribute>() || + !tileInfo.untiledUpperBound.is<Attribute>() || + !tileInfo.untiledStep.is<Attribute>()) { continue; } - int64_t lb = tileInfo.lb.get<Attribute>().cast<IntegerAttr>().getInt(); - int64_t ub = tileInfo.ub.get<Attribute>().cast<IntegerAttr>().getInt(); - int64_t step = tileInfo.step.get<Attribute>().cast<IntegerAttr>().getInt(); + int64_t lb = tileInfo.untiledLowerBound.get<Attribute>() + .cast<IntegerAttr>() + .getInt(); + int64_t ub = tileInfo.untiledUpperBound.get<Attribute>() + .cast<IntegerAttr>() + .getInt(); + int64_t step = + tileInfo.untiledStep.get<Attribute>().cast<IntegerAttr>().getInt(); if (step == 0) return SmallVector<int64_t>(); - workloadSize[tileInfo.distributionDim] = (ub - lb) / step; + workloadSize[tileInfo.processorDistributionDim] = (ub - lb) / step; } auto translationInfo = getTranslationInfo(entryPointOp); if (!translationInfo) return SmallVector<int64_t>();
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp index 3692cf6..1f59872 100644 --- a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp +++ b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -385,7 +385,8 @@ // 1) distributing to as many threads as possible, and 2) avoid assigning too // many threads to handle out-of-bound elements (thus idle). - SmallVector<TiledLoopInfo> tiledLoopInfo = getTiledLoopInfo(funcOp); + SmallVector<LoopTilingAndDistributionInfo> tiledLoopInfo = + getTiledAndDistributedLoopInfo(funcOp); // The number of linalg implicit loops to partition and tiled loops // surrounding the op should match. Otherwise, something is incorrect. assert(partitionedLoops.size() == tiledLoopInfo.size()); @@ -395,8 +396,9 @@ // tiledLoopInfo uses the reverse order of partitionedLoops. for (auto pair : llvm::zip(llvm::reverse(partitionedLoops), tiledLoopInfo)) { unsigned loopIndex = std::get<0>(pair); - const TiledLoopInfo &loopInfo = std::get<1>(pair); - Optional<int64_t> attrValue = getConstantIntValue(loopInfo.ub); + const LoopTilingAndDistributionInfo &loopInfo = std::get<1>(pair); + Optional<int64_t> attrValue = + getConstantIntValue(loopInfo.untiledUpperBound); if (attrValue) { loopBounds[loopIndex] = *attrValue; } else { @@ -624,7 +626,7 @@ if (getTranslationInfo(entryPointOp)) continue; SmallVector<Operation *> computeOps; - SmallVector<TiledLoopInfo> tiledLoops; + SmallVector<LoopTilingAndDistributionInfo> tiledLoops; if (failed(getComputeOps(funcOp, computeOps, tiledLoops))) { return funcOp.emitOpError("failed to get compute ops"); }
diff --git a/iree/compiler/Codegen/SPIRV/SPIRVRemoveOneTripTiledLoops.cpp b/iree/compiler/Codegen/SPIRV/SPIRVRemoveOneTripTiledLoops.cpp index 21b7e62..3a67961 100644 --- a/iree/compiler/Codegen/SPIRV/SPIRVRemoveOneTripTiledLoops.cpp +++ b/iree/compiler/Codegen/SPIRV/SPIRVRemoveOneTripTiledLoops.cpp
@@ -103,7 +103,7 @@ // This pass seems to be only needed for the convolution vectorization. So // filter out the necessary conv ops. SmallVector<Operation *> rootOp; - SmallVector<TiledLoopInfo> tiledLoops; + SmallVector<LoopTilingAndDistributionInfo> tiledLoops; auto isConvOp = [](Operation *op) { return isa<linalg::DepthwiseConv2DNhwOp, linalg::Conv2DNhwcHwcfOp>(op); };
diff --git a/iree/compiler/Codegen/Utils/Utils.cpp b/iree/compiler/Codegen/Utils/Utils.cpp index 96fa627..c371447 100644 --- a/iree/compiler/Codegen/Utils/Utils.cpp +++ b/iree/compiler/Codegen/Utils/Utils.cpp
@@ -12,10 +12,13 @@ #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallSet.h" #include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/SymbolTable.h" +#define DEBUG_TYPE "iree-codegen-utils" + namespace mlir { namespace iree_compiler { @@ -163,7 +166,7 @@ return isaAffineExprOfType<T2, T3...>(expr); } -/// Returns a Value that repreesnts the value for symbol or dim expr for the map +/// Returns a Value that represents the value for symbol or dim expr for the map /// in the `applyOp`. static Value getValueForDimOrSymbol(AffineApplyOp applyOp, AffineExpr expr) { unsigned numDims = applyOp.getAffineMap().getNumDims(); @@ -233,7 +236,8 @@ class LowerBoundExprVisitor : public AffineExprVisitor<LowerBoundExprVisitor, LogicalResult> { public: - LowerBoundExprVisitor(AffineApplyOp applyOp, TiledLoopInfo &loopInfo) + LowerBoundExprVisitor(AffineApplyOp applyOp, + LoopTilingAndDistributionInfo &loopInfo) : applyOp(applyOp), loopInfo(loopInfo) {} LogicalResult visitSymbolExpr(AffineSymbolExpr /*expr*/) { return failure(); } @@ -260,10 +264,10 @@ if (!v) { return failure(); } - loopInfo.lb = getAsOpFoldResult(v); + loopInfo.untiledLowerBound = getAsOpFoldResult(v); } else if (auto constExpr = lbExpr.dyn_cast<AffineConstantExpr>()) { - loopInfo.lb = IntegerAttr::get(IndexType::get(applyOp.getContext()), - constExpr.getValue()); + loopInfo.untiledLowerBound = IntegerAttr::get( + IndexType::get(applyOp.getContext()), constExpr.getValue()); } else { return failure(); } @@ -279,7 +283,7 @@ if (vals.size() != 1 || !vals[0]) { return failure(); } - loopInfo.workgroupSize = wgSize.getValue(); + loopInfo.tileSize = wgSize.getValue(); dimension = checkDimensions<IREE::HAL::InterfaceWorkgroupIDOp>(vals); } else { vals = getValuesForDimsOrSymbols(applyOp, {expr.getLHS(), expr.getRHS()}); @@ -292,16 +296,17 @@ if (!dimension) { return failure(); } - loopInfo.distributionDim = dimension.getValue(); - if (!loopInfo.lb) { - loopInfo.lb = IntegerAttr::get(IndexType::get(applyOp.getContext()), 0); + loopInfo.processorDistributionDim = dimension.getValue(); + if (!loopInfo.untiledLowerBound) { + loopInfo.untiledLowerBound = + IntegerAttr::get(IndexType::get(applyOp.getContext()), 0); } return success(); } private: AffineApplyOp applyOp; - TiledLoopInfo &loopInfo; + LoopTilingAndDistributionInfo &loopInfo; }; /// Visitor to walk the `step` of a distributed loop. Expected the expression to @@ -311,7 +316,8 @@ class StepExprVisitor : public AffineExprVisitor<StepExprVisitor, LogicalResult> { public: - StepExprVisitor(AffineApplyOp applyOp, TiledLoopInfo &loopInfo) + StepExprVisitor(AffineApplyOp applyOp, + LoopTilingAndDistributionInfo &loopInfo) : applyOp(applyOp), loopInfo(loopInfo) {} LogicalResult visitSymbolExpr(AffineSymbolExpr /*expr*/) { return failure(); } @@ -335,28 +341,28 @@ } expr = e.cast<AffineBinaryOpExpr>(); } else { - // Check if WorkgroupSizeOp is folded. - if (loopInfo.workgroupSize) { - if (auto stepBySize = expr.getRHS().dyn_cast<AffineConstantExpr>()) { - loopInfo.step = + // Check if the workgroup tile size is folded into the affine map itself. + if (loopInfo.tileSize) { + if (auto stepCst = expr.getRHS().dyn_cast<AffineConstantExpr>()) { + loopInfo.untiledStep = IntegerAttr::get(IndexType::get(applyOp.getContext()), - stepBySize.getValue() / *loopInfo.workgroupSize); + stepCst.getValue() / *loopInfo.tileSize); } } else { - loopInfo.step = + loopInfo.untiledStep = IntegerAttr::get(IndexType::get(applyOp.getContext()), 1); } } if (failed(processSentinel(expr.getLHS(), sentinels)) || - (!loopInfo.workgroupSize && + (!loopInfo.tileSize && failed(processSentinel(expr.getRHS(), sentinels)))) { return failure(); } // Either there are 3 sentinels and step isnt set, or there are two // sentinels and the step is set. if (sentinels.size() == 3) { - if (loopInfo.step) { + if (loopInfo.untiledStep) { return failure(); } auto it = sentinels.begin(); @@ -364,7 +370,7 @@ Value v = getValueForDimOrSymbol(applyOp, *it); if (!v.getDefiningOp<IREE::HAL::InterfaceWorkgroupSizeOp>() && !v.getDefiningOp<IREE::HAL::InterfaceWorkgroupCountOp>()) { - loopInfo.step = getAsOpFoldResult(v); + loopInfo.untiledStep = getAsOpFoldResult(v); break; } } @@ -373,18 +379,18 @@ } } - if ((sentinels.size() != 2 || !loopInfo.step) && - (sentinels.size() != 1 || !loopInfo.workgroupSize)) { + if ((sentinels.size() != 2 || !loopInfo.untiledStep) && + (sentinels.size() != 1 || !loopInfo.tileSize)) { return failure(); } SmallVector<Value> vals = getValuesForDimsOrSymbols(applyOp, sentinels); - if ((loopInfo.workgroupSize && + if ((loopInfo.tileSize && !checkDimensions<IREE::HAL::InterfaceWorkgroupCountOp>( - vals, loopInfo.distributionDim)) || - (!loopInfo.workgroupSize && + vals, loopInfo.processorDistributionDim)) || + (!loopInfo.tileSize && !checkDimensions<IREE::HAL::InterfaceWorkgroupCountOp, IREE::HAL::InterfaceWorkgroupSizeOp>( - vals, loopInfo.distributionDim))) { + vals, loopInfo.processorDistributionDim))) { return failure(); } return success(); @@ -397,18 +403,18 @@ sentinels.push_back(e); return success(); } else if (auto constExpr = e.dyn_cast<AffineConstantExpr>()) { - if (loopInfo.step) { + if (loopInfo.untiledStep) { return failure(); } - loopInfo.step = IntegerAttr::get(IndexType::get(applyOp.getContext()), - constExpr.getValue()); + loopInfo.untiledStep = IntegerAttr::get( + IndexType::get(applyOp.getContext()), constExpr.getValue()); return success(); } return failure(); } AffineApplyOp applyOp; - TiledLoopInfo &loopInfo; + LoopTilingAndDistributionInfo &loopInfo; }; } // namespace @@ -419,17 +425,16 @@ /// %id = flow.dispatch.workgroup.id[%dim] /// %count = flow.dispatch.workgroup.count[%dim] /// %size = flow.dispatch.workgroup.size[%dim] -/// %offset = affine.apply affine_map<(d0)[s0, s1] -> (d0 + s0 * -/// s1)>(%lb)[%id, %size] %new_step = affine.apply affine_map<(d0)[s0, s1] -> -/// (d0 * s0 * s1)>(%step)[%id, %size] scf.for %iv = %offset to %ub step -/// %new_step { -/// ... -/// } +/// %offset = affine.apply +/// affine_map<(d0)[s0, s1] -> (d0 + s0 * s1)>(%lb)[%id, %size] +/// %new_step = affine.apply +/// affine_map<(d0)[s0, s1] -> (d0 * s0 * s1)>(%step)[%id, %size] +/// scf.for %iv = %offset to %ub step %new_step { ... } /// ``` -static Optional<TiledLoopInfo> isTiledLoop(MLIRContext *context, - scf::ForOp forOp) { - TiledLoopInfo loopInfo; - loopInfo.tiledLoop = forOp; +Optional<LoopTilingAndDistributionInfo> isTiledAndDistributedLoop( + scf::ForOp forOp) { + LoopTilingAndDistributionInfo loopInfo; + loopInfo.loop = forOp; auto lbApplyOp = forOp.lowerBound().getDefiningOp<AffineApplyOp>(); if (!lbApplyOp) { return llvm::None; @@ -444,27 +449,27 @@ failed(stepVisitor.visit(stepApplyOp.getAffineMap().getResults()[0]))) { return llvm::None; } - if (!loopInfo.lb || !loopInfo.step) { + if (!loopInfo.untiledLowerBound || !loopInfo.untiledStep) { return llvm::None; } - loopInfo.ub = getAsOpFoldResult(forOp.upperBound()); + loopInfo.untiledUpperBound = getAsOpFoldResult(forOp.upperBound()); return loopInfo; } -LogicalResult getFilteredOps(FuncOp funcOp, RootOpFilteringFn filteringFn, - SmallVectorImpl<Operation *> &filteredOps, - SmallVectorImpl<TiledLoopInfo> &tiledLoops) { +LogicalResult getFilteredOps( + FuncOp funcOp, RootOpFilteringFn filteringFn, + SmallVectorImpl<Operation *> &filteredOps, + SmallVectorImpl<LoopTilingAndDistributionInfo> &tiledLoops) { Region ®ion = funcOp.body(); if (!llvm::hasSingleElement(region)) { return funcOp.emitError("unable dispatch function with multiple blocks"); } Block *body = ®ion.front(); - MLIRContext *context = funcOp.getContext(); auto forOps = body->getOps<scf::ForOp>(); while (!forOps.empty()) { if (!llvm::hasSingleElement(forOps)) return failure(); scf::ForOp forOp = *(forOps.begin()); - if (auto tiledLoopInfo = isTiledLoop(context, forOp)) { + if (auto tiledLoopInfo = isTiledAndDistributedLoop(forOp)) { tiledLoops.emplace_back(std::move(tiledLoopInfo.getValue())); } body = forOp.getBody(); @@ -478,9 +483,9 @@ return success(); } -LogicalResult getComputeOps(FuncOp funcOp, - SmallVectorImpl<Operation *> &computeOps, - SmallVectorImpl<TiledLoopInfo> &tiledLoops) { +LogicalResult getComputeOps( + FuncOp funcOp, SmallVectorImpl<Operation *> &computeOps, + SmallVectorImpl<LoopTilingAndDistributionInfo> &tiledLoops) { if (failed(getFilteredOps( funcOp, [](Operation *op) { @@ -492,10 +497,11 @@ return success(); } -SmallVector<TiledLoopInfo> getTiledLoopInfo(FuncOp funcOp) { - SmallVector<TiledLoopInfo> info; +SmallVector<LoopTilingAndDistributionInfo> getTiledAndDistributedLoopInfo( + FuncOp funcOp) { + SmallVector<LoopTilingAndDistributionInfo> info; funcOp.walk([&](scf::ForOp forOp) { - if (auto tiledLoopInfo = isTiledLoop(forOp.getContext(), forOp)) { + if (auto tiledLoopInfo = isTiledAndDistributedLoop(forOp)) { info.emplace_back(std::move(tiledLoopInfo.getValue())); } });
diff --git a/iree/compiler/Codegen/Utils/Utils.h b/iree/compiler/Codegen/Utils/Utils.h index f0c563a..e5203b0 100644 --- a/iree/compiler/Codegen/Utils/Utils.h +++ b/iree/compiler/Codegen/Utils/Utils.h
@@ -18,7 +18,7 @@ static constexpr unsigned kNumMaxParallelDims = 3; //===----------------------------------------------------------------------===// -// Utility functions to get entry point(s) +// Utility functions to get entry points //===----------------------------------------------------------------------===// /// Returns true if the given `func` is a kernel dispatch entry point. @@ -32,7 +32,7 @@ IREE::HAL::ExecutableEntryPointOp getEntryPoint(FuncOp funcOp); //===----------------------------------------------------------------------===// -// Utility functions used in setting default configurations. +// Utility functions to set configurations //===----------------------------------------------------------------------===// /// Returns the loops that are partitioned during dispatch region formations, in @@ -57,6 +57,47 @@ ArrayRef<int64_t> getUntiledResultShape(linalg::LinalgOp linalgOp, unsigned resultNum); +/// Information about a tiled and distributed loop. +/// +/// Right now distribution is happening as the same time when we tile the linalg +/// op. 0) Given an original loop: +/// +/// ``` +/// scf.for %iv = %init_lb to %init_ub step %init_step { ... } +/// ``` +// +/// 1) After tiling with tile size `%tile_size`, we have: +// +/// ``` +/// %tiled_step = %init_step * %tile_size +/// scf.for %iv = %init_lb to %init_ub step %tiled_step { ... } +/// ``` +/// +/// 2) After distribution with processor `%id` and `%count`, we have: +// +/// ``` +/// %dist_lb = %init_lb + %id * %tiled_step +/// %dist_step = %init_step * %tile_size * %count +/// scf.for %iv = %dist_lb to %init_ub step %dist_step { ... } +/// ``` +/// +/// Given a loop already after 2), this struct contains recovered information +/// about 0) and 1). +struct LoopTilingAndDistributionInfo { + // The tiled and distributed loop. + Operation *loop; + // The lower bound for the original untiled loop. + OpFoldResult untiledLowerBound; + // The upper bound for the original untiled loop. + OpFoldResult untiledUpperBound; + // The step for the original untiled loop. + OpFoldResult untiledStep; + // The tile size used to tile (and not distribute) the original untiled loop. + Optional<int64_t> tileSize; + // The processor dimension this loop is distributed to. + unsigned processorDistributionDim; +}; + /// Assuming that `funcOp` contains a single nested scf.for that represented the /// tiled+fused+distributed loops with the distribution being across workgroups, /// i.e. @@ -76,19 +117,11 @@ /// `scf.for` operations in the function return the linalg operations in the /// body of the function if it has a single basic block. Return failure in all /// other cases. - -struct TiledLoopInfo { - Operation *tiledLoop; - OpFoldResult lb; - OpFoldResult ub; - OpFoldResult step; - Optional<int64_t> workgroupSize; - unsigned distributionDim; -}; using RootOpFilteringFn = std::function<bool(Operation *)>; -LogicalResult getFilteredOps(FuncOp funcOp, RootOpFilteringFn filteringFn, - SmallVectorImpl<Operation *> &filteredOps, - SmallVectorImpl<TiledLoopInfo> &tiledLoops); +LogicalResult getFilteredOps( + FuncOp funcOp, RootOpFilteringFn filteringFn, + SmallVectorImpl<Operation *> &filteredOps, + SmallVectorImpl<LoopTilingAndDistributionInfo> &tiledLoops); /// Specialization of `getFilteredOps` for filtering `LinalgOp`s and /// `LinagExtOp`s. @@ -96,12 +129,13 @@ /// within the loop. The marker is the way to tie into rest of the /// codegen. Refactor the downstream passes and get rid of the markers once and /// for all. -LogicalResult getComputeOps(FuncOp funcOp, - SmallVectorImpl<Operation *> &computeOps, - SmallVectorImpl<TiledLoopInfo> &tiledLoops); +LogicalResult getComputeOps( + FuncOp funcOp, SmallVectorImpl<Operation *> &computeOps, + SmallVectorImpl<LoopTilingAndDistributionInfo> &tiledLoops); -/// Collect information about loops matching tiled+distribute pattern. -SmallVector<TiledLoopInfo> getTiledLoopInfo(FuncOp funcOp); +/// Collects information about loops matching tiled+distribute pattern. +SmallVector<LoopTilingAndDistributionInfo> getTiledAndDistributedLoopInfo( + FuncOp funcOp); } // namespace iree_compiler } // namespace mlir