tile_copy_node.h

Go to the documentation of this file.

#pragma once
 
#include "sdfg/data_flow/library_node.h"
 
#include <functional>
#include <string>
 
#include "sdfg/codegen/dispatchers/block_dispatcher.h"
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/tiles/tiled_copy.h"
 
namespace sdfg {
namespace tiles {
 
inline data_flow::LibraryNodeCode LibraryNodeType_TileCopy{"tile_copy"};

class TileCopyNode : public data_flow::LibraryNode {
    TiledCopy plan_;
    CopyDirection direction_;
    size_t bytes_; 
    TileGuard guard_;
    std::vector<int> coop_axes_;
    symbolic::Expression coop_threads_;
    size_t coop_lanes_ = 1;
    CopyPhase phase_ = CopyPhase::Full;
 
public:
    TileCopyNode(
        size_t element_id,
        const DebugInfo& debug_info,
        const graph::Vertex vertex,
        data_flow::DataFlowGraph& parent,
        const data_flow::ImplementationType& implementation_type,
        TiledCopy plan,
        CopyDirection direction,
        size_t bytes,
        TileGuard guard = {},
        std::vector<int> coop_axes = {},
        symbolic::Expression coop_threads = {},
        size_t coop_lanes = 1
    );
 
    const TiledCopy& plan() const {
        return plan_;
    }
    void set_plan(TiledCopy plan) {
        plan_ = std::move(plan);
    }
 
    CopyDirection direction() const {
        return direction_;
    }
    CopyAtom atom() const {
        return plan_.atom;
    }
    void set_atom(CopyAtom atom) {
        plan_.atom = atom;
    }
 
    size_t bytes() const {
        return bytes_;
    }
    void set_bytes(size_t bytes) {
        bytes_ = bytes;
    }
 
    const TileGuard& guard() const {
        return guard_;
    }

    bool normalize_guard(const symbolic::SymbolSet& parameters, const symbolic::Assumptions& assumptions) {
        return guard_.discharge(parameters, assumptions);
    }
 
    const std::vector<int>& coop_axes() const {
        return coop_axes_;
    }

    const symbolic::Expression& coop_threads() const {
        return coop_threads_;
    }
 
    size_t coop_lanes() const {
        return coop_lanes_;
    }
 
    CopyPhase phase() const {
        return phase_;
    }
    void set_phase(CopyPhase phase) {
        phase_ = phase;
    }

    int element_input() const {
        return phase_ == CopyPhase::LoadRegs ? 1 : 0;
    }
 
 
    void validate(const Function& function) const override;
 
    symbolic::SymbolSet symbols() const override;
 
    std::unique_ptr<DataFlowNode>
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const override;
 
    void replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) override;
 
    void replace(const symbolic::ExpressionMapping& replacements) override;
 
    std::string toStr() const override;

    data_flow::PointerAccessType pointer_access_type(int input_idx) const override;
};
 
class TileCopyNodeSerializer : public serializer::LibraryNodeSerializer {
public:
    nlohmann::json serialize(const sdfg::data_flow::LibraryNode& library_node) override;
    data_flow::LibraryNode& deserialize(
        const nlohmann::json& j,
        sdfg::builder::StructuredSDFGBuilder& builder,
        sdfg::structured_control_flow::Block& parent
    ) override;
};

bool provably_multiple_of(symbolic::Expression e, long long m, const TileCopyNode& node);

bool guard_axis_aligned(const TileCopyNode& node, size_t axis, long long factor);

size_t widest_legal_transfer(const TileCopyNode& node, CopyAtom atom, const std::vector<size_t>& widths);

std::string emit_tile_copy_stmt(
    codegen::LanguageExtension& language_extension,
    const TileCopyNode& node,
    const std::string& dst_expr,
    const std::string& src_expr,
    types::PrimitiveType elem,
    const symbolic::Expression& index
);

using GuardedLoadFn = std::function<std::string(
    const std::string& type,
    const std::string& var,
    const std::string& base,
    const std::string& index,
    size_t elem_bytes,
    size_t bytes
)>;

std::string emit_predicated_guarded_load(
    const std::string& type, const std::string& var, const std::string& base, const std::string& index
);

void emit_cooperative_copy_loop(
    codegen::LanguageExtension& language_extension,
    codegen::PrettyPrinter& stream,
    const TileCopyNode& node,
    const std::string& dst_expr,
    const std::string& src_expr,
    types::PrimitiveType elem,
    const std::function<std::string(const std::string& dst_addr, const std::string& src_addr, size_t bytes)>&
        async_stmt = {},
    const GuardedLoadFn& guarded_load = {}
);

class TileCopyNodeDispatcher : public codegen::LibraryNodeDispatcher {
public:
    TileCopyNodeDispatcher(
        codegen::LanguageExtension& language_extension,
        const Function& function,
        const data_flow::DataFlowGraph& data_flow_graph,
        const TileCopyNode& node
    );
 
    void dispatch_code_with_edges(
        codegen::CodegenOutput& out,
        std::vector<codegen::DispatchInput>& inputs,
        std::vector<codegen::DispatchOutput>& outputs
    ) override;
};
 
} // namespace tiles
} // namespace sdfg