Refactor Reduction Tree Pass

[msifontes]

Summary

Refactor the way the reduction tree pass works in the MLIR Reduce tool by introducing a set of utilities that facilitate the implementation of new Reducer classes to be used in the passes.

This will allow for the fast implementation of general transformations to operate on all mlir modules as well as custom transformations for different dialects.

These utilities allow for the implementation of Reducer classes by simply defining a method that indexes the operations/blocks/regions to be transformed and a method to perform the deletion or transformation based on the indexes.

Create the transformSpace class member in the ReductionNode class to keep track of the indexes that have already been transformed or deleted at the current level.

Delete the FunctionReducer class and replace it with the OpReducer class to reflect the use of the new API while performing the same transformation and allowing the instantiation of a reduction pass for different types of operations at the module's highest hierarchichal level.

Modify the SinglePath Traversal method to reflect the use of the new API.

[msifontes]

Fix indexing bug in createTransformedVariants().

[msifontes]

Merge createTransformedVariants and createErasedVariants utilities into single createVariants utility.

[msifontes]

NFC: clang-format

[jpienaar]

mlir/tools/mlir-reduce/Passes/OpReducer.cpp

32

So seems like template is only used for getOps really. So you could have your implementation have functionref Module -> iterator_range<OpIterator> filter function and then avoid templating inside here. You could still have the simple templated top class, but then the instantiated one & implementation could just call the filter function and so in you cpp file you'd have no templates.

69

This seems less ideal than having the above in a header as I think this restricts you to only use these ops.

[msifontes]

Apply Requested changes

• Modify OpReducer class to use PImpl idiom.
• Modify Reduction Tree Utilities to use llvm::function_ref instead of std::function

[msifontes]

NFC: Formatting changes

[msifontes]

Merge confilcts in working tree.

[msifontes]

Change mlir-reduce.cpp to use new OpReducer

[msifontes]

Remove unused headers in OpReducer.h

[msifontes]

Modify the reduction node transform space to use a bool vector with the transformed indices marked as true.

[msifontes]

Fix clang-format error

[jpienaar]

Nice thanks, LG with comments addressed

mlir/include/mlir/Reducer/Passes/OpReducer.h

39

Why not const Tester& ? Is null allowed?

43

Could you expand to document what transform arg does/how it will be used?

53

These would all be of OpType?

55

Lets make this OpType as refers to class

mlir/include/mlir/Reducer/ReductionTreePass.h

127–128

Could we change 1 and 2 (below) to something more informative?

142

Same here

mlir/include/mlir/Reducer/ReductionTreeUtils.h

32

Mmm, if these are all static, why not just have these in a namespace rather than class?

36

processNode is too general as a utility function name. Lets find something more descriptive.

mlir/tools/mlir-reduce/Passes/OpReducer.cpp

2

Line update?

mlir/tools/mlir-reduce/ReductionTreeUtils.cpp

36

Why not just range based for?

70

Should this be wrapped in LLVM_DEBUG too?

123

Where is this used?

[msifontes]

Apply requested changes

• Use constant Tester reference throughout.
• Modify debugging messages to explain how the variants are being created.
• Change the ReductionTreeUtils class to be a namespace.
• Replace processNode method name with updateSmallestNode.

mlir/include/mlir/Reducer/Passes/FunctionReducer.h

-//===- FunctionReducer.h - MLIR Reduce Function Reducer ---------*- 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 defines the FunctionReducer class. It defines a variant generator
-// method with the purpose of producing different variants by eliminating
-// functions from the  parent module.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_REDUCER_PASSES_FUNCTIONREDUCER_H
-#define MLIR_REDUCER_PASSES_FUNCTIONREDUCER_H
-
-#include "mlir/Reducer/ReductionNode.h"
-#include "mlir/Reducer/Tester.h"
-
-namespace mlir {
-
-/// The FunctionReducer class defines a variant generator method that produces
-/// multiple variants by eliminating different operations from the
-/// parent module.
-class FunctionReducer {
-public:
-  /// Generate variants by removing functions from the module in the parent
-  /// Reduction Node and link the variants as children in the Reduction Tree
-  /// Pass.
-  void generateVariants(ReductionNode *parent, const Tester *test);
-};
-
-} // end namespace mlir
-
-#endif

