local_storage.h

Go to the documentation of this file.

#pragma once
 
#include <optional>
#include <string>
#include <unordered_set>
#include <vector>
 
#include "sdfg/analysis/analysis.h"
#include "sdfg/analysis/memory_layout_analysis.h"
#include "sdfg/data_flow/access_node.h"
#include "sdfg/data_flow/data_flow_graph.h"
#include "sdfg/data_flow/memlet.h"
#include "sdfg/structured_control_flow/control_flow_node.h"
#include "sdfg/structured_control_flow/reduce.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"
#include "sdfg/tiles/analysis/reduction_analysis.h"
#include "sdfg/tiles/analysis/tile_analysis.h"
#include "sdfg/tiles/buffer_layout.h"
#include "sdfg/tiles/locality.h"
#include "sdfg/tiles/tiled_copy.h"
#include "sdfg/transformations/transformation.h"
#include "sdfg/types/type.h"
 
namespace sdfg {
namespace transformations {

class LocalStorage : public Transformation {
public:
    struct TileInfo {
        std::vector<symbolic::Expression> dimensions;
        std::vector<symbolic::Expression> bases;
        std::vector<symbolic::Expression> maxes;
        std::vector<symbolic::Expression> strides;
        symbolic::Expression offset = symbolic::integer(0);

        std::vector<size_t> varying_dims() const;
        std::vector<symbolic::Expression> varying_sizes() const;
        tiles::Layout source_layout() const;
        std::vector<symbolic::Expression> original_subset(const std::vector<symbolic::Expression>& tile_indices) const;
        std::vector<symbolic::Expression> local_index(const std::vector<symbolic::Expression>& access_subset) const;
    };

    static bool has_side_effect(structured_control_flow::StructuredLoop& loop);
 
private:
    structured_control_flow::StructuredLoop& loop_;
    const data_flow::AccessNode& access_node_;
    std::string container_;
    std::string local_name_; 
    types::StorageType storage_type_; 
    TileInfo tile_info_; 
    tiles::LocalityPlan plan_; 
    bool swizzle_layout_ = false; 
    bool lane_contiguous_ = false; 
    bool transpose_layout_ = false; 
    std::unordered_set<const data_flow::Memlet*> group_memlets_; 
    std::vector<structured_control_flow::Reduce*> reduce_retargets_; 
    std::vector<structured_control_flow::Reduce*> grid_reduce_owners_; 
    bool atomic_merge_ = false; 
    bool container_read_ = false; 
    bool container_written_ = false; 

    bool prepare(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager);

    bool reduction_buffers_supported(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager);

    void apply_prepared(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager);

    bool needs_copy_in() const {
        return container_read_;
    }

    bool needs_copy_out() const {
        return container_written_;
    }

    size_t max_tile_elements() const {
        return 1u << 16;
    }

    symbolic::Condition boundary_guard(
        const data_flow::Subset& tile_indices,
        const symbolic::SymbolSet& params = {},
        const symbolic::Assumptions& assums = {}
    ) const;

    void emit_private_copy(
        builder::StructuredSDFGBuilder& builder,
        analysis::AnalysisManager& analysis_manager,
        structured_control_flow::Sequence& parent,
        const tiles::PackedBuffer& buffer,
        const types::IType& buffer_type,
        const types::IType& pointer_type,
        bool writeback
    );

    struct BuiltCopy {
        tiles::TiledCopy plan;
        std::vector<int> coop_axes;
        tiles::TileGuard guard;
        symbolic::Expression coop_threads; 
        size_t coop_lanes = 1; 
    };

    BuiltCopy build_tiled_copy(
        analysis::AnalysisManager& analysis_manager,
        structured_control_flow::Sequence& guard_scope,
        const tiles::PackedBuffer& buffer,
        const types::IType& pointer_type,
        const std::vector<symbolic::Expression>& slot_sizes,
        const std::vector<symbolic::Expression>& slot_indices,
        const std::vector<symbolic::Expression>& slot_indvars,
        const std::vector<symbolic::Expression>& slot_inits,
        const std::vector<symbolic::Expression>& slot_strides
    );

    void emit_copy_node(
        builder::StructuredSDFGBuilder& builder,
        structured_control_flow::Sequence& scope,
        structured_control_flow::ControlFlowNode& anchor,
        bool after,
        const BuiltCopy& copy,
        const data_flow::ImplementationType& impl,
        const types::IType& pointer_type,
        tiles::CopyDirection direction,
        bool leading_barrier,
        bool trailing_barrier
    );

    tiles::TileGuard tile_boundary_guard(
        analysis::AnalysisManager& analysis_manager,
        structured_control_flow::Sequence& guard_scope,
        const std::vector<symbolic::Expression>& vsizes
    );

    void rewrite_body(
        builder::StructuredSDFGBuilder& builder,
        analysis::AnalysisManager& analysis_manager,
        const tiles::PackedBuffer& buffer,
        const types::IType& buffer_type,
        const std::vector<symbolic::Expression>& slot_indices
    );

    std::optional<math::tensor::TensorLayout> operand_view(
        const analysis::MemoryTile& access,
        const tiles::PackedBuffer& buffer,
        const std::vector<symbolic::Expression>& slot_indices
    ) const;

    void choose_conflict_policy(
        analysis::AnalysisManager& analysis_manager, tiles::PackedBuffer& buffer, types::PrimitiveType element_type
    );
 
public:
    LocalStorage(
        structured_control_flow::StructuredLoop& loop,
        const data_flow::AccessNode& access_node,
        bool swizzle_layout = false,
        bool lane_contiguous = false,
        bool transpose_layout = false
    )
        : loop_(loop), access_node_(access_node), container_(access_node.data()),
          storage_type_(types::StorageType::CPU_Stack()), swizzle_layout_(swizzle_layout),
          lane_contiguous_(lane_contiguous), transpose_layout_(transpose_layout) {
    }
 
    std::string name() const override {
        return "LocalStorage";
    }

    bool can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) override;

    void apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) override;

    void to_json(nlohmann::json& j) const override;

    static LocalStorage from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& j);

    const std::string& local_container() const {
        return local_name_;
    }

    const TileInfo& tile_info() const {
        return tile_info_;
    }

    const types::StorageType& storage_type() const {
        return storage_type_;
    }
};
 
} // namespace transformations
} // namespace sdfg