Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Refactor] Decouple concrete passes from driver #305

Merged
merged 1 commit into from
Apr 18, 2021
Merged

[Refactor] Decouple concrete passes from driver #305

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 7 additions & 25 deletions armory/analyzer/driver.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,8 @@
#include "halo/lib/ir/ir_builder.h"
#include "halo/lib/parser/parser.h"
#include "halo/lib/pass/pass_manager.h"
#include "halo/lib/transforms/analyzer.h"
#include "halo/lib/transforms/caffeextension_legalizer.h"
#include "halo/lib/transforms/dce.h"
#include "halo/lib/transforms/input_legalizer.h"
#include "halo/lib/transforms/inst_simplify.h"
#include "halo/lib/transforms/onnxextension_legalizer.h"
#include "halo/lib/transforms/tfextension_legalizer.h"
#include "halo/lib/transforms/type_legalizer.h"
#include "halo/utils/cl_options.h"
#include "halo/utils/passes_helper.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Path.h"
Expand Down Expand Up @@ -60,24 +53,13 @@ static void PopulatePassesAndRun(GlobalContext& ctx, Module& m,
Parser::Format format) {
PassManager pm(ctx);
std::vector<std::string> input_shapes(InputsShape.begin(), InputsShape.end());
pm.AddPass<InputLegalizer>(batch.getValue(), input_shapes,
PreprocessScale.getValue());
if (format == Parser::Format::CAFFE) {
pm.AddPass<CAFFEExtensionLegalizer>();
} else if (format == Parser::Format::TENSORFLOW) {
pm.AddPass<TFExtensionLegalizer>();
} else {
HLCHECK(format == Parser::Format::ONNX);
pm.AddPass<ONNXExtensionLegalizer>();
}
pm.AddPass<DCE>();
pm.AddPass<TypeLegalizer>(true);
pm.AddPass<InstSimplify>(true, true, false, false, false, false);
auto analyzer = pm.AddPass<Analyzer>();
Fusion::Options fusion_opts;
Opts opts;
PopulateOptPasses(&pm, "cxx", input_shapes, {}, {}, batch, "",
ReorderChannel::ChannelOrder::None, false, false, format,
opts, fusion_opts);
pm.AddAnalyzerPass(&std::cout);
pm.Run(&m);
if (PrintAnalysisReport) {
analyzer->WriteCSVReport(std::cout);
}
}