mlir/include/mlir/Reducer/Passes/OpReducer.h

+//===- OpReducer.h - MLIR Reduce Operation Reducer ------------*- 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 defines the OpReducer class. It defines a variant generator method
+// with the purpose of producing different variants by eliminating a
+// parametarizable type of operations from the  parent module.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_REDUCER_PASSES_OPREDUCER_H
+#define MLIR_REDUCER_PASSES_OPREDUCER_H
+
+#include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/IR/Region.h"
+#include "mlir/Reducer/ReductionNode.h"
+#include "mlir/Reducer/ReductionTreeUtils.h"
+#include "mlir/Reducer/Tester.h"
+
+namespace mlir {
+
+class OpReducerImpl {
+public:
+  OpReducerImpl(
+      llvm::function_ref<std::vector<Operation *>(ModuleOp)> getSpecificOps);
+
+  /// Return the name of this reducer class.
+  StringRef getName();
+
+  /// Return the initial transformSpace cointaing the transformable indices.
+  std::vector<int> initTransformSpace(ModuleOp module);
+
+  /// Generate variants by removing opType operations from the module in the
+  /// parent and link the variants as childs in the Reduction Tree Pass.
+  void generateVariants(ReductionNode *parent, const Tester *test,
+                        int numVariants);
+
+  /// Generate variants by removing opType operations from the module in the
+  /// parent and link the variants as childs in the Reduction Tree Pass.
+  void generateVariants(ReductionNode *parent, const Tester *test,
+                        int numVariants,
+                        llvm::function_ref<void(ModuleOp, int, int)> transform);
+
+private:
+  llvm::function_ref<std::vector<Operation *>(ModuleOp)> getSpecificOps;
+};
+
+/// The OpReducer class defines a variant generator method that produces
+/// multiple variants by eliminating different operations from the
+/// parent module.
+template <typename opType>
+class OpReducer {
+public:
+  OpReducer() : impl(new OpReducerImpl(getSpecificOps)) {}
+
+  /// Returns the vector of pointer to the opType operations in the module.
+  static std::vector<Operation *> getSpecificOps(ModuleOp module) {
+    std::vector<Operation *> ops;
+    for (auto op : module.getOps<opType>()) {
+      ops.push_back(op);
+    }
+    return ops;
+  }
+
+  /// Deletes the opType operations in the module in the specified index.
+  static void deleteOps(ModuleOp module, int start, int end) {
+    std::vector<Operation *> opsToRemove;
+
+    for (auto op : enumerate(getSpecificOps(module))) {
+      int index = op.index();
+      if (index >= start && index < end)
+        opsToRemove.push_back(op.value());
+    }
+
+    for (Operation *o : opsToRemove) {
+      o->dropAllUses();
+      o->erase();
+    }
+  }
+
+  /// Return the name of this reducer class.
+  StringRef getName() { return impl->getName(); }
+
+  /// Return the initial transformSpace cointaing the transformable indices.
+  std::vector<int> initTransformSpace(ModuleOp module) {
+    return impl->initTransformSpace(module);
+  }
+
+  /// Generate variants by removing opType operations from the module in the
+  /// parent and link the variants as childs in the Reduction Tree Pass.
+  void generateVariants(ReductionNode *parent, const Tester *test,
+                        int numVariants) {
+    impl->generateVariants(parent, test, numVariants, deleteOps);
+  }
+
+private:
+  std::unique_ptr<OpReducerImpl> impl;
+};
+
+} // end namespace mlir
+
+#endif

