gpu_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/structured_control_flow/map.h"
#include "sdfg/targets/gpu/gpu_types.h"
 
namespace sdfg {
namespace gpu {

template<typename Derived>
class ScheduleType_GPU_Base {
public:
    static void dimension(structured_control_flow::ScheduleType& schedule, const GPUDimension& dimension) {
        schedule.set_property("dimension", std::to_string(dimension));
    }

    static GPUDimension dimension(const structured_control_flow::ScheduleType& schedule) {
        return static_cast<GPUDimension>(std::stoi(schedule.properties().at("dimension")));
    }

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

    static symbolic::Integer block_size(const structured_control_flow::ScheduleType& schedule) {
        if (schedule.properties().find("block_size") == schedule.properties().end()) {
            if (dimension(schedule) == GPUDimension::X) {
                return Derived::default_block_size_x();
            } else if (dimension(schedule) == GPUDimension::Y) {
                return symbolic::integer(8);
            } else if (dimension(schedule) == GPUDimension::Z) {
                return symbolic::integer(4);
            } else {
                throw InvalidSDFGException("Invalid GPU dimension");
            }
        }
        std::string expr_str = schedule.properties().at("block_size");
        return symbolic::integer(std::stoi(expr_str));
    }

    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 structured_control_flow::ScheduleType create() {
        auto schedule_type = structured_control_flow::
            ScheduleType(Derived::value(), structured_control_flow::ScheduleTypeCategory::Offloader);
        dimension(schedule_type, GPUDimension::X);
        return schedule_type;
    }
};

inline GPUDimension gpu_dimension(const structured_control_flow::ScheduleType& schedule) {
    return static_cast<GPUDimension>(std::stoi(schedule.properties().at("dimension")));
}

inline void gpu_dimension(structured_control_flow::ScheduleType& schedule, const GPUDimension& dimension) {
    schedule.set_property("dimension", std::to_string(dimension));
}

inline symbolic::Integer gpu_block_size(const structured_control_flow::ScheduleType& schedule) {
    if (schedule.properties().find("block_size") == schedule.properties().end()) {
        auto dim = gpu_dimension(schedule);
        if (dim == GPUDimension::X) {
            // Default block size depends on the backend
            if (schedule.value() == "ROCM") {
                return symbolic::integer(64);
            }
            return symbolic::integer(32);
        } else if (dim == GPUDimension::Y) {
            return symbolic::integer(8);
        } else if (dim == GPUDimension::Z) {
            return symbolic::integer(4);
        } else {
            throw InvalidSDFGException("Invalid GPU dimension");
        }
    }
    std::string expr_str = schedule.properties().at("block_size");
    return symbolic::integer(std::stoi(expr_str));
}

inline void gpu_block_size(structured_control_flow::ScheduleType& schedule, const symbolic::Expression block_size) {
    serializer::JSONSerializer serializer;
    schedule.set_property("block_size", serializer.expression(block_size));
}
 
} // namespace gpu
} // namespace sdfg