int main(int argc, char** argv) {
Expand Down
234 changes: 41 additions & 193 deletions driver/driver.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,29 +23,10 @@
#include "halo/lib/ir/ir_builder.h"
#include "halo/lib/parser/parser.h"
#include "halo/lib/pass/pass_manager.h"
#include "halo/lib/quantizer/weights_quantizer.h"
#include "halo/lib/target/cpu/arm/binary/arm_llvmir_codegen.h"
#include "halo/lib/target/cpu/riscv/binary/riscv_llvmir_codegen.h"
#include "halo/lib/target/cpu/x86/binary/x86_llvmir_codegen.h"
#include "halo/lib/target/generic_cxx/generic_cxx_codegen.h"
#include "halo/lib/target/generic_llvmir/generic_llvmir_codegen.h"
#include "halo/lib/target/triton/triton_config_writer.h"
#include "halo/lib/transforms/caffeextension_legalizer.h"
#include "halo/lib/transforms/dce.h"
#include "halo/lib/transforms/device_placement.h"
#include "halo/lib/transforms/fusion.h"
#include "halo/lib/transforms/input_legalizer.h"
#include "halo/lib/transforms/input_rewriter.h"
#include "halo/lib/transforms/inst_simplify.h"
#include "halo/lib/transforms/onnxextension_legalizer.h"
#include "halo/lib/transforms/output_rewriter.h"
#include "halo/lib/transforms/reorder_channel.h"
#include "halo/lib/transforms/splitting.h"
#include "halo/lib/transforms/tfextension_legalizer.h"
#include "halo/lib/transforms/tfliteextension_legalizer.h"
#include "halo/lib/transforms/type_legalizer.h"
#include "halo/lib/transforms/typecast.h"
#include "halo/utils/cl_options.h"
#include "halo/utils/passes_helper.h"
#include "halo/version.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringSwitch.h"
Expand Down Expand Up @@ -248,177 +229,6 @@ static llvm::cl::opt<bool> CheckModel("check-model",
#include "halo/lib/ir/fusion.cc.inc"
#undef HALO_FUSION_CMD_OPTIONS_DECL

static void PopulateCodeGenPasses(PassManager* pm, std::ostream* out_code,
std::ostream* out_constants,
std::ostream* out_header,
std::ostream* out_dynamic_check,
bool is_c_or_cxx_output,
bool is_binary_output) {
auto constant_storage =
GenericLLVMIRCodeGen::ConstantDataStorage::DefinedAsStatic;
if (SeparateConstants) {
constant_storage =
GenericLLVMIRCodeGen::ConstantDataStorage::DeclaredAsExternal;
}

CodeGen* cg = nullptr;
if (is_c_or_cxx_output) {
Opts opts(BF16Mode);
if (llvm::StringRef(Target).startswith_lower("cc")) {
opts.dialect = Dialect::C99;
}
opts.print_mem_stats = PrintMemStats;
opts.emit_value_reset = EmitValueReset;
opts.exec_mode = ExecMode.getValue();
opts.emit_value_id_as_int = EmitValueIDAsInt;
opts.emit_inference_func_sig = EmitInferenceFunctionSignature;
opts.emit_dynamic_batch = (Batch.getValue() == kDynamicBatchSize);
opts.fp16_mode = EnableFP16;
opts.max_batch_size = MaxBatch.getValue();
opts.min_batch_size = MinBatch.getValue();
opts.opt_batch_size = OptBatch.getValue();
opts.check_model = CheckModel;
opts.enable_ipu_device = EnableIpuDevice;
opts.use_ipu_model = UseIpuModel;
opts.ipu_num = IpuNum;
opts.batches_per_step = BatchesPerStep;

pm->AddPass<WeightsQuantizer>(QuantWeights.getValue(), PGQFile.getValue());
cg = pm->AddPass<GenericCXXCodeGen>(std::ref(*out_code),
std::ref(*out_header),
std::ref(*out_dynamic_check), opts);
cg->SetAPI(Api);

if (EmitDataAsC) {
pm->AddPass<GenericCXXConstantWriter>(std::ref(*out_constants));
} else {
pm->AddPass<X86ConstantWriter>(std::ref(*out_constants));
}
if (EmitTritonConfig) {
pm->AddPass<TritonConfigWriter>(
TritonConfigFile.getValue(),
opts.emit_dynamic_batch ? MaxBatch.getValue() : 0);
}
return;
}

if (EmitLLVMIR) {
pm->AddPass<WeightsQuantizer>(QuantWeights.getValue(), PGQFile.getValue());
cg = pm->AddPass<GenericLLVMIRCodeGen>(constant_storage);
pm->AddPass<GenericLLVMIRWriter>(std::ref(*out_code), is_binary_output);
if (SeparateConstants && !EmitCodeOnly) {
pm->AddPass<GenericConstantWriter>(std::ref(*out_constants),
is_binary_output);
}
} else {
llvm::Triple triple(Target);
switch (triple.getArch()) {
case llvm::Triple::ArchType::x86:
case llvm::Triple::ArchType::x86_64: {
pm->AddPass<X86LLVMIRCodeGen>(
GenericLLVMIRCodeGen::ConstantDataStorage::DeclaredAsExternal);
pm->AddPass<X86BinaryWriter>(std::ref(*out_code));
if (SeparateConstants && !EmitCodeOnly) {
pm->AddPass<WeightsQuantizer>(QuantWeights.getValue(),
PGQFile.getValue());
pm->AddPass<X86ConstantWriter>(std::ref(*out_constants));
}
break;
}
case llvm::Triple::ArchType::aarch64: {
pm->AddPass<ARMLLVMIRCodeGen>(
GenericLLVMIRCodeGen::ConstantDataStorage::DeclaredAsExternal);
pm->AddPass<ARMBinaryWriter>(std::ref(*out_code));
if (SeparateConstants && !EmitCodeOnly) {
pm->AddPass<WeightsQuantizer>(QuantWeights.getValue(),
PGQFile.getValue());
pm->AddPass<ARMConstantWriter>(std::ref(*out_constants));
}
break;
}
case llvm::Triple::ArchType::riscv32:
case llvm::Triple::ArchType::riscv64: {
if (RISCVOpt) {
pm->AddPass<RISCVLLVMIRCodeGen>(
GenericLLVMIRCodeGen::ConstantDataStorage::DeclaredAsExternal,
"libRT_RISCV.a");
} else {
pm->AddPass<RISCVLLVMIRCodeGen>(
GenericLLVMIRCodeGen::ConstantDataStorage::DeclaredAsExternal);
}
pm->AddPass<RISCVBinaryWriter>(std::ref(*out_code));
if (SeparateConstants && !EmitCodeOnly) {
pm->AddPass<WeightsQuantizer>(QuantWeights.getValue(),
PGQFile.getValue());
pm->AddPass<RISCVConstantWriter>(std::ref(*out_constants));
}

break;
}

default: {
HLCHECK(0 && "Unsupported");
}
}
}
if (cg != nullptr) {
cg->SetAPI(Api);
}
}

static void PopulatePasses(PassManager* pm, std::ostream* out_code,
std::ostream* out_constants,
std::ostream* out_header,
std::ostream* out_dynamic_check,
bool is_c_or_cxx_output, bool is_binary_output,
Parser::Format format) {
std::vector<std::string> input_shapes(InputsShape.begin(), InputsShape.end());
pm->AddPass<InputLegalizer>(Batch.getValue(), input_shapes,
PreprocessScale.getValue());
if (!Outputs.empty()) {
std::vector<std::string> outputs(Outputs.begin(), Outputs.end());
pm->AddPass<OutputRewriter>(outputs);
}
if (format == Parser::Format::CAFFE) {
pm->AddPass<CAFFEExtensionLegalizer>();
} else if (format == Parser::Format::TENSORFLOW) {
pm->AddPass<TFExtensionLegalizer>();
} else if (format == Parser::Format::TFLITE) {
HLCHECK(format == Parser::Format::TFLITE);
pm->AddPass<TFLITEExtensionLegalizer>();
} else {
HLCHECK(format == Parser::Format::ONNX);
pm->AddPass<ONNXExtensionLegalizer>();
}
pm->AddPass<DCE>();
pm->AddPass<TypeLegalizer>(true);
if (!Inputs.empty()) {
std::vector<std::string> inputs(Inputs.begin(), Inputs.end());
pm->AddPass<InputRewriter>(inputs);
}
auto fusion_opts = GetFusionOptions();
pm->AddPass<InstSimplify>(
llvm::StringRef(Target).startswith("cxx"), DisableBroadcasting.getValue(),
RemoveInputTranspose.getValue(), RemoveOutputTranspose.getValue(),
DisableConvBN.getValue(), fusion_opts.ConvBias);
if (ReorderChannelLayout != ReorderChannel::ChannelOrder::None) {
pm->AddPass<ReorderChannel>(ReorderChannelLayout ==
ReorderChannel::ChannelOrder::ChannelFirst);
}
pm->AddPass<Fusion>(fusion_opts);
if (SplitFunction) {
pm->AddPass<Splitting>();
pm->AddPass<DevicePlacement>();
}
if (!DisableTypeCast) {
pm->AddPass<TypeCast>();
}

PopulateCodeGenPasses(pm, out_code, out_constants, out_header,
out_dynamic_check, is_c_or_cxx_output,
is_binary_output);
}

static bool FormatCode(const std::string& filename) {
if (filename.empty() || filename == "-") {
return false;
Expand Down Expand Up @@ -540,11 +350,49 @@ int main(int argc, char** argv) {
out_dynamic_check = &of_dynamic_check;
}

PopulatePasses(&pm, out_code, out_constants, out_header, out_dynamic_check,
is_c_or_cxx_output, is_binary_output, format);
Opts cg_opts;
cg_opts.bf16_mode = BF16Mode;
cg_opts.print_mem_stats = PrintMemStats;
cg_opts.emit_value_reset = EmitValueReset;
cg_opts.exec_mode = ExecMode.getValue();
cg_opts.emit_value_id_as_int = EmitValueIDAsInt;
cg_opts.emit_inference_func_sig = EmitInferenceFunctionSignature;
cg_opts.emit_dynamic_batch = (Batch.getValue() == kDynamicBatchSize);
cg_opts.fp16_mode = EnableFP16;
cg_opts.max_batch_size = MaxBatch.getValue();
cg_opts.min_batch_size = MinBatch.getValue();
cg_opts.opt_batch_size = OptBatch.getValue();
cg_opts.check_model = CheckModel;
cg_opts.enable_ipu_device = EnableIpuDevice;
cg_opts.use_ipu_model = UseIpuModel;
cg_opts.ipu_num = IpuNum;
cg_opts.batches_per_step = BatchesPerStep;
cg_opts.api = Api;
cg_opts.disable_broadcasting = DisableBroadcasting;
cg_opts.separate_constants = SeparateConstants;
cg_opts.disable_conv_bn = DisableConvBN;
cg_opts.remove_input_transpose = RemoveInputTranspose;
cg_opts.remove_output_transpose = RemoveOutputTranspose;

if (is_c_or_cxx_output) {
ctx.SetTargetTriple("x86_64"); // For binary constant writer.
if (llvm::StringRef(Target).startswith_lower("cc")) {
cg_opts.dialect = Dialect::C99;
}
}
std::vector<std::string> input_shapes(InputsShape.begin(), InputsShape.end());
std::vector<std::string> inputs(Inputs.begin(), Inputs.end());
std::vector<std::string> outputs(Outputs.begin(), Outputs.end());
const auto& fusion_opts = GetFusionOptions();

PopulateOptPasses(&pm, Target, input_shapes, inputs, outputs, Batch,
PreprocessScale, ReorderChannelLayout, SplitFunction,
DisableTypeCast, format, cg_opts, fusion_opts);
PopulateCodeGenPasses(&pm, out_code, out_constants, out_header,
out_dynamic_check, Target, is_c_or_cxx_output,
is_binary_output, EmitDataAsC, EmitCodeOnly, EmitLLVMIR,
EmitTritonConfig, TritonConfigFile, QuantWeights,
PGQFile, RISCVOpt, cg_opts);

auto status = pm.Run(&m);

Expand Down
63 changes: 62 additions & 1 deletion include/halo/lib/pass/pass_manager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@
#include "halo/lib/framework/global_context.h"
#include "halo/lib/ir/module.h"
#include "halo/lib/pass/pass.h"
#include "halo/lib/target/generic_cxx/generic_cxx_codegen.h"
#include "halo/lib/target/generic_llvmir/generic_llvmir_codegen.h"
#include "halo/lib/transforms/fusion.h"

namespace halo {

Expand All @@ -39,7 +42,7 @@ class PassManager final {

/// Add a pass to the pass manager.
template <typename T, typename... TS>
T* AddPass(TS... args) {
T* AddPass(TS&... args) {
auto pass = std::make_unique<T>(args...);
T* ret = static_cast<T*>(pass.get());
Add(std::move(pass));
Expand All @@ -53,6 +56,64 @@ class PassManager final {
void Print(std::ostream& os) const;

void Dump() const;
Pass* AddAnalyzerPass(std::ostream* os);
Pass* AddARMBinaryWriterPass(std::ostream& os);
Pass* AddARMConstantWriterPass(std::ostream& os);
Pass* AddARMLLVMIRCodeGenPass(
GenericLLVMIRCodeGen::ConstantDataStorage constant_data_storage);
Pass* AddCAFFEExtensionLegalizerPass();
Pass* AddDCEPass();
Pass* AddDevicePlacementPass();
Pass* AddFusionPass(const Fusion::Options& opts);
Pass* AddGenericConstantWriterPass(std::ostream& os, bool bitcode_format);
Pass* AddGenericCXXConstantWriterPass(std::ostream& os);
Pass* AddGenericCXXCodeGenPass(std::ostream& os, std::ostream& header_os);
Pass* AddGenericCXXCodeGenPass(std::ostream& os, std::ostream& header_os,
std::ostream& dynamic_check_os,
const Opts& opts);
Pass* AddGenericLLVMIRCodeGenPass();
Pass* AddGenericLLVMIRCodeGenPass(
GenericLLVMIRCodeGen::ConstantDataStorage constant_data_storage);
Pass* AddGenericLLVMIRCodeGenPass(
const std::string& name,
GenericLLVMIRCodeGen::ConstantDataStorage constant_data_storage);
Pass* AddGenericLLVMIRWriterPass(std::ostream& os, bool bitcode_format);
Pass* AddInputLegalizerPass(int batch_size,
const std::vector<std::string>& inputs_shapes,
const std::string& scale_str);
Pass* AddInputRewriterPass(const std::vector<std::string>& inputs);
Pass* AddInstSimplifyPass();
Pass* AddInstSimplifyPass(bool simplify_for_preprocess,
bool disable_broadcasting,
bool remove_input_transpose,
bool remove_output_transpose, bool disable_conv_bn,
bool fuse_conv_bias);
Pass* AddONNXExtensionLegalizerPass();
Pass* AddOutputRewriterPass(const std::vector<std::string>& outputs);
Pass* AddReorderChannelPass(bool channel_first);
Pass* AddRISCVBinaryWriterPass(std::ostream& os);
Pass* AddRISCVConstantWriterPass(std::ostream& os);
Pass* AddRISCVLLVMIRCodeGenPass(
GenericLLVMIRCodeGen::ConstantDataStorage constant_data_storage);

Pass* AddRISCVLLVMIRCodeGenPass(
GenericLLVMIRCodeGen::ConstantDataStorage constant_data_storage,
std::string rt_lib_name);
Pass* AddSplittingPass();
Pass* AddTFExtensionLegalizerPass();
Pass* AddTFLiteExtensionLegalizerPass();
Pass* AddTritonConfigWriterPass(const std::string& filename,
int max_batch_size);
Pass* AddTypeCastPass();
Pass* AddTypeLegalizerPass();
Pass* AddTypeLegalizerPass(bool relaxed);
Pass* AddWeightsQuantizerPass(CodeGen::Quantization quant,
const std::string& file);
Pass* AddX86BinaryWriterPass(std::ostream& os);
Pass* AddX86ConstantWriterPass(std::ostream& os);
Pass* AddX86LLVMIRCodeGenPass();
Pass* AddX86LLVMIRCodeGenPass(
GenericLLVMIRCodeGen::ConstantDataStorage constant_data_storage);

private:
Pass* Add(std::unique_ptr<ModulePass> pass);
Expand Down
Loading