[Reducer] Add bytecode support to iree-reduce (#15079)
diff --git a/compiler/src/iree/compiler/API/Internal/IREEReduceToolEntryPoint.cpp b/compiler/src/iree/compiler/API/Internal/IREEReduceToolEntryPoint.cpp index 584e65d..e2a60aa 100644 --- a/compiler/src/iree/compiler/API/Internal/IREEReduceToolEntryPoint.cpp +++ b/compiler/src/iree/compiler/API/Internal/IREEReduceToolEntryPoint.cpp
@@ -12,6 +12,7 @@ #include "llvm/Support/Process.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Support/ToolOutputFile.h" +#include "mlir/Bytecode/BytecodeWriter.h" #include "mlir/IR/AsmState.h" #include "mlir/IR/Dialect.h" #include "mlir/IR/MLIRContext.h" @@ -51,19 +52,28 @@ llvm::cl::OptionCategory ireeReduceCategory("iree-reduce options"); - static llvm::cl::opt<std::string> testScript(cl::Positional, cl::Required, - cl::desc("<test script>"), - cl::cat(ireeReduceCategory)); + llvm::cl::opt<std::string> testScript(cl::Positional, cl::Required, + cl::desc("<test script>"), + cl::cat(ireeReduceCategory)); - static cl::opt<std::string> inputFilename( - cl::Positional, cl::desc("<input file>"), cl::init("-"), - llvm::cl::cat(ireeReduceCategory)); + cl::opt<std::string> inputFilename(cl::Positional, cl::desc("<input file>"), + cl::init("-"), + llvm::cl::cat(ireeReduceCategory)); - static cl::opt<std::string> outputFilename( + cl::opt<std::string> outputFilename( "o", cl::desc("Output filename for the reduced test case."), cl::value_desc("filename"), cl::init("-"), llvm::cl::cat(ireeReduceCategory)); + cl::opt<bool> useBytecodeForTesting( + "use-bytecode", + cl::desc("Use bytecode as input to the interesting script."), + cl::init(false), llvm::cl::cat(ireeReduceCategory)); + + cl::opt<bool> outputAsBytecode( + "output-bytecode", cl::desc("Output the final output as bytecode."), + cl::init(false), llvm::cl::cat(ireeReduceCategory)); + llvm::cl::HideUnrelatedOptions(ireeReduceCategory); InitLLVM y(argc, argv); @@ -88,12 +98,22 @@ return failure(); } - Operation *newModule = - ireeRunReducingStrategies(std::move(module), testScript); + ReducerConfig config(testScript, useBytecodeForTesting); + Operation *newModule = ireeRunReducingStrategies(std::move(module), config); module = OwningOpRef<Operation *>(newModule); - // Print module to output file. - module->print(output->os()); + if (outputAsBytecode) { + // Write bytecode to output file. + BytecodeWriterConfig config; + LogicalResult result = + writeBytecodeToFile(module.get(), output->os(), config); + if (failed(result)) { + llvm::report_fatal_error("Failed to write bytecode to output file"); + } + } else { + // Write MLIR to file. + module->print(output->os()); + } // Keep the output file if the invocation of MlirOptMain was successful. output->keep();
diff --git a/compiler/src/iree/compiler/Reducer/BUILD.bazel b/compiler/src/iree/compiler/Reducer/BUILD.bazel index 4cb8bcb..2af2408 100644 --- a/compiler/src/iree/compiler/Reducer/BUILD.bazel +++ b/compiler/src/iree/compiler/Reducer/BUILD.bazel
@@ -23,6 +23,7 @@ deps = [ "//compiler/src/iree/compiler/Reducer/Framework:ReducerFramework", "//compiler/src/iree/compiler/Reducer/Strategies:DeltaStrategies", + "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ],
diff --git a/compiler/src/iree/compiler/Reducer/CMakeLists.txt b/compiler/src/iree/compiler/Reducer/CMakeLists.txt index ca5afb9..b6de442 100644 --- a/compiler/src/iree/compiler/Reducer/CMakeLists.txt +++ b/compiler/src/iree/compiler/Reducer/CMakeLists.txt
@@ -18,6 +18,7 @@ SRCS "iree_reduce_lib.cc" DEPS + LLVMSupport MLIRIR MLIRSupport iree::compiler::Reducer::Framework::ReducerFramework
diff --git a/compiler/src/iree/compiler/Reducer/Framework/Oracle.cpp b/compiler/src/iree/compiler/Reducer/Framework/Oracle.cpp index 5e944cb..6789f0f 100644 --- a/compiler/src/iree/compiler/Reducer/Framework/Oracle.cpp +++ b/compiler/src/iree/compiler/Reducer/Framework/Oracle.cpp
@@ -10,6 +10,7 @@ #include "llvm/Support/FileSystem.h" #include "llvm/Support/Program.h" #include "llvm/Support/ToolOutputFile.h" +#include "mlir/Bytecode/BytecodeWriter.h" using namespace mlir; using namespace mlir::iree_compiler; @@ -29,8 +30,9 @@ // Print module to a temporary file. SmallString<128> filepath; int fd; + std::string extension = useBytecode ? "mlirbc" : "mlir"; std::error_code ec = - llvm::sys::fs::createTemporaryFile("oracle", "mlir", fd, filepath); + llvm::sys::fs::createTemporaryFile("oracle", extension, fd, filepath); if (ec) { llvm::report_fatal_error(llvm::Twine("Failed to create temporary file: ") + @@ -38,7 +40,21 @@ } llvm::ToolOutputFile output(filepath, fd); - workItem.getModule()->print(output.os()); + + if (useBytecode) { + // Write bytecode to file. + BytecodeWriterConfig config; + LogicalResult result = + writeBytecodeToFile(workItem.getModule(), output.os(), config); + if (failed(result)) { + llvm::report_fatal_error( + llvm::Twine("Failed to write bytecode to file: ") + filepath); + } + } else { + // Write MLIR to file. + workItem.getModule()->print(output.os()); + } + output.os().close(); if (output.os().has_error()) {
diff --git a/compiler/src/iree/compiler/Reducer/Framework/Oracle.h b/compiler/src/iree/compiler/Reducer/Framework/Oracle.h index d40ef04..e045a8b 100644 --- a/compiler/src/iree/compiler/Reducer/Framework/Oracle.h +++ b/compiler/src/iree/compiler/Reducer/Framework/Oracle.h
@@ -15,12 +15,14 @@ class Oracle { public: - Oracle(StringRef testScript) : testScript(testScript) {} + Oracle(StringRef testScript, bool useBytecode) + : testScript(testScript), useBytecode(useBytecode) {} bool isInteresting(WorkItem &workItem); private: StringRef testScript; + bool useBytecode; }; } // namespace mlir::iree_compiler::Reducer
diff --git a/compiler/src/iree/compiler/Reducer/iree_reduce_lib.cc b/compiler/src/iree/compiler/Reducer/iree_reduce_lib.cc index 0090640..c0a8afc 100644 --- a/compiler/src/iree/compiler/Reducer/iree_reduce_lib.cc +++ b/compiler/src/iree/compiler/Reducer/iree_reduce_lib.cc
@@ -14,10 +14,10 @@ using namespace llvm; Operation *mlir::iree_compiler::Reducer::ireeRunReducingStrategies( - OwningOpRef<Operation *> module, StringRef testScript) { + OwningOpRef<Operation *> module, ReducerConfig &config) { ModuleOp root = dyn_cast<ModuleOp>(module.release()); WorkItem workItem(root); - Oracle oracle(testScript); + Oracle oracle(config.testScript, config.useBytecode); Delta delta(oracle, workItem); delta.runDeltaPass(reduceFlowDispatchOperandToResultDelta,
diff --git a/compiler/src/iree/compiler/Reducer/iree_reduce_lib.h b/compiler/src/iree/compiler/Reducer/iree_reduce_lib.h index a77fa9d..7ec8bee 100644 --- a/compiler/src/iree/compiler/Reducer/iree_reduce_lib.h +++ b/compiler/src/iree/compiler/Reducer/iree_reduce_lib.h
@@ -7,6 +7,7 @@ #ifndef IREE_COMPILER_TOOLS_IREE_REDUCER_LIB_H #define IREE_COMPILER_TOOLS_IREE_REDUCER_LIB_H +#include "llvm/ADT/StringRef.h" #include "mlir/IR/OwningOpRef.h" #include "mlir/Support/LLVM.h" @@ -16,8 +17,20 @@ namespace iree_compiler::Reducer { +struct ReducerConfig { + ReducerConfig() = delete; + + explicit ReducerConfig(StringRef testScript, bool useBytecode) + : testScript(testScript), useBytecode(useBytecode) {} + + // Path to the test script to run on the reduced program. + StringRef testScript; + // Flag to indicate whether the test script can use bytecode or not. + bool useBytecode; +}; + Operation *ireeRunReducingStrategies(OwningOpRef<Operation *> module, - StringRef testScript); + ReducerConfig &config); } // namespace iree_compiler::Reducer } // namespace mlir