gpu_offload_schedule_type.h

Go to the documentation of this file.

#pragma once
 
#include <string>
 
#include "sdfg/exceptions.h"
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/targets/gpu/gpu_arch.h"
#include "sdfg/targets/gpu/gpu_types.h"
 
namespace sdfg {
namespace gpu {
 
enum class TargetLevel {
    X_GRID,
    Y_GRID,
    Z_GRID,
    X_BLOCK,
    Y_BLOCK,
    Z_BLOCK,
    WARP,
};
 
inline std::string to_string(const TargetLevel& target_level) {
    switch (target_level) {
        case TargetLevel::X_GRID:
            return "X_GRID";
        case TargetLevel::Y_GRID:
            return "Y_GRID";
        case TargetLevel::Z_GRID:
            return "Z_GRID";
        case TargetLevel::X_BLOCK:
            return "X_BLOCK";
        case TargetLevel::Y_BLOCK:
            return "Y_BLOCK";
        case TargetLevel::Z_BLOCK:
            return "Z_BLOCK";
        case TargetLevel::WARP:
            return "WARP";
    }
    throw InvalidSDFGException("Invalid TargetLevel");
}
 
inline TargetLevel target_level_from_string(const std::string& value) {
    if (value == "X_GRID") {
        return TargetLevel::X_GRID;
    } else if (value == "Y_GRID") {
        return TargetLevel::Y_GRID;
    } else if (value == "Z_GRID") {
        return TargetLevel::Z_GRID;
    } else if (value == "X_BLOCK") {
        return TargetLevel::X_BLOCK;
    } else if (value == "Y_BLOCK") {
        return TargetLevel::Y_BLOCK;
    } else if (value == "Z_BLOCK") {
        return TargetLevel::Z_BLOCK;
    } else if (value == "WARP") {
        return TargetLevel::WARP;
    }
    throw InvalidSDFGException("Invalid TargetLevel: " + value);
}

enum class ReduceStrategy { Register, Shared, Global };
 
inline std::string reduce_strategy_to_string(ReduceStrategy strategy) {
    switch (strategy) {
        case ReduceStrategy::Register:
            return "register";
        case ReduceStrategy::Shared:
            return "shared";
        case ReduceStrategy::Global:
            return "global";
    }
    throw InvalidSDFGException("Invalid ReduceStrategy");
}
 
inline ReduceStrategy reduce_strategy_from_string(const std::string& value) {
    if (value == "register") {
        return ReduceStrategy::Register;
    } else if (value == "shared") {
        return ReduceStrategy::Shared;
    } else if (value == "global") {
        return ReduceStrategy::Global;
    }
    throw InvalidSDFGException("Invalid ReduceStrategy: " + value);
}
 
// The mechanism a target level uses by default: warp -> register+shuffle,
// block -> shared tree, grid -> global atomics. Reproduces the historical
// (pre-partial_storage) codegen when the property is unset.
inline ReduceStrategy default_reduce_strategy(TargetLevel target_level) {
    switch (target_level) {
        case TargetLevel::WARP:
            return ReduceStrategy::Register;
        case TargetLevel::X_BLOCK:
        case TargetLevel::Y_BLOCK:
        case TargetLevel::Z_BLOCK:
            return ReduceStrategy::Shared;
        default:
            return ReduceStrategy::Global;
    }
}

 
class ScheduleType_GPU_Offload {
public:
    static void target_level(structured_control_flow::ScheduleType& schedule, const TargetLevel& target_level) {
        schedule.set_property("target_level", gpu::to_string(target_level));
    }

    static TargetLevel target_level(const structured_control_flow::ScheduleType& schedule) {
        return gpu::target_level_from_string(schedule.properties().at("target_level"));
    }
 
    static std::optional<TargetLevel> target_level_if_exists(const structured_control_flow::ScheduleType& schedule) {
        auto it = schedule.properties().find("target_level");
        if (it == schedule.properties().end()) {
            return std::nullopt;
        }
        return gpu::target_level_from_string(it->second);
    }

