reduction_buffer_analysis.h

Go to the documentation of this file.

#pragma once
 
#include <optional>
#include <string>
#include <vector>
 
#include "sdfg/analysis/analysis.h"
#include "sdfg/structured_control_flow/reduce.h"
#include "sdfg/targets/gpu/gpu_reduce_layout.h"
#include "sdfg/transformations/loop_header.h"
 
namespace sdfg {
namespace tiles {

enum class ReductionBufferStatus {
    Exact, 
    ConservativeBound, 
    Unsupported 
};

struct ReductionScheduleProposal {
    structured_control_flow::StructuredLoop& loop;
    const structured_control_flow::ScheduleType& schedule;
};

struct ReductionLoopDomain {
    symbolic::Symbol indvar;
    symbolic::Expression init;
    symbolic::Expression count;
    symbolic::Integer stride;
 
    static ReductionLoopDomain from_header(const symbolic::Symbol& indvar, const transformations::LoopHeader& header);
};

struct ReductionBufferInfo {
    ReductionBufferStatus status = ReductionBufferStatus::Unsupported;
    std::string diagnostic;
    std::optional<gpu::ReductionLayout> layout;
    symbolic::Expression accumulator_index = SymEngine::null;
    std::optional<types::PrimitiveType> primitive;
    std::optional<int64_t> element_bytes;
    std::optional<int64_t> private_bytes;
    std::optional<int64_t> shared_bytes;
    std::optional<size_t> shared_owner;
    symbolic::Expression linear_thread_index = SymEngine::null;
    std::string private_buffer;
    std::string shared_buffer;
    bool materialized = false;
    bool multi_output = false;
};

struct ReductionKernelInfo {
    ReductionBufferStatus status = ReductionBufferStatus::Unsupported;
    std::optional<int64_t> shared_bytes;
    std::string diagnostic;
};

class ReductionBufferAnalysis : public analysis::Analysis {
    analysis::AnalysisManager* analysis_manager_ = nullptr;

    ReductionBufferInfo compute(structured_control_flow::Reduce& reduction, const std::string& container) const;
 
protected:
    void run(analysis::AnalysisManager& analysis_manager) override;
 
public:
    explicit ReductionBufferAnalysis(StructuredSDFG& sdfg);

    std::string name() const override;

    ReductionBufferInfo buffer(structured_control_flow::Reduce& reduction, const std::string& container) const;

    ReductionBufferInfo require(structured_control_flow::Reduce& reduction, const std::string& container) const;

    ReductionBufferInfo estimate_schedule(
        structured_control_flow::Reduce& reduction,
        const std::string& container,
        ReductionBufferInfo footprint,
        const ReductionScheduleProposal& proposal
    ) const;

    ReductionBufferInfo estimate_interchange(
        structured_control_flow::Reduce& reduction,
        const std::string& container,
        ReductionBufferInfo footprint,
        const transformations::LoopSwap& proposal
    ) const;

    ReductionBufferInfo estimate_geometry(
        structured_control_flow::Reduce& reduction,
        const std::string& container,
        ReductionBufferInfo footprint,
        const std::vector<ReductionLoopDomain>& inner_domains
    ) const;

    bool supports_interchange(const transformations::LoopSwap& proposal) const;

    bool supports_schedule(
        structured_control_flow::StructuredLoop& loop, const structured_control_flow::ScheduleType& schedule
    ) const;
    std::vector<structured_control_flow::Reduce*> affected_reductions(structured_control_flow::StructuredLoop& loop) const;

    ReductionKernelInfo kernel(const structured_control_flow::StructuredLoop& root) const;

    bool is_partial_buffer(const std::string& container) const;
};
 
} // namespace tiles
} // namespace sdfg