Create Opt Reduction Pass

[msifontes]

Summary

Implement a pass for the MLIR Reduce tool that runs any MLIR transformation pass and only replaces the output if the transformed version is smaller and still interesting to the tool. Depends on the commit: 05859f7e00ff Create Reduction Tree Pass in the following revision: https://reviews.llvm.org/D83969.

Call the pass with SymbolDCEPass as optimization pass in the tool's main function.

[msifontes]

Include only the new changes in this commit.

[jpienaar]

mlir/include/mlir/Reducer/OptReductionPass.h

55

Nit: prefer member functions before data

mlir/include/mlir/Reducer/Passes.td

25

I don't see this constructor? Wouldn't one need to pass in a pass here?

28

Typo

mlir/tools/mlir-reduce/OptReductionPass.cpp

35

What about

if (ReductionNode* transformed = singleTransform())

utils.updateGoldenModule(module, transformed->getModule().clone());

?

52

Where is this memory reclaimed?

55

Prefer early exits:

if (initNode->variantsEmpty()) return initNode;

[msifontes]

Refactor Optimization Pass

Create two static reduction nodes instead of a two level reduction tree.

Add -test-pass option to main in order to test individual passes and different pipelines.

Modify tests to use the -test-pass option.

[msifontes]

NFC: Fix Typo

[msifontes]

Add missing new line in simple-test.mlir

[jpienaar]

Nice, thanks

mlir/include/mlir/Reducer/OptReductionPass.h

2

Need C++ file type designator here

19

I don't think this is needed here

22

Prefer to use include path here (e.g., mlir/Reducer/ReductionNode.h)

30

Move this to the cpp file

mlir/include/mlir/Reducer/Passes.td

23

s/mlir-opt/optimization/ ? (mlir-opt is the tool that can be used to test the passes, but the passes are not the tool's)

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

103–115

Could PassPipelineCLParser be reused here? (at least for the optimization passes). See ir-printing.mlir test as an example usage with specific prefix. Can also be refined later so add a TODO for now in that case. It feels like one needs the same concept here of a reduction pipeline, where I'd think we could reuse a lot of that parsing with just some extra pass registrations here so that the pass pipeline parsing could be used here.

[msifontes]

Apply requested changes

• Add TODO to use PassPipelineCLParser instead of -test-pass option to define specific pipelines in the command line.

• Formatting changes.

[msifontes]

NFC: Amend commit message

mlir/include/mlir/Reducer/OptReductionPass.h

+//===- OptReductionPass.h - Optimization Reduction Pass Wrapper -*- 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 Opt Reduction Pass Wrapper. It creates a pass to run
+// any optimization pass within it and only replaces the output module with the
+// transformed version if it is smaller and interesting.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_REDUCER_OPTREDUCTIONPASS_H
+#define MLIR_REDUCER_OPTREDUCTIONPASS_H
+
+#include "PassDetail.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Reducer/ReductionNode.h"
+#include "mlir/Reducer/ReductionTreePass.h"
+#include "mlir/Reducer/Tester.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/Support/Debug.h"
+
+namespace mlir {
+
+class OptReductionPass : public OptReductionBase<OptReductionPass> {
+public:
+  OptReductionPass(const Tester *test, MLIRContext *context,
+                   std::unique_ptr<Pass> optPass);
+
+  OptReductionPass(const OptReductionPass &srcPass);
+
+  /// Runs the pass instance in the pass pipeline.
+  void runOnOperation() override;
+
+private:
+  // Points to the context to be used in the pass manager.
+  MLIRContext *context;
+
+  // This is used to test the interesting behavior of the transformed module.
+  const Tester *test;
+
+  // Points to the mlir-opt pass to be called.
+  std::unique_ptr<Pass> optPass;
+};
+
+} // end namespace mlir
+
+#endif

mlir/include/mlir/Reducer/Passes.td

 
 #ifndef MLIR_REDUCER_PASSES
 #define MLIR_REDUCER_PASSES
 
 include "mlir/Pass/PassBase.td"
 
 def ReductionTree : Pass<"reduction-tree", "ModuleOp"> {
   let summary = "A general reduction tree pass for the MLIR Reduce Tool";
-  let constructor = "mlir::createReductionTreePass()";
 }
 
-#endif // MLIR_REDUCE_PASSES
+def OptReduction : Pass<"opt-reduction-pass", "ModuleOp"> {
+  let summary = "A reduction pass wrapper for optimization passes";
+}
+
+#endif // MLIR_REDUCER_PASSES

mlir/include/mlir/Reducer/ReductionTreePass.h

 /// 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);
+  static void updateGoldenModule(ModuleOp &golden, ModuleOp reduced);
 };
 
 /// 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>> {

mlir/test/mlir-reduce/dce-test.mlir

+// UNSUPPORTED: -windows-
+// RUN: mlir-reduce %s -test %S/failure-test.sh -pass-test DCE | FileCheck %s
+// This input should be reduced by the pass pipeline so that only
+// the @simple1 function remains as the other fucntions should be
+// removed by the dead code elimination pass.
+// CHECK-LABEL: func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
+
+// CHECK-NOT: func @dead_nested_function
+func @dead_private_function() attributes { sym_visibility = "private" }
+
+// CHECK-NOT: func @dead_nested_function
+func @dead_nested_function() attributes { sym_visibility = "nested" }
+
+func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
+  "test.crashOp" () : () -> ()
+  return
+}

