plugin.h

Go to the documentation of this file.

#pragma once
 
#include <memory>
#include <optional>
#include <string>
 
#include "sdfg/codegen/dispatchers/node_dispatcher_registry.h"
#include "sdfg/passes/scheduler/omp_scheduler.h"
#include "sdfg/passes/scheduler/scheduler_registry.h"
#include "sdfg/plugins/plugins.h"
#include "sdfg/targets/omp/codegen/omp_map_dispatcher.h"
#include "sdfg/targets/omp/schedule.h"
#include "sdfg/tiles/tile_target_registry.h"
 
namespace sdfg {
namespace omp {
 
class OMPTileTarget : public tiles::TileTarget {
public:
    std::optional<tiles::AxisSchedule> classify(const structured_control_flow::ScheduleType& sched) const override {
        if (sched.category() == structured_control_flow::ScheduleTypeCategory::None) {
            return std::nullopt;
        }
        return tiles::AxisSchedule(tiles::Level::Device, space(tiles::Level::Device), has_scratchpad());
    }
 
    types::StorageType storage_type(tiles::Space space) const override {
        return space == tiles::Space::Register ? types::StorageType::CPU_Stack() : types::StorageType::CPU_Heap();
    }

    tiles::Space space(tiles::Level) const override {
        return tiles::Space::Global;
    }
 
    bool supports_cooperative_staging(const structured_control_flow::ScheduleType&) const override {
        return false;
    }
 
    unsigned lane_width() const override {
        return 1;
    }

    data_flow::ImplementationType implementation_type() const override {
        return data_flow::ImplementationType_NONE;
    }
};
 
inline void register_omp_plugin(plugins::Context& context) {
    context.get_map_dispatcher_registry().register_map_dispatcher(
        ScheduleType_OMP::value(),
        [](codegen::LanguageExtension& language_extension,
           StructuredSDFG& sdfg,
           analysis::AnalysisManager& analysis_manager,
           structured_control_flow::Map& node,
           codegen::InstrumentationPlan& instrumentation_plan,
           codegen::ArgCapturePlan& arg_capture_plan) {
            return std::make_unique<OMPMapDispatcher>(
                language_extension, sdfg, analysis_manager, node, instrumentation_plan, arg_capture_plan
            );
        }
    );
 
    context.get_scheduler_registry()
        .register_loop_scheduler<passes::scheduler::OMPScheduler>(passes::scheduler::OMPScheduler::target());
 
    context.get_tile_target_registry().register_target(ScheduleType_OMP::value(), std::make_shared<OMPTileTarget>());
}
 
inline void register_omp_plugin() {
    auto ctx = plugins::Context::global_context();
    register_omp_plugin(ctx);
}
 
} // namespace omp
} // namespace sdfg