mlir/include/mlir/Reducer/ReductionNode.h

 /// This class defines the ReductionNode which is used to wrap the module of
 /// a generated variant and keep track of the necessary metadata for the
 /// reduction pass. The nodes are linked together in a reduction tree stucture
 /// which defines the relationship between all the different generated variants.
 class ReductionNode {
 public:
   ReductionNode(ModuleOp module, ReductionNode *parent);
 
+  ReductionNode(ModuleOp module, ReductionNode *parent,
+                std::vector<int> transformSpace);
+
   /// Calculates and initializes the size and interesting values of the node.
   void measureAndTest(const Tester *test);
 
   /// Returns the module.
   ModuleOp getModule() const { return module; }
 
   /// Returns true if the size and interestingness have been calculated.
   bool isEvaluated() const;
 
   /// Returns the size in bytes of the module.
   int getSize() const;
 
   /// Returns true if the module exhibits the interesting behavior.
   bool isInteresting() const;
 
   /// Returns the pointer to a child variant by index.
   ReductionNode *getVariant(unsigned long index) const;
 
+  /// Returns the number of child variants.
+  int variantsSize() const;
+
   /// Returns true if the vector containing the child variants is empty.
   bool variantsEmpty() const;
 
   /// Sort the child variants and remove the uninteresting ones.
   void organizeVariants(const Tester *test);
 
+  /// Returns the number of child variants.
+  int transformSpaceSize();
+
+  /// Returns a vector of the transformable indices.
+  const std::vector<int> getTransformSpace();
+
 private:
   /// Link a child variant node.
   void linkVariant(ReductionNode *newVariant);
 
   // This is the MLIR module of this variant.
   ModuleOp module;
 
   // This is true if the module has been evaluated and it exhibits the
   // interesting behavior.
   bool interesting;
 
   // This indicates the number of characters in the printed module if the module
   // has been evaluated.
   int size;
 
   // This indicates if the module has been evalueated (measured and tested).
   bool evaluated;
 
+  // Contains the indices in the node that haven't been transformed in previous
+  // levels of the reduction tree.
+  std::vector<int> transformSpace;
+
   // This points to the child variants that were created using this node as a
   // starting point.
   std::vector<std::unique_ptr<ReductionNode>> variants;
 };
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/Reducer/ReductionTreePass.h

 
 #ifndef MLIR_REDUCER_REDUCTIONTREEPASS_H
 #define MLIR_REDUCER_REDUCTIONTREEPASS_H
 
 #include <vector>
 
 #include "PassDetail.h"
 #include "ReductionNode.h"
-#include "mlir/Reducer/Passes/FunctionReducer.h"
+#include "mlir/Reducer/Passes/OpReducer.h"
+#include "mlir/Reducer/ReductionTreeUtils.h"
 #include "mlir/Reducer/Tester.h"
 
