stream_k.h

Go to the documentation of this file.

#pragma once
 
#include <map>
 
#include "sdfg/structured_control_flow/reduce.h"
#include "sdfg/transformations/transformation.h"
 
namespace sdfg::gpu {
class GpuMmaMatmulNode;
}
 
namespace sdfg {
namespace transformations {

class StreamK : public Transformation {
    // Outermost GPU-offloaded parallel output-tile loop (the band converted to a
    // 1-D persistent worker grid).
    structured_control_flow::StructuredLoop& grid_loop_;
 
    // Fixed persistent grid size (absolute block count).
    size_t num_blocks_;
 
    void apply_fragment_split(
        builder::StructuredSDFGBuilder& builder,
        analysis::AnalysisManager& analysis_manager,
        structured_control_flow::Reduce& reduce
    );
 
    void build_fragment_merge(
        builder::StructuredSDFGBuilder& builder,
        structured_control_flow::Reduce& reduce,
        const std::vector<structured_control_flow::ReductionInfo>& reductions,
        const std::map<std::string, const gpu::GpuMmaMatmulNode*>& accumulators,
        const symbolic::Symbol& bid,
        const symbolic::Expression& tile,
        const symbolic::Expression& wave,
        int64_t num_waves,
        int64_t splits,
        int64_t tiles
    );
 
public:
    explicit StreamK(structured_control_flow::StructuredLoop& grid_loop, size_t num_blocks = 336);
 
    virtual std::string name() const override;
 
    virtual bool
    can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) override;
 
    virtual void apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) override;
 
    virtual void to_json(nlohmann::json& j) const override;
 
    static StreamK from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& j);
};
 
} // namespace transformations
} // namespace sdfg