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

31 ↗

(On Diff #284747)

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.

68 ↗

(On Diff #284747)

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

38

Why not const Tester& ? Is null allowed?

42

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

52

These would all be of OpType?

54

Lets make this OpType as refers to class

mlir/include/mlir/Reducer/ReductionTreePass.h

101 ↗

(On Diff #286398)

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

133 ↗

(On Diff #286398)

Same here

mlir/include/mlir/Reducer/ReductionTreeUtils.h

31 ↗

(On Diff #286398)

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

35 ↗

(On Diff #286398)

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

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

1 ↗

(On Diff #286398)

Line update?

mlir/tools/mlir-reduce/ReductionTreeUtils.cpp

35 ↗

(On Diff #286398)

Why not just range based for?

69 ↗

(On Diff #286398)

Should this be wrapped in LLVM_DEBUG too?

122 ↗

(On Diff #286398)

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/OpReducer.h

 // 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>()) {

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

 //===- mlir-reduce.cpp - The MLIR 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 implements the general framework of the MLIR reducer tool. It
 // parses the command line arguments, parses the initial MLIR test case and sets
 // up the testing environment. It  outputs the most reduced test case variant
 // after executing the reduction passes.
 //
 //===----------------------------------------------------------------------===//
 
-#include <vector>
-
 #include "mlir/InitAllDialects.h"
 #include "mlir/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Reducer/OptReductionPass.h"
+#include "mlir/Reducer/Passes/OpReducer.h"
 #include "mlir/Reducer/ReductionNode.h"
 #include "mlir/Reducer/ReductionTreePass.h"
 #include "mlir/Reducer/Tester.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/ToolOutputFile.h"
if (passTestSpecifier == "DCE") {
     // Opt Reduction Pass with SymbolDCEPass as opt pass.
     pm.addPass(std::make_unique<OptReductionPass>(&test, &context,
                                                   createSymbolDCEPass()));
 
   } else if (passTestSpecifier == "function-reducer") {
 
     // Reduction tree pass with OpReducer variant generation and single path
     // traversal.
-    pm.addPass(std::make_unique<ReductionTreePass<FunctionReducer, SinglePath>>(
+    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;
 }