+#define DEBUG_TYPE "mlir-reduce"
+
 namespace mlir {
 
-/// Defines the traversal method options to be used in the reduction tree
+// Defines the traversal method options to be used in the reduction tree
 /// traversal.
-enum TraversalMode { SinglePath, MultiPath, Concurrent, Backtrack };
-
-// This class defines the non- templated utilities used by the ReductionTreePass
-// class.
-class ReductionTreeUtils {
-public:
-  void updateGoldenModule(ModuleOp &golden, ModuleOp reduced);
-};
+enum TraversalMode { SinglePath, Backtrack, MultiPath };
 
 /// This class defines the Reduction Tree Pass. It provides a framework to
 /// to implement a reduction pass using a tree structure to keep track of the
 /// generated reduced variants.
 template <typename Reducer, TraversalMode mode>
 class ReductionTreePass
     : public ReductionTreeBase<ReductionTreePass<Reducer, mode>> {
 public:
-  ReductionTreePass(const Tester *test) : test(test) {}
-
   ReductionTreePass(const ReductionTreePass &pass)
-      : root(new ReductionNode(pass.root->getModule().clone(), nullptr)),
+      : root(new ReductionNode(pass.root->getModule().clone(), nullptr,
+                               pass.root->getTransformSpace())),
         test(pass.test) {}
 
+  ReductionTreePass(const Tester *test) : test(test) {}
+
   /// Runs the pass instance in the pass pipeline.
   void runOnOperation() override {
     ModuleOp module = this->getOperation();
-    this->root = std::make_unique<ReductionNode>(module, nullptr);
+    Reducer reducer;
+    std::vector<int> transformSpace = reducer.initTransformSpace(module);
     ReductionNode *reduced;
 
+    this->root =
+        std::make_unique<ReductionNode>(module, nullptr, transformSpace);
+
+    root->measureAndTest(test);
+
+    LLVM_DEBUG(llvm::dbgs() << "\nReduction Tree Pass: " << reducer.getName(););
     switch (mode) {
     case SinglePath:
+      LLVM_DEBUG(llvm::dbgs() << " (Single Path)\n";);
       reduced = singlePathTraversal();
       break;
     default:
       llvm::report_fatal_error("Traversal method not currently supported.");
     }
 
-    ReductionTreeUtils utils;
-    utils.updateGoldenModule(module, reduced->getModule());
+    ReductionTreeUtils::updateGoldenModule(module,
+                                           reduced->getModule().clone());
   }
 
 private:
   // Points to the root node in this reduction tree.
   std::unique_ptr<ReductionNode> root;
 
   // This object defines the variant generation at each level of the reduction
   // tree.
   Reducer reducer;
 
   // This is used to test the interesting behavior of the reduction nodes in the
   // tree.
   const Tester *test;
 
   /// Traverse the most reduced path in the reduction tree by generating the
   /// variants at each level using the Reducer parameter's generateVariants
   /// function. Stops when no new successful variants can be created at the
   /// current level.
   ReductionNode *singlePathTraversal() {
-    ReductionNode *currLevel = root.get();
+    ReductionNode *currNode = root.get();
+    ReductionNode *smallestNode = currNode;
+    int tSpaceSize = currNode->transformSpaceSize();
+    int numVariants = 2;
+    std::vector<int> path;
+
+    // Attempt to create atleast one succesful variant by creating the least
+    // number of variants.
+    while (numVariants <= tSpaceSize) {
+      ReductionTreeUtils::processNode(currNode, smallestNode, path);
+
+      LLVM_DEBUG(llvm::dbgs() << "\nGenerating " << numVariants);
+      LLVM_DEBUG(llvm::dbgs() << " variants\nTesting\n");
+
+      reducer.generateVariants(currNode, test, numVariants);
+      currNode->organizeVariants(test);
+
+      if (currNode->variantsEmpty()) {
+        numVariants = numVariants * 2;
+        continue;
+      }
 
-    while (true) {
-      reducer.generateVariants(currLevel, test);
-      currLevel->organizeVariants(test);
+      currNode = currNode->getVariant(0);
+      tSpaceSize = currNode->transformSpaceSize();
+      path.push_back(0);
+    }
 
-      if (currLevel->variantsEmpty())
-        break;
+    // Attempt to create a single veriant by trasforming/erasing the unique
+    // index in the transform space.
+    if (tSpaceSize == 1) {
+      ReductionTreeUtils::processNode(currNode, smallestNode, path);
 
-      currLevel = currLevel->getVariant(0);
+      LLVM_DEBUG(llvm::dbgs() << "\nGenerating 1 variant\nTesting\n";);
+
+      reducer.generateVariants(currNode, test, numVariants);
+      currNode->organizeVariants(test);
+
+      if (!currNode->variantsEmpty()) {
+        currNode = currNode->getVariant(0);
+        path.push_back(0);
+
+        ReductionTreeUtils::processNode(currNode, smallestNode, path);
+      }
     }
 
-    return currLevel;
+    return smallestNode;
   }
 };
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/Reducer/ReductionTreeUtils.h

+//===- ReductionTreeUtils.h - Reduction Tree utilities ----------*- 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 defines the Reduction Tree Utilities. It defines pass independent
+// methods that help in the reduction passes of the MLIR Reduce tool.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_REDUCER_REDUCTIONTREEUTILS_H
+#define MLIR_REDUCER_REDUCTIONTREEUTILS_H
+
+#include "PassDetail.h"
+#include "ReductionNode.h"
+#include "mlir/Reducer/Tester.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "mlir-reduce"
+
+namespace mlir {
+
+// This class defines the utilities  for the implementation of custom reduction
+// passes using the ReductionTreePass framework.
+class ReductionTreeUtils {
+public:
+  /// Update the golden module's content with that of the reduced module.
+  static void updateGoldenModule(ModuleOp &golden, ModuleOp reduced);
+
+  /// Update the the smallest node traversed so far in the reduction tree and
+  /// print the debugging information for the currNode being traversed.
+  static void processNode(ReductionNode *currNode, ReductionNode *&smallestNode,
+                          std::vector<int> path);
+
+  /// Create a transform space index vector based on the specified number of
+  /// indices.
+  static std::vector<int> createTransformSpace(ModuleOp module, int numIndices);
+
+  /// Create the specified number of variants by applying the transform method
+  /// to different ranges of indices in the parent module.
+  static void
+  createVariants(ReductionNode *parent, const Tester *test, int numVariants,
+                 llvm::function_ref<void(ModuleOp, int, int)> transform);
+};
+
+} // end namespace mlir
+
+#endif

mlir/tools/mlir-reduce/CMakeLists.txt

set(LIBS
   MLIRPass
   MLIRReduce
   MLIRSupport
   MLIRTransforms
   MLIRTransformUtils
   )
 
 add_llvm_tool(mlir-reduce
-  Passes/FunctionReducer.cpp
+  Passes/OpReducer.cpp
   ReductionNode.cpp
-  ReductionTreePass.cpp
+  ReductionTreeUtils.cpp
   mlir-reduce.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Reducer
 
   DEPENDS
   MLIRReducerIncGen
   )
 
 target_link_libraries(mlir-reduce PRIVATE ${LIBS})
 llvm_update_compile_flags(mlir-reduce)
 
 mlir_check_all_link_libraries(mlir-reduce)

mlir/tools/mlir-reduce/Passes/FunctionReducer.cpp

-//===- FunctionReducer.cpp - MLIR Reduce Function Reducer -----------------===//
-//
-// 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 FunctionReducer class. It defines a variant generator
-// class to be used in a Reduction Tree Pass instantiation with the aim of
-// reducing the number of function operations in an MLIR Module.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Reducer/Passes/FunctionReducer.h"
-#include "mlir/IR/Function.h"
-
-using namespace mlir;
-
-/// Return the number of function operations in the module's body.
-int countFunctions(ModuleOp module) {
-  auto ops = module.getOps<FuncOp>();
-  return std::distance(ops.begin(), ops.end());
-}
-
-/// Generate variants by removing function operations from the module in the
-/// parent and link the variants as children in the Reduction Tree Pass.
-void FunctionReducer::generateVariants(ReductionNode *parent,
-                                       const Tester *test) {
-  ModuleOp module = parent->getModule();
-  int opCount = countFunctions(module);
-  int sectionSize = opCount / 2;
-  std::vector<Operation *> opsToRemove;
-
-  if (opCount == 0)
-    return;
-
-  // Create a variant by deleting all ops.
-  if (opCount == 1) {
-    opsToRemove.clear();
-    ModuleOp moduleVariant = module.clone();
-
-    for (FuncOp op : moduleVariant.getOps<FuncOp>())
-      opsToRemove.push_back(op);
-
-    for (Operation *o : opsToRemove)
-      o->erase();
-
-    new ReductionNode(moduleVariant, parent);
-
-    return;
-  }
-
-  // Create two variants by bisecting the module.
-  for (int i = 0; i < 2; ++i) {
-    opsToRemove.clear();
-    ModuleOp moduleVariant = module.clone();
-
-    for (auto op : enumerate(moduleVariant.getOps<FuncOp>())) {
-      int index = op.index();
-      if (index >= sectionSize * i && index < sectionSize * (i + 1))
-        opsToRemove.push_back(op.value());
-    }
-
-    for (Operation *o : opsToRemove)
-      o->erase();
-
-    new ReductionNode(moduleVariant, parent);
-  }
-
-  return;
-}

mlir/tools/mlir-reduce/Passes/OpReducer.cpp

+//===- OpReducer.h - Operation Reducer --------------------------*- 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 defines the OpReducer class. It defines a variant generator method
+// with the purpose of producing different variants by eliminating a
+// parametarizable type of operations from the  parent module.
+//
+//===----------------------------------------------------------------------===//
+#include "mlir/Reducer/Passes/OpReducer.h"
+
+using namespace mlir;
+
+OpReducerImpl::OpReducerImpl(
+    llvm::function_ref<std::vector<Operation *>(ModuleOp)> getSpecificOps)
+    : getSpecificOps(getSpecificOps) {}
+
+/// Return the name of this reducer class.
+StringRef OpReducerImpl::getName() {
+  return StringRef("High Level Operation Reduction");
+}
+
+/// Return the initial transformSpace cointaing the transformable indices.
+std::vector<int> OpReducerImpl::initTransformSpace(ModuleOp module) {
+  auto ops = getSpecificOps(module);
+  int numOps = std::distance(ops.begin(), ops.end());
+  return ReductionTreeUtils::createTransformSpace(module, numOps);
+}
+
+/// Generate variants by removing opType operations from the module in the
+/// parent and link the variants as childs in the Reduction Tree Pass.
+void OpReducerImpl::generateVariants(
+    ReductionNode *parent, const Tester *test, int numVariants,
+    llvm::function_ref<void(ModuleOp, int, int)> transform) {
+  ReductionTreeUtils::createVariants(parent, test, numVariants, transform);
+}

mlir/tools/mlir-reduce/ReductionNode.cpp

 /// Sets up the metadata and links the node to its parent.
 ReductionNode::ReductionNode(ModuleOp module, ReductionNode *parent)
     : module(module), evaluated(false) {
 
   if (parent != nullptr)
     parent->linkVariant(this);
 }
 
+ReductionNode::ReductionNode(ModuleOp module, ReductionNode *parent,
+                             std::vector<int> transformSpace)
+    : module(module), evaluated(false), transformSpace(transformSpace) {
+
+  if (parent != nullptr)
+    parent->linkVariant(this);
+}
+
 /// Calculates and updates the size and interesting values of the module.
 void ReductionNode::measureAndTest(const Tester *test) {
   SmallString<128> filepath;
   int fd;
 
   // Print module to temprary file.
   std::error_code ec =
       llvm::sys::fs::createTemporaryFile("mlir-reduce", "mlir", fd, filepath);

 /// Returns the pointers to the child variants.
 ReductionNode *ReductionNode::getVariant(unsigned long index) const {
   if (index < variants.size())
     return variants[index].get();
 
   return nullptr;
 }
 
+/// Returns the number of child variants.
+int ReductionNode::variantsSize() const { return variants.size(); }
+
 /// Returns true if the child variants vector is empty.
 bool ReductionNode::variantsEmpty() const { return variants.empty(); }
 
 /// Link a child variant node.
 void ReductionNode::linkVariant(ReductionNode *newVariant) {
   std::unique_ptr<ReductionNode> ptrVariant(newVariant);
   variants.push_back(std::move(ptrVariant));
 }
for (auto &var : variants) {
     } else {
       break;
     }
   }
 
   // Remove uninteresting variants.
   variants.resize(interestingCount);
 }
+
+/// Returns the number of child variants.
+int ReductionNode::transformSpaceSize() { return transformSpace.size(); }
+
+/// Returns a vector of the transformable indices in the Module.
+const std::vector<int> ReductionNode::getTransformSpace() {
+  return transformSpace;
+}

mlir/tools/mlir-reduce/ReductionTreePass.cpp

-//===- ReductionTreePass.cpp - Reduction Tree Pass Implementation ---------===//
-//
-// 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 Pass. It provides a framework for
-// the implementation of different reduction passes in the MLIR Reduce tool. It
-// allows for custom specification of the variant generation behavior. It
-// implements methods that define the different possible traversals of the
-// reduction tree.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Reducer/ReductionTreePass.h"
-
-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());
-}

