tensor_layout.h

Go to the documentation of this file.

#pragma once
#include <memory>
#include <nlohmann/json_fwd.hpp>
#include <optional>
 
#include "sdfg/data_flow/memlet.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/types/type.h"
 
namespace sdfg::math::tensor {

class TensorLayout {
private:
    symbolic::MultiExpression shape_;
    symbolic::MultiExpression strides_;
    symbolic::Expression offset_;
 
public:
    TensorLayout(
        const symbolic::MultiExpression& shape,
        const symbolic::MultiExpression& strides = {},
        const symbolic::Expression offset = symbolic::integer(0)
    );
 
    const symbolic::MultiExpression& shape() const {
        return shape_;
    }
    const symbolic::MultiExpression& strides() const {
        return strides_;
    }
 
    const symbolic::Expression& offset() const {
        return offset_;
    }
 
    void serialize_to_json(nlohmann::json& j) const;
 
    std::string toStr() const;
 
    symbolic::Expression total_elements() const;
 
    symbolic::Expression memory_span() const;
 
    symbolic::MultiExpression linear_strides() const;

    symbolic::Expression resolve_element(const symbolic::MultiExpression& indices, bool require_to_element = true) const;

    std::optional<symbolic::MultiExpression> partition_offset_into_dimensions() const;

    symbolic::Expression max_accessed_byte_offset_from_ptr(types::PrimitiveType element_type) const;
 
    static symbolic::MultiExpression linear_strides(const symbolic::MultiExpression& shape);
 
    void collect_symbols(symbolic::SymbolSet& set) const;
 
    void replace_symbols(const symbolic::Expression& old, const symbolic::Expression& new_expr);
 
    void replace_symbols(const symbolic::ExpressionMapping& replacements);

    const symbolic::Expression& get_dim(int i) const {
        return shape_.at(i);
    }

    symbolic::Expression get_dim_innermost(int i) const {
        return shape_.at(shape_.size() - 1 - i);
    }

    symbolic::Expression get_stride(int i) const {
        return strides_.at(i);
    }

    symbolic::Expression get_stride_innermost(int i) const {
        return strides_.at(strides_.size() - 1 - i);
    }
 
    int dims() const {
        return shape_.size();
    }
 
    bool is_scalar() const;
 
    static TensorLayout deserialize_from_json(const nlohmann::json& j);
 
    static bool has_linear_accesses(symbolic::MultiExpression shape, symbolic::MultiExpression strides);
 
    static bool has_linear_accesses_no_padding(
        symbolic::MultiExpression shape, symbolic::MultiExpression strides, symbolic::Expression offset
    );
 
    bool has_linear_accesses() const;
 
    bool has_linear_accesses_no_padding() const;
 
    bool has_transposed_strides_no_padding() const;
 
    bool operator==(const TensorLayout& other) const;
 
    std::unique_ptr<TensorLayout> newaxis(size_t axis) const;
 
    std::unique_ptr<TensorLayout> flip(size_t axis) const;
 
    std::unique_ptr<TensorLayout> unsqueeze(size_t axis) const;
 
    std::unique_ptr<TensorLayout> squeeze(size_t axis) const;
 
    std::unique_ptr<TensorLayout> squeeze() const;
 
    std::unique_ptr<TensorLayout> reshape(const symbolic::MultiExpression& new_shape) const;

    std::unique_ptr<TensorLayout> broadcast(const symbolic::MultiExpression& ref_shape) const;
 
    static std::ostream& emit_symbolic_list(std::ostream& stream, const symbolic::MultiExpression& list);
 
    static types::PrimitiveType get_tensor_indvar_type_for_shape(const std::vector<symbolic::Expression>& shape);
 
    enum TensorLayoutType { LAYOUT_OTHER = 0, LAYOUT_ROW_MAJOR = 1, LAYOUT_COL_MAJOR = 2, LAYOUT_ROW_OR_COL_MAJOR = 3 };
    TensorLayoutType is_2d_col_or_row_major() const;

    bool is_last_dims_expressible_as_row_major() const;

    static TensorLayoutType is_last_dims_col_or_row_major(const symbolic::MultiExpression& strides);
};
 
std::ostream& operator<<(std::ostream& stream, const TensorLayout& layout);
 
 
} // namespace sdfg::math::tensor