mlir/test/mlir-reduce/multiple-function.mlir

 // UNSUPPORTED: -windows-
-// RUN: mlir-reduce %s -test %S/failure-test.sh | FileCheck %s
+// RUN: mlir-reduce %s -test %S/failure-test.sh -pass-test function-reducer | FileCheck %s
 // This input should be reduced by the pass pipeline so that only 
 // the @simple5 function remains as this is the shortest function 
 // containing the interesting behavior.
 
 // CHECK-NOT: func @simple1() {
 func @simple1() {
   return
 }

mlir/test/mlir-reduce/simple-test.mlir

 // UNSUPPORTED: -windows-
-// RUN: mlir-reduce %s -test %S/test.sh
+// RUN: mlir-reduce %s -test %S/test.sh -pass-test function
 
 func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = alloc() : memref<2xf32>
   br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   return
-}
+}
 

mlir/test/mlir-reduce/single-function.mlir

 // RUN: mlir-opt %s
-// RUN: not mlir-opt %s -test-mlir-reducer
+// RUN: not mlir-opt %s -test-mlir-reducer -pass-test function-reducer
 
 func @test() {
   "test.crashOp"() : () -> ()
   return
 }

mlir/tools/mlir-reduce/CMakeLists.txt

set(LIBS
   MLIRPass
   MLIRReduce
   MLIRSupport
   MLIRTransforms
   MLIRTransformUtils
   )
 
 add_llvm_tool(mlir-reduce
+  OptReductionPass.cpp
   Passes/FunctionReducer.cpp
   ReductionNode.cpp
   ReductionTreePass.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/OptReductionPass.cpp

+//===- OptReductionPass.cpp - Optimization Reduction Pass Wrapper ---------===//
+//
+// 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 Opt Reduction Pass class. It creates a pass to run
+// any optimization pass within it and only replaces the output module with the
+// transformed version if it is smaller and interesting.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Reducer/OptReductionPass.h"
+
+#define DEBUG_TYPE "mlir-reduce"
+
+using namespace mlir;
+
+OptReductionPass::OptReductionPass(const Tester *test, MLIRContext *context,
+                                   std::unique_ptr<Pass> optPass)
+    : context(context), test(test), optPass(std::move(optPass)) {}
+
+OptReductionPass::OptReductionPass(const OptReductionPass &srcPass)
+    : test(srcPass.test), optPass(srcPass.optPass.get()) {}
+
+/// Runs the pass instance in the pass pipeline.
+void OptReductionPass::runOnOperation() {
+  LLVM_DEBUG(llvm::dbgs() << "\nOptimization Reduction pass: ");
+  LLVM_DEBUG(llvm::dbgs() << optPass.get()->getName() << "\nTesting:\n");
+
+  ModuleOp module = this->getOperation();
+  ModuleOp moduleVariant = module.clone();
+  PassManager pmTransform(context);
+  pmTransform.addPass(std::move(optPass));
+
+  if (failed(pmTransform.run(moduleVariant)))
+    return;
+
+  ReductionNode original(module, nullptr);
+  original.measureAndTest(test);
+
+  ReductionNode reduced(moduleVariant, nullptr);
+  reduced.measureAndTest(test);
+
+  if (reduced.isInteresting() && reduced.getSize() < original.getSize()) {
+    ReductionTreeUtils::updateGoldenModule(module, reduced.getModule().clone());
+    LLVM_DEBUG(llvm::dbgs() << "\nSuccessful Transformed version\n\n");
+  } else {
+    LLVM_DEBUG(llvm::dbgs() << "\nUnsuccessful Transformed version\n\n");
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "Pass Complete\n\n");
+}

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

 //===----------------------------------------------------------------------===//
 
 #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/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"
static llvm::cl::list<std::string>
     testArguments("test-args", llvm::cl::ZeroOrMore,
                   llvm::cl::desc("Testing script arguments"));
 
 static llvm::cl::opt<std::string>
     outputFilename("o",
                    llvm::cl::desc("Output filename for the reduced test case"),
                    llvm::cl::init("-"));
 
+// TODO: Use PassPipelineCLParser to define pass pieplines in the command line.
+static llvm::cl::opt<std::string>
+    passTestSpecifier("pass-test",
+                      llvm::cl::desc("Indicate a specific pass to be tested"));
+
 // Parse and verify the input MLIR file.
 static LogicalResult loadModule(MLIRContext &context, OwningModuleRef &module,
                                 StringRef inputFilename) {
   module = parseSourceFile(inputFilename, &context);
   if (!module)
     return failure();
 
   return success();
int main(int argc, char **argv) {
 
   if (!test.isInteresting(inputFilename))
     llvm::report_fatal_error(
         "Input test case does not exhibit interesting behavior");
 
   // Reduction pass pipeline.
   PassManager pm(&context);
 
+  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>>(&test));
+    pm.addPass(std::make_unique<ReductionTreePass<FunctionReducer, 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;
 }