From 1bdc117ddb07d7f2087a9413654be628364be087 Mon Sep 17 00:00:00 2001 From: Hanchen Ye Date: Sat, 6 Nov 2021 02:43:47 -0500 Subject: [PATCH] [ArrayPartition] Add array partition API with explicit factors and kinds (cyclic or block) --- include/scalehls/Transforms/Utils.h | 6 +- lib/Bindings/Python/ScaleHLSModule.cpp | 23 +++- lib/Transforms/Directive/ArrayPartition.cpp | 113 +++++++++++++------- lib/Transforms/Utils.cpp | 6 +- 4 files changed, 102 insertions(+), 46 deletions(-) diff --git a/include/scalehls/Transforms/Utils.h b/include/scalehls/Transforms/Utils.h index db989b5..e944b9b 100644 --- a/include/scalehls/Transforms/Utils.h +++ b/include/scalehls/Transforms/Utils.h @@ -82,7 +82,11 @@ bool applyFullyUnrollAndPartition(Block &block, FuncOp func); bool applyMemoryAccessOpt(FuncOp func); -bool applyArrayPartition(FuncOp func); +bool applyArrayPartition(Value array, ArrayRef factors, + ArrayRef kinds, + bool updateFuncSignature = true); + +bool applyAutoArrayPartition(FuncOp func); /// Apply optimization strategy to a loop band. The ancestor function is /// also passed in because the post-tiling optimizations have to take diff --git a/lib/Bindings/Python/ScaleHLSModule.cpp b/lib/Bindings/Python/ScaleHLSModule.cpp index 2938e6b..77f83cb 100644 --- a/lib/Bindings/Python/ScaleHLSModule.cpp +++ b/lib/Bindings/Python/ScaleHLSModule.cpp @@ -176,11 +176,29 @@ static bool memoryAccessOpt(MlirOperation op) { static bool autoArrayPartition(MlirOperation op) { py::gil_scoped_release(); if (auto func = dyn_cast(unwrap(op))) - return applyArrayPartition(func); + return applyAutoArrayPartition(func); throw py::raiseValueError("targeted operation not a function"); return false; } +//===----------------------------------------------------------------------===// +// Array transform APIs +//===----------------------------------------------------------------------===// + +/// TODO: Support to apply different partition kind to different dimension. +static bool arrayPartition(MlirValue array, py::object factorsObject, + std::string kind) { + py::gil_scoped_release(); + llvm::SmallVector factors; + if (!getVectorFromUnsignedNpArray(factorsObject.ptr(), factors)) + return false; + llvm::SmallVector kinds( + factors.size(), kind == "cyclic" ? hlscpp::PartitionKind::CYCLIC + : kind == "block" ? hlscpp::PartitionKind::BLOCK + : hlscpp::PartitionKind::NONE); + return applyArrayPartition(unwrap(array), factors, kinds); +} + //===----------------------------------------------------------------------===// // Emission APIs //===----------------------------------------------------------------------===// @@ -224,6 +242,9 @@ PYBIND11_MODULE(_scalehls, m) { m.def("memory_access_opt", &memoryAccessOpt); m.def("auto_array_partition", &autoArrayPartition); + // Array transform APIs. + m.def("array_partition", &arrayPartition); + // Emission APIs. m.def("emit_hlscpp", &emitHlsCpp); diff --git a/lib/Transforms/Directive/ArrayPartition.cpp b/lib/Transforms/Directive/ArrayPartition.cpp index 346a62b..68208ea 100644 --- a/lib/Transforms/Directive/ArrayPartition.cpp +++ b/lib/Transforms/Directive/ArrayPartition.cpp @@ -43,8 +43,70 @@ static void updateSubFuncs(FuncOp func, Builder builder) { }); } -/// TODO: support to pass in partition strategy. -bool scalehls::applyArrayPartition(FuncOp func) { +bool scalehls::applyArrayPartition(Value array, ArrayRef factors, + ArrayRef kinds, + bool updateFuncSignature) { + auto builder = Builder(array.getContext()); + auto arrayType = array.getType().dyn_cast(); + if (!arrayType || !arrayType.hasStaticShape() || + factors.size() != arrayType.getRank() || + kinds.size() != arrayType.getRank()) + return false; + + // Walk through each dimension of the current memory. + SmallVector partitionIndices; + SmallVector addressIndices; + + for (int64_t dim = 0; dim < arrayType.getRank(); ++dim) { + auto kind = kinds[dim]; + auto factor = factors[dim]; + + if (kind == PartitionKind::CYCLIC) { + partitionIndices.push_back(builder.getAffineDimExpr(dim) % factor); + addressIndices.push_back(builder.getAffineDimExpr(dim).floorDiv(factor)); + + } else if (kind == PartitionKind::BLOCK) { + auto blockFactor = (arrayType.getShape()[dim] + factor - 1) / factor; + partitionIndices.push_back( + builder.getAffineDimExpr(dim).floorDiv(blockFactor)); + addressIndices.push_back(builder.getAffineDimExpr(dim) % blockFactor); + + } else { + partitionIndices.push_back(builder.getAffineConstantExpr(0)); + addressIndices.push_back(builder.getAffineDimExpr(dim)); + } + } + + // Construct new layout map. + partitionIndices.append(addressIndices.begin(), addressIndices.end()); + auto layoutMap = AffineMap::get(arrayType.getRank(), 0, partitionIndices, + builder.getContext()); + + // Construct new array type. + auto newType = + MemRefType::get(arrayType.getShape(), arrayType.getElementType(), + layoutMap, arrayType.getMemorySpace()); + + // Set new type. + array.setType(newType); + + if (updateFuncSignature) + if (auto func = dyn_cast(array.getParentBlock()->getParentOp())) { + // Align function type with entry block argument types only if the array + // is defined as an argument of the function. + if (!array.getDefiningOp()) { + auto resultTypes = func.front().getTerminator()->getOperandTypes(); + auto inputTypes = func.front().getArgumentTypes(); + func.setType(builder.getFunctionType(inputTypes, resultTypes)); + } + + // Update the types of all sub-functions. + updateSubFuncs(func, builder); + } + return true; +} + +bool scalehls::applyAutoArrayPartition(FuncOp func) { // Check whether the input function is pipelined. bool funcPipeline = false; if (auto attr = func->getAttrOfType("pipeline")) @@ -201,7 +263,7 @@ bool scalehls::applyArrayPartition(FuncOp func) { assert(subFunc && "callable is not a function operation"); // Apply array partition to the sub-function. - applyArrayPartition(subFunc); + applyAutoArrayPartition(subFunc); auto subFuncType = subFunc.getType(); unsigned index = 0; @@ -243,47 +305,16 @@ bool scalehls::applyArrayPartition(FuncOp func) { auto builder = Builder(func); for (auto pair : partitionsMap) { auto memref = pair.first; - auto memrefType = memref.getType().cast(); auto partitions = pair.second; - // Walk through each dimension of the current memory. - SmallVector partitionIndices; - SmallVector addressIndices; - - for (int64_t dim = 0; dim < memrefType.getRank(); ++dim) { - auto partition = partitions[dim]; - auto kind = partition.first; - auto factor = partition.second; - - if (kind == PartitionKind::CYCLIC) { - partitionIndices.push_back(builder.getAffineDimExpr(dim) % factor); - addressIndices.push_back( - builder.getAffineDimExpr(dim).floorDiv(factor)); - - } else if (kind == PartitionKind::BLOCK) { - auto blockFactor = (memrefType.getShape()[dim] + factor - 1) / factor; - partitionIndices.push_back( - builder.getAffineDimExpr(dim).floorDiv(blockFactor)); - addressIndices.push_back(builder.getAffineDimExpr(dim) % blockFactor); - - } else { - partitionIndices.push_back(builder.getAffineConstantExpr(0)); - addressIndices.push_back(builder.getAffineDimExpr(dim)); - } + SmallVector kinds; + SmallVector factors; + for (auto info : partitions) { + kinds.push_back(info.first); + factors.push_back(info.second); } - // Construct new layout map. - partitionIndices.append(addressIndices.begin(), addressIndices.end()); - auto layoutMap = AffineMap::get(memrefType.getRank(), 0, partitionIndices, - builder.getContext()); - - // Construct new memref type. - auto newType = - MemRefType::get(memrefType.getShape(), memrefType.getElementType(), - layoutMap, memrefType.getMemorySpace()); - - // Set new type. - memref.setType(newType); + applyArrayPartition(memref, factors, kinds, /*updateFuncSignature=*/false); } // Align function type with entry block argument types. @@ -303,7 +334,7 @@ struct ArrayPartition : public ArrayPartitionBase { for (auto func : getOperation().getOps()) { if (auto funcDirect = getFuncDirective(func)) if (funcDirect.getTopFunc()) - applyArrayPartition(func); + applyAutoArrayPartition(func); } } }; diff --git a/lib/Transforms/Utils.cpp b/lib/Transforms/Utils.cpp index f47be07..2ced7a8 100644 --- a/lib/Transforms/Utils.cpp +++ b/lib/Transforms/Utils.cpp @@ -138,7 +138,7 @@ bool scalehls::applyFullyUnrollAndPartition(Block &block, FuncOp func) { return false; // Apply the best suitable array partition strategy to the function. - applyArrayPartition(func); + applyAutoArrayPartition(func); return true; } @@ -171,7 +171,7 @@ bool scalehls::applyOptStrategy(AffineLoopBand &band, FuncOp func, return false; // Apply the best suitable array partition strategy to the function. - applyArrayPartition(func); + applyAutoArrayPartition(func); return true; } @@ -206,7 +206,7 @@ bool scalehls::applyOptStrategy(FuncOp func, ArrayRef tileLists, return false; // Apply the best suitable array partition strategy to the function. - applyArrayPartition(func); + applyAutoArrayPartition(func); return true; }