gpu_tile_target.h

Go to the documentation of this file.

#pragma once
 
#include <optional>
#include <string>
#include <utility>
 
#include "sdfg/targets/gpu/gpu_map_utils.h"
#include "sdfg/targets/gpu/gpu_offload_schedule_type.h"
#include "sdfg/targets/gpu/gpu_schedule_type.h"
#include "sdfg/tiles/tile_target.h"
 
namespace sdfg {
namespace tiles {

class GPUTileTarget : public TileTarget {
    unsigned lane_width_;
    std::string offload_value_;
    data_flow::ImplementationType implementation_type_;
    size_t min_async_bytes_;
 
    static unsigned spatial_axis_of(gpu::TargetLevel tl) {
        switch (tl) {
            case gpu::TargetLevel::Y_GRID:
            case gpu::TargetLevel::Y_BLOCK:
                return 1;
            case gpu::TargetLevel::Z_GRID:
            case gpu::TargetLevel::Z_BLOCK:
                return 2;
            default:
                return 0;
        }
    }
 
public:
    GPUTileTarget(
        unsigned lane_width,
        std::string offload_value,
        data_flow::ImplementationType implementation_type,
        size_t min_async_bytes = 4
    )
        : lane_width_(lane_width), offload_value_(std::move(offload_value)),
          implementation_type_(std::move(implementation_type)), min_async_bytes_(min_async_bytes) {
    }
 
    std::optional<AxisSchedule> classify(const structured_control_flow::ScheduleType& sched) const override {
        if (sched.value() == offload_value_) {
            gpu::TargetLevel tl = gpu::gpu_target_level(sched);
            Level level = gpu::is_grid_level(tl)    ? Level::Device
                          : gpu::is_block_level(tl) ? Level::Group
                                                    : Level::Subgroup;
            return AxisSchedule(
                level,
                space(level),
                has_scratchpad(),
                spatial_axis_of(tl),
                gpu::ScheduleType_GPU_Offload::parallel_size(sched),
                gpu::ScheduleType_GPU_Offload::nested_sync(sched),
                gpu::ScheduleType_GPU_Offload::lanes(sched)
            );
        }
        // Legacy fused block-thread schedule.
        return AxisSchedule(
            Level::Group, space(Level::Group), has_scratchpad(), /*spatial_axis=*/0, gpu::gpu_block_size(sched)
        );
    }

    Space space(Level level) const override {
        switch (level) {
            case Level::Device:
                return Space::Global; // device-wide cooperation (GPU grid): only global memory
            case Level::Group:
                return Space::Shared;
            case Level::Subgroup:
                return Space::Register;
        }
        return Space::Register;
    }
 
    bool supports_cooperative_staging(const structured_control_flow::ScheduleType& sched) const override {
        // Only the offload schedule carries a separable block-thread copy Map; the
        // legacy fused schedule cannot host cooperative shared-memory staging.
        return sched.value() == offload_value_;
    }
 
    types::StorageType storage_type(Space space) const override {
        switch (space) {
            case Space::Shared:
                return types::StorageType::NV_Shared();
            case Space::Global:
                return types::StorageType::NV_Global();
            case Space::Register:
                return types::StorageType::CPU_Stack();
        }
        return types::StorageType::CPU_Stack();
    }
 
    unsigned lane_width() const override {
        return lane_width_;
    }

    std::optional<tiles::BankModel> bank_model(tiles::Space space) const override {
        if (space == tiles::Space::Shared) {
            return tiles::BankModel{32, 4};
        }
        return std::nullopt;
    }
 
    data_flow::ImplementationType implementation_type() const override {
        return implementation_type_;
    }
 
    size_t min_async_bytes() const override {
        return min_async_bytes_;
    }
};
 
} // namespace tiles
} // namespace sdfg