    static void parallel_size(structured_control_flow::ScheduleType& schedule, const symbolic::Integer parallel_size) {
        serializer::JSONSerializer serializer;
        schedule.set_property("parallel_size", serializer.expression(parallel_size));
    }

    static symbolic::Integer parallel_size(const structured_control_flow::ScheduleType& schedule) {
        if (schedule.properties().find("parallel_size") == schedule.properties().end()) {
            throw InvalidSDFGException("Parallel size not set for schedule type: " + schedule.value());
        }
        std::string expr_str = schedule.properties().at("parallel_size");
        return symbolic::integer(std::stoi(expr_str));
    }

    static symbolic::Integer lanes(const structured_control_flow::ScheduleType& schedule) {
        auto it = schedule.properties().find("lanes");
        if (it == schedule.properties().end()) {
            return symbolic::integer(1);
        }
        return symbolic::integer(std::stoi(it->second));
    }

    static void lanes(structured_control_flow::ScheduleType& schedule, const symbolic::Integer lanes) {
        if (target_level(schedule) != TargetLevel::X_BLOCK) {
            throw InvalidSDFGException("lanes can only be set on an X_BLOCK schedule");
        }
        if (lanes->as_int() < 1) {
            throw InvalidSDFGException("lanes must be positive");
        }
        serializer::JSONSerializer serializer;
        schedule.set_property("lanes", serializer.expression(lanes));
    }

    static symbolic::Integer threads(const structured_control_flow::ScheduleType& schedule) {
        return symbolic::integer(parallel_size(schedule)->as_int() * lanes(schedule)->as_int());
    }

    static bool nested_sync(const structured_control_flow::ScheduleType& schedule) {
        if (schedule.properties().find("nested_sync") == schedule.properties().end()) {
            return false;
        }
        std::string val = schedule.properties().at("nested_sync");
        return val == "true";
    }

    static void nested_sync(structured_control_flow::ScheduleType& schedule, const bool nested_sync) {
        schedule.set_property("nested_sync", nested_sync ? "true" : "false");
    }

    static ReduceStrategy partial_storage(const structured_control_flow::ScheduleType& schedule) {
        auto it = schedule.properties().find("partial_storage");
        if (it == schedule.properties().end()) {
            return default_reduce_strategy(target_level(schedule));
        }
        return reduce_strategy_from_string(it->second);
    }

    static void partial_storage(structured_control_flow::ScheduleType& schedule, const ReduceStrategy strategy) {
        schedule.set_property("partial_storage", reduce_strategy_to_string(strategy));
    }

    static std::string partial_container(const structured_control_flow::ScheduleType& schedule) {
        auto it = schedule.properties().find("partial_container");
        return it == schedule.properties().end() ? std::string() : it->second;
    }

    static void partial_container(structured_control_flow::ScheduleType& schedule, const std::string& container) {
        schedule.set_property("partial_container", container);
    }

    template<typename Derived>
    static structured_control_flow::ScheduleType
    create(const TargetLevel& target_level_, const symbolic::Integer& parallel_size_) {
        auto schedule_type = structured_control_flow::
            ScheduleType(Derived::value(), structured_control_flow::ScheduleTypeCategory::Offloader);
        target_level(schedule_type, target_level_);
        parallel_size(schedule_type, parallel_size_);
        return schedule_type;
    }
 
    static structured_control_flow::ScheduleType
    create(const GpuArch& arch, const TargetLevel& target_level_, const symbolic::Integer& parallel_size_) {
        auto sched_type = arch.create_schedule_type();
        target_level(sched_type, target_level_);
        parallel_size(sched_type, parallel_size_);
        return sched_type;
    }
};

inline TargetLevel gpu_target_level(const structured_control_flow::ScheduleType& schedule) {
    return target_level_from_string(schedule.properties().at("target_level"));
}

inline void gpu_target_level(structured_control_flow::ScheduleType& schedule, const TargetLevel& target_level) {
    schedule.set_property("target_level", to_string(target_level));
}
 
} // namespace gpu
} // namespace sdfg