| //===- ReductionTreeUtils.cpp - Reduction Tree Utilities ------------------===// |
| // |
| // 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 defines the Reduction Tree Utilities. It defines pass independent |
| // methods that help in a reduction pass of the MLIR Reduce tool. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Reducer/ReductionTreeUtils.h" |
| |
| #define DEBUG_TYPE "mlir-reduce" |
| |
| using namespace mlir; |
| |
| /// Update the golden module's content with that of the reduced module. |
| void ReductionTreeUtils::updateGoldenModule(ModuleOp &golden, |
| ModuleOp reduced) { |
| golden.getBody()->clear(); |
| |
| golden.getBody()->getOperations().splice(golden.getBody()->begin(), |
| reduced.getBody()->getOperations()); |
| } |
| |
| /// Update the smallest node traversed so far in the reduction tree and |
| /// print the debugging information for the currNode being traversed. |
| void ReductionTreeUtils::updateSmallestNode(ReductionNode *currNode, |
| ReductionNode *&smallestNode, |
| std::vector<int> path) { |
| LLVM_DEBUG(llvm::dbgs() << "\nTree Path: root"); |
| #ifndef NDEBUG |
| for (int nodeIndex : path) |
| LLVM_DEBUG(llvm::dbgs() << " -> " << nodeIndex); |
| #endif |
| |
| LLVM_DEBUG(llvm::dbgs() << "\nSize (chars): " << currNode->getSize()); |
| if (currNode->getSize() < smallestNode->getSize()) { |
| LLVM_DEBUG(llvm::dbgs() << " - new smallest node!"); |
| smallestNode = currNode; |
| } |
| } |
| |
| /// Create a transform space index vector based on the specified number of |
| /// indices. |
| std::vector<bool> ReductionTreeUtils::createTransformSpace(ModuleOp module, |
| int numIndices) { |
| std::vector<bool> transformSpace; |
| for (int i = 0; i < numIndices; ++i) |
| transformSpace.push_back(false); |
| |
| return transformSpace; |
| } |
| |
| /// Translate section start and end into a vector of ranges specifying the |
| /// section in the non transformed indices in the transform space. |
| static std::vector<std::tuple<int, int>> getRanges(std::vector<bool> tSpace, |
| int start, int end) { |
| std::vector<std::tuple<int, int>> ranges; |
| int rangeStart = 0; |
| int rangeEnd = 0; |
| bool inside = false; |
| int transformableCount = 0; |
| |
| for (auto element : llvm::enumerate(tSpace)) { |
| int index = element.index(); |
| bool value = element.value(); |
| |
| if (start <= transformableCount && transformableCount < end) { |
| if (!value && !inside) { |
| inside = true; |
| rangeStart = index; |
| } |
| if (value && inside) { |
| rangeEnd = index; |
| ranges.push_back(std::make_tuple(rangeStart, rangeEnd)); |
| inside = false; |
| } |
| } |
| |
| if (!value) |
| transformableCount++; |
| |
| if (transformableCount == end && inside) { |
| ranges.push_back(std::make_tuple(rangeStart, index + 1)); |
| inside = false; |
| break; |
| } |
| } |
| |
| return ranges; |
| } |
| |
| /// Create the specified number of variants by applying the transform method |
| /// to different ranges of indices in the parent module. The isDeletion boolean |
| /// specifies if the transformation is the deletion of indices. |
| void ReductionTreeUtils::createVariants( |
| ReductionNode *parent, const Tester &test, int numVariants, |
| llvm::function_ref<void(ModuleOp, int, int)> transform, bool isDeletion) { |
| std::vector<bool> newTSpace; |
| ModuleOp module = parent->getModule(); |
| |
| std::vector<bool> parentTSpace = parent->getTransformSpace(); |
| int indexCount = parent->transformSpaceSize(); |
| std::vector<std::tuple<int, int>> ranges; |
| |
| // No new variants can be created. |
| if (indexCount == 0) |
| return; |
| |
| // Create a single variant by transforming the unique index. |
| if (indexCount == 1) { |
| ModuleOp variantModule = module.clone(); |
| if (isDeletion) { |
| transform(variantModule, 0, 1); |
| } else { |
| ranges = getRanges(parentTSpace, 0, parentTSpace.size()); |
| transform(variantModule, std::get<0>(ranges[0]), std::get<1>(ranges[0])); |
| } |
| |
| new ReductionNode(variantModule, parent, newTSpace); |
| |
| return; |
| } |
| |
| // Create the specified number of variants. |
| for (int i = 0; i < numVariants; ++i) { |
| ModuleOp variantModule = module.clone(); |
| newTSpace = parent->getTransformSpace(); |
| int sectionSize = indexCount / numVariants; |
| int sectionStart = sectionSize * i; |
| int sectionEnd = sectionSize * (i + 1); |
| |
| if (i == numVariants - 1) |
| sectionEnd = indexCount; |
| |
| if (isDeletion) |
| transform(variantModule, sectionStart, sectionEnd); |
| |
| ranges = getRanges(parentTSpace, sectionStart, sectionEnd); |
| |
| for (auto range : ranges) { |
| int rangeStart = std::get<0>(range); |
| int rangeEnd = std::get<1>(range); |
| |
| for (int x = rangeStart; x < rangeEnd; ++x) |
| newTSpace[x] = true; |
| |
| if (!isDeletion) |
| transform(variantModule, rangeStart, rangeEnd); |
| } |
| |
| // Create Reduction Node in the Reduction tree |
| new ReductionNode(variantModule, parent, newTSpace); |
| } |
| } |