gpu_map_utils.h

Go to the documentation of this file.

#pragma once
 
#include <unordered_map>
#include <vector>
 
#include "sdfg/analysis/analysis.h"
#include "sdfg/structured_control_flow/map.h"
#include "sdfg/structured_control_flow/sequence.h"
#include "sdfg/structured_control_flow/structured_loop.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/targets/gpu/gpu_offload_schedule_type.h"
#include "sdfg/targets/gpu/gpu_types.h"
 
namespace sdfg {
 
// Forward declarations for explicit template instantiation
namespace cuda {
class ScheduleType_CUDA;
class ScheduleType_CUDA_Offload;
} // namespace cuda
namespace rocm {
class ScheduleType_ROCM;
class ScheduleType_ROCM_Offload;
} // namespace rocm
 
namespace gpu {


template<typename ScheduleT>
symbolic::Expression find_nested_gpu_blocksize(
    structured_control_flow::Map& node, analysis::AnalysisManager& analysis_manager, GPUDimension dimension
);

template<typename ScheduleT>
symbolic::Expression find_nested_gpu_iterations(
    structured_control_flow::Map& node, analysis::AnalysisManager& analysis_manager, GPUDimension dimension
);

template<typename ScheduleT>
bool is_outermost_gpu_map(structured_control_flow::Map& node, analysis::AnalysisManager& analysis_manager);

template<typename ScheduleT>
symbolic::SymbolSet get_gpu_indvars(
    structured_control_flow::Map& node, analysis::AnalysisManager& analysis_manager, GPUDimension dimension
);

bool is_gpu_schedule(const structured_control_flow::ScheduleType& schedule);
 
// Warp/wavefront size of the target, mirroring the offload dispatchers'
// get_warp_size() (cuda::CUDA_WARP_SIZE / rocm::ROCM_WARP_SIZE).
template<typename GPUType>
int64_t gpu_warp_size();
 
template<>
int64_t gpu_warp_size<sdfg::cuda::ScheduleType_CUDA_Offload>();
template<>
int64_t gpu_warp_size<sdfg::rocm::ScheduleType_ROCM_Offload>();
 
// Runtime variant dispatching on an offload schedule's target (defaults to the
// CUDA warp size for any non-ROCm schedule).
int64_t gpu_warp_size(const structured_control_flow::ScheduleType& schedule);

template<typename ScheduleT>
std::vector<structured_control_flow::Map*>
get_gpu_maps(structured_control_flow::Map& node, analysis::AnalysisManager& analysis_manager, GPUDimension dimension);

bool nested_parallelization_is_unsafe(
    structured_control_flow::StructuredLoop& loop, analysis::AnalysisManager& analysis_manager
);
 
struct FoldBarrier {
    structured_control_flow::Sequence* sequence;
    structured_control_flow::ControlFlowNode* before;
};
 
struct NestedFoldPlan {
    bool unsafe = false;
    std::vector<FoldBarrier> barriers;
};
 
NestedFoldPlan analyze_nested_fold(
    structured_control_flow::StructuredLoop& loop, TargetLevel target_level, analysis::AnalysisManager& analysis_manager
);
 
symbolic::Expression get_target_level_dim(TargetLevel target_level, int warp_size);
 
symbolic::Expression get_target_level_idx(TargetLevel target_level);
 
// Coarse GPU cooperation level of a target axis.
bool is_grid_level(TargetLevel target_level);
bool is_block_level(TargetLevel target_level);
bool is_warp_level(TargetLevel target_level);
 
// Induction variables of every offloaded loop (Map or Reduce) at @p target_level
// in @p node's subtree (including @p node itself).
symbolic::SymbolSet target_level_indvars(
    structured_control_flow::StructuredLoop& node, analysis::AnalysisManager& analysis_manager, TargetLevel target_level
);
 
// Per-level launch schedule for @p node's subtree: for each occupied target level,
// the offloaded schedule with the largest thread count (parallel_size * lanes), so a
// single launch dimension covers every sibling at that level.
void get_nested_schedule_types(
    structured_control_flow::StructuredLoop& node,
    analysis::AnalysisManager& analysis_manager,
    std::unordered_map<TargetLevel, structured_control_flow::ScheduleType>& output
);
 
// Per-level launch map for @p node's subtree: for each occupied target level, the
// offloaded loop with the largest parallel_size (matching get_nested_schedule_types).
// Exposes the loop node so callers can read its num_iterations for the launch size.
void get_nested_level_maps(
    structured_control_flow::StructuredLoop& node,
    analysis::AnalysisManager& analysis_manager,
    std::unordered_map<TargetLevel, structured_control_flow::StructuredLoop*>& output
);
 
bool nested_warp_dim(structured_control_flow::StructuredLoop& loop, analysis::AnalysisManager& analysis_manager);
 
// Throws unless every lanes > 1 offload loop in @p node's subtree is an X_BLOCK Map with
// lanes == the warp size of its ARCH (or @p warp_size without one) and no nested WARP loop.
void validate_wave_maps(
    structured_control_flow::StructuredLoop& node, analysis::AnalysisManager& analysis_manager, int64_t warp_size
);
 
structured_control_flow::StructuredLoop* find_x_block_owning_warp_level(
    structured_control_flow::StructuredLoop& node, analysis::AnalysisManager& analysis_manager
);
 
// Counts the depth of the perfectly nested loop chain starting at `loop`.
// A loop is perfectly nested when its body sequence has exactly one child and
// that child is another map or reduce node.
size_t perfectly_nested_depth(structured_control_flow::StructuredLoop* loop);
 
// Extern template declarations to prevent implicit instantiation
extern template symbolic::Expression find_nested_gpu_blocksize<
    sdfg::cuda::ScheduleType_CUDA>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
extern template symbolic::Expression find_nested_gpu_blocksize<
    sdfg::rocm::ScheduleType_ROCM>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
 
extern template symbolic::Expression find_nested_gpu_iterations<
    sdfg::cuda::ScheduleType_CUDA>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
extern template symbolic::Expression find_nested_gpu_iterations<
    sdfg::rocm::ScheduleType_ROCM>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
 
extern template bool is_outermost_gpu_map<
    sdfg::cuda::ScheduleType_CUDA>(structured_control_flow::Map&, analysis::AnalysisManager&);
extern template bool is_outermost_gpu_map<
    sdfg::rocm::ScheduleType_ROCM>(structured_control_flow::Map&, analysis::AnalysisManager&);
 
extern template symbolic::SymbolSet get_gpu_indvars<
    sdfg::cuda::ScheduleType_CUDA>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
extern template symbolic::SymbolSet get_gpu_indvars<
    sdfg::rocm::ScheduleType_ROCM>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
 
extern template std::vector<structured_control_flow::Map*> get_gpu_maps<
    sdfg::cuda::ScheduleType_CUDA>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
extern template std::vector<structured_control_flow::Map*> get_gpu_maps<
    sdfg::rocm::ScheduleType_ROCM>(structured_control_flow::Map&, analysis::AnalysisManager&, GPUDimension);
 
} // namespace gpu
} // namespace sdfg