mlir/tools/mlir-reduce/ReductionTreeUtils.cpp

+//===- 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"
+
+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 the smallest node traversed so far in the reduction tree and
+/// print the debugging information for the currNode being traversed.
+void ReductionTreeUtils::processNode(ReductionNode *currNode,
+                                     ReductionNode *&smallestNode,
+                                     std::vector<int> path) {
+  LLVM_DEBUG(llvm::dbgs() << "\nTree Path: root");
+  for (auto i = path.begin(); i != path.end(); ++i)
+    LLVM_DEBUG(llvm::dbgs() << " -> " << *i);
+
+  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<int> ReductionTreeUtils::createTransformSpace(ModuleOp module,
+                                                          int numIndices) {
+  std::vector<int> transformSpace;
+
+  for (int i = 0; i < numIndices; ++i)
+    transformSpace.push_back(i);
+
+  return transformSpace;
+}
+
+/// Create the specified number of variants by applying the erase method
+/// to different ranges of indices in the parent module.
+void ReductionTreeUtils::createVariants(
+    ReductionNode *parent, const Tester *test, int numVariants,
+    llvm::function_ref<void(ModuleOp, int, int)> transform) {
+
+  std::vector<int> newTransformSpace;
+  ModuleOp module = parent->getModule();
+  int indexCount = parent->transformSpaceSize();
+
+  // No new variants can be created.
+  if (indexCount == 0)
+    return;
+
+  // Create a single variant by erasing the unique index.
+  if (indexCount == 1) {
+    ModuleOp variantModule = module.clone();
+    transform(variantModule, 0, 1);
+    new ReductionNode(variantModule, parent, newTransformSpace);
+
+    return;
+  }
+
+  // Create the specified number of variants.
+  for (int i = 0; i < numVariants; ++i) {
+    ModuleOp variantModule = module.clone();
+    newTransformSpace = parent->getTransformSpace();
+    int sectionSize = indexCount / numVariants;
+    int start = sectionSize * i;
+    int end = sectionSize * (i + 1);
+
+    if (i == numVariants - 1)
+      end = indexCount;
+
+    transform(variantModule, start, end);
+
+    for (int x = start; x < end; ++x) {
+      newTransformSpace.erase(newTransformSpace.begin() + start);
+    }
+
+    new ReductionNode(variantModule, parent, newTransformSpace);
+  }
+}
+

mlir/tools/mlir-reduce/mlir-reduce.cpp

if (!test.isInteresting(inputFilename))
     llvm::report_fatal_error(
         "Input test case does not exhibit interesting behavior");
 
   // Reduction pass pipeline.
   PassManager pm(&context);
 
   // Reduction tree pass with OpReducer variant generation and single path
   // traversal.
-  pm.addPass(
-      std::make_unique<ReductionTreePass<FunctionReducer, SinglePath>>(&test));
+  pm.addPass(std::make_unique<ReductionTreePass<OpReducer<FuncOp>, SinglePath>>(
+      &test));
 
   ModuleOp m = moduleRef.get().clone();
 
   if (failed(pm.run(m)))
     llvm::report_fatal_error("Error running the reduction pass pipeline");
 
   m.print(output->os());
   output->keep();
 
   return 0;
 }