Skip to content

Commit

Permalink
tower: update tower api to enable run passes on core:module and enabl…
Browse files Browse the repository at this point in the history
…e location transform
  • Loading branch information
kumarak committed Jan 23, 2025
1 parent e9b159c commit 060edb9
Show file tree
Hide file tree
Showing 4 changed files with 267 additions and 249 deletions.
37 changes: 17 additions & 20 deletions include/vast/CodeGen/CodeGenDriver.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@ VAST_UNRELAX_WARNINGS
#include "vast/Dialect/HighLevel/HighLevelDialect.hpp"
#include "vast/Dialect/HighLevel/HighLevelOps.hpp"

#include "vast/CodeGen/ScopeContext.hpp"
#include "vast/CodeGen/CodeGenModule.hpp"
#include "vast/CodeGen/ScopeContext.hpp"

#include "vast/Frontend/Options.hpp"

Expand All @@ -26,43 +26,39 @@ namespace vast::cg {

std::unique_ptr< codegen_builder > mk_codegen_builder(mcontext_t &mctx);

std::shared_ptr< meta_generator > mk_meta_generator(
acontext_t &actx, mcontext_t &mctx, const cc::vast_args &vargs
);
std::shared_ptr< meta_generator >
mk_meta_generator(acontext_t &actx, mcontext_t &mctx, const cc::vast_args &vargs);

std::shared_ptr< symbol_generator > mk_symbol_generator(
acontext_t &actx, mcontext_t &mctx, const cc::vast_args &vargs
);
std::shared_ptr< symbol_generator >
mk_symbol_generator(acontext_t &actx, mcontext_t &mctx, const cc::vast_args &vargs);

std::unique_ptr< mcontext_t > mk_mcontext();

void set_target_triple(core::module mod, std::string triple);
void set_source_language(core::module mod, cc::source_language lang);

owning_mlir_module_ref mk_wrapping_module(mcontext_t &mctx);
owning_mlir_module_ref mk_wrapping_module(acontext_t &actx, mcontext_t &mctx);

core::module mk_module(acontext_t &actx, mlir_module top);
core::module mk_module_with_attrs(acontext_t &actx, mlir_module top, cc::source_language lang);
core::module
mk_module_with_attrs(acontext_t &actx, mlir_module top, cc::source_language lang);

struct driver
{
explicit driver(
acontext_t &_actx
, mcontext_t &_mctx
, std::unique_ptr< codegen_builder > _bld
, std::shared_ptr< visitor_base > _visitor
acontext_t &_actx, mcontext_t &_mctx, std::unique_ptr< codegen_builder > _bld,
std::shared_ptr< visitor_base > _visitor
)
: actx(_actx)
, mctx(_mctx)
, bld(std::move(_bld))
, visitor(std::move(_visitor))
, top(mk_wrapping_module(mctx))
, top(mk_wrapping_module(actx, mctx))
, mod(mk_module_with_attrs(
actx, top.get(), cc::get_source_language(actx.getLangOpts())
))
actx, top.get(), cc::get_source_language(actx.getLangOpts())
))
, scope(symbols)
, generator(*bld, scoped_visitor_view(*visitor, scope))
{
, generator(*bld, scoped_visitor_view(*visitor, scope)) {
bld->module = mod;
bld->set_insertion_point_to_start(&mod.getBodyRegion());
}
Expand All @@ -80,6 +76,7 @@ namespace vast::cg {
owning_mlir_module_ref freeze();

mcontext_t &mcontext() { return mctx; }

acontext_t &acontext() { return actx; }

virtual bool verify();
Expand Down Expand Up @@ -115,6 +112,6 @@ namespace vast::cg {
};

std::unique_ptr< driver > mk_default_driver(
cc::action_options &opts, const cc::vast_args &vargs,
acontext_t &actx, mcontext_t &mctx);
cc::action_options &opts, const cc::vast_args &vargs, acontext_t &actx, mcontext_t &mctx
);
} // namespace vast::cg
53 changes: 24 additions & 29 deletions lib/vast/CodeGen/CodeGenDriver.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -34,13 +34,13 @@ namespace vast::cg {

owning_mlir_module_ref driver::freeze() { return std::move(top); }

// TODO this should not be needed the data layout should be emitted from cached types directly
dl::DataLayoutBlueprint emit_data_layout_blueprint(
const acontext_t &actx, const type_caching_proxy &types
) {
// TODO this should not be needed the data layout should be emitted from cached types
// directly
dl::DataLayoutBlueprint
emit_data_layout_blueprint(const acontext_t &actx, const type_caching_proxy &types) {
dl::DataLayoutBlueprint dl;

auto store_layout = [&] (const clang_type *orig, mlir_type vast_type) {
auto store_layout = [&](const clang_type *orig, mlir_type vast_type) {
if (orig->isFunctionType()) {
return;
}
Expand Down Expand Up @@ -83,7 +83,9 @@ namespace vast::cg {
auto list = std::dynamic_pointer_cast< visitor_list >(visitor);
for (auto node = list->head; node; node = node->next) {
if (auto types = std::dynamic_pointer_cast< type_caching_proxy >(node)) {
::vast::cg::emit_data_layout(mctx, mod, emit_data_layout_blueprint(actx, *types));
::vast::cg::emit_data_layout(
mctx, mod, emit_data_layout_blueprint(actx, *types)
);
}
}
}
Expand All @@ -94,9 +96,8 @@ namespace vast::cg {
return std::make_unique< codegen_builder >(&mctx);
}

std::shared_ptr< meta_generator > mk_meta_generator(
acontext_t *actx, mcontext_t *mctx, const cc::vast_args &vargs
) {
std::shared_ptr< meta_generator >
mk_meta_generator(acontext_t *actx, mcontext_t *mctx, const cc::vast_args &vargs) {
if (vargs.has_option(cc::opt::locs_as_meta_ids)) {
return std::make_shared< id_meta_gen >(actx, mctx);
}
Expand All @@ -123,28 +124,23 @@ namespace vast::cg {
// setup visitor list
const bool enable_unsupported = !vargs.has_option(cc::opt::disable_unsupported);

auto mg = mk_meta_generator(&actx, &mctx, vargs);
auto mg = mk_meta_generator(&actx, &mctx, vargs);
auto invalid_mg = mk_invalid_meta_generator(&mctx);
auto sg = mk_symbol_generator(actx);
auto policy = mk_codegen_policy(opts);
auto sg = mk_symbol_generator(actx);
auto policy = mk_codegen_policy(opts);

auto visitors = std::make_shared< visitor_list >()
| as_node_with_list_ref< attr_visitor_proxy >()
| as_node< type_caching_proxy >()
| as_node_with_list_ref< attr_visitor_proxy >() | as_node< type_caching_proxy >()
| as_node_with_list_ref< default_visitor >(
mctx, actx, *bld, std::move(mg), std::move(sg), std::move(policy)
mctx, actx, *bld, std::move(mg), std::move(sg), std::move(policy)
)
| optional(enable_unsupported,
as_node_with_list_ref< unsup_visitor >(
mctx, *bld, std::move(invalid_mg)
)
as_node_with_list_ref< unsup_visitor >(mctx, *bld, std::move(invalid_mg))
)
| as_node< unreach_visitor >();

// setup driver
auto drv = std::make_unique< driver >(
actx, mctx, std::move(bld), visitors
);
auto drv = std::make_unique< driver >(actx, mctx, std::move(bld), visitors);

drv->enable_verifier(!vargs.has_option(cc::opt::disable_vast_verifier));
return drv;
Expand Down Expand Up @@ -175,30 +171,29 @@ namespace vast::cg {
std::pair< loc_t, std::string > module_loc_name(mcontext_t &mctx, acontext_t &actx) {
// TODO use meta generator
if (auto path = get_path_to_source(actx); !path.empty()) {
return { mlir::FileLineColLoc::get(&mctx, path, 0, 0), path.str() };
return mlir::FileLineColLoc::get(&mctx, path, 0, 0);
}
return { mlir::UnknownLoc::get(&mctx), "unknown" };
}
} // namespace detail

owning_mlir_module_ref mk_wrapping_module(mcontext_t &mctx) {
return mlir::ModuleOp::create(mlir::UnknownLoc::get(&mctx));
owning_mlir_module_ref mk_wrapping_module(acontext_t &actx, mcontext_t &mctx) {
auto [loc, _] = detail::module_loc_name(mctx, actx);
return mlir::ModuleOp::create(loc);
}

core::module mk_module(acontext_t &actx, mlir_module top) {
mlir::OpBuilder bld(top);
bld.setInsertionPointToStart(top.getBody());

// TODO use symbol generator
auto mctx = top.getContext();
auto mctx = top.getContext();
auto [loc, name] = detail::module_loc_name(*mctx, actx);
return bld.create< core::module >(loc, name);
}

core::module mk_module_with_attrs(
acontext_t &actx, mlir_module top,
cc::source_language lang
) {
core::module
mk_module_with_attrs(acontext_t &actx, mlir_module top, cc::source_language lang) {
auto mod = mk_module(actx, top);

set_target_triple(mod, actx.getTargetInfo().getTriple().str());
Expand Down
49 changes: 34 additions & 15 deletions lib/vast/Tower/Tower.cpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
// Copyright (c) 2022-present, Trail of Bits, Inc.

#include "vast/Tower/Tower.hpp"
#include "vast/Dialect/Core/CoreOps.hpp"
#include "vast/Tower/Link.hpp"
#include "vast/Util/Common.hpp"

#include <gap/core/ranges.hpp>

Expand All @@ -12,28 +15,37 @@ namespace vast::tw {
module_storage &storage;

// Start empty and after each callback add to it.
conversion_passes_t path = {};
conversion_passes_t path = {};
// TODO: Remove.
conversion_path_t str_path = {};

std::vector< handle_t > handles;
link_vector steps;


explicit link_builder(location_info_t &li, module_storage &storage, handle_t root)
: li(li), storage(storage), handles{ root } {}

void runAfterPass(pass_ptr pass, operation op) override {
auto mod = mlir::dyn_cast< mlir_module >(op);
auto current = op;
while (current && !mlir::isa< mlir_module >(current)) {
current = op->getParentOp();
}
auto mod = mlir::dyn_cast< mlir_module >(current);
VAST_CHECK(mod, "Pass inside tower was not run on module!");

// Update locations so each operation now has a unique loc that also
// encodes backlink.
path.emplace_back(pass);
str_path.emplace_back(pass->getArgument().str());
transform_locations(li, str_path, mod);

owning_mlir_module_ref persistent = mlir::dyn_cast< mlir_module >(op->clone());
// Location transformation depends on the command line argument
// of passes. If it is empty, don't perform location transformation.
if (!pass->getArgument().empty()) {
str_path.emplace_back(pass->getArgument().str());
transform_locations(li, str_path, op);
}

// Clone the module to make it persistent
owning_mlir_module_ref persistent = mlir::dyn_cast< mlir_module >(mod->clone());

auto from = handles.back();
handles.emplace_back(storage.store(path, std::move(persistent)));
Expand All @@ -45,11 +57,15 @@ namespace vast::tw {

namespace {

link_vector construct_steps(const std::vector< handle_t > &handles, location_info_t &li) {
link_vector
construct_steps(const std::vector< handle_t > &handles, location_info_t &li) {
VAST_ASSERT(handles.size() >= 2);
link_vector out;
for (std::size_t i = 1; i < handles.size(); ++i)
out.emplace_back(std::make_unique< conversion_step >(handles[i - 1], handles[i], li));
for (std::size_t i = 1; i < handles.size(); ++i) {
out.emplace_back(
std::make_unique< conversion_step >(handles[i - 1], handles[i], li)
);
}
return out;
}

Expand All @@ -73,29 +89,32 @@ namespace vast::tw {

link_ptr tower::apply(handle_t root, location_info_t &li, mlir::PassManager &requested_pm) {
std::vector< mlir::Pass * > requested_passes;
for (auto &p : requested_pm.getPasses())
for (auto &p : requested_pm.getPasses()) {
requested_passes.push_back(&p);
}
auto [handles, suffix] = storage.get_maximum_prefix_path(requested_passes, root);

// This path is completely new.
if (handles.empty())
if (handles.empty()) {
return std::make_unique< fat_link >(mk_full_path(root, li, requested_pm));
}

auto as_steps = construct_steps(handles, li);

// This path is already present - construct a link.
if (suffix.empty())
if (suffix.empty()) {
return std::make_unique< fat_link >(std::move(as_steps));
}

auto pm = mlir::PassManager(requested_pm.getContext());
copy_passes(pm, suffix);

auto new_steps = mk_full_path(handles.back(), li, pm);
// TODO: Update with newer stdlib
as_steps.insert(
as_steps.end(),
std::make_move_iterator(new_steps.begin()),
std::make_move_iterator(new_steps.end()));
as_steps.end(), std::make_move_iterator(new_steps.begin()),
std::make_move_iterator(new_steps.end())
);

return std::make_unique< fat_link >(std::move(as_steps));
}
Expand Down
Loading

0 comments on commit 060edb9

Please sign in to comment.