SDFGTranslator.h

Go to the documentation of this file.

#pragma once
 
#include <list>
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/ScopedHashTable.h>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
 
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/element.h"
#include "sdfg/function.h"
#include "sdfg/structured_control_flow/sequence.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/types/scalar.h"
#include "sdfg/types/tensor.h"
 
namespace mlir {
namespace sdfg {
 
class TensorInfo {
    std::vector<int64_t> shape_;
    std::vector<int64_t> strides_;
    int64_t offset_;
 
public:
    TensorInfo();
    TensorInfo(std::vector<int64_t> shape, std::vector<int64_t> strides, int64_t offset = 0);
 
    const std::vector<int64_t>& shape() const;
    const std::vector<int64_t>& strides() const;
    int64_t offset() const;

    static std::vector<int64_t> compute_strides(const std::vector<int64_t>& shape);

    static TensorInfo from_tensor_type(TensorType type);

    static bool has_basic_strides(ArrayRef<int64_t> shape, ArrayRef<int64_t> strides);

    TensorInfo transpose(ArrayRef<int64_t> permutation) const;

    TensorInfo flip(ArrayRef<int64_t> axes) const;

    bool is_reshape_valid(ArrayRef<int64_t> new_shape) const;

    bool has_basic_strides() const;

    bool has_transposed_strides_last_two_dims() const;

    TensorInfo reshape(ArrayRef<int64_t> new_shape) const;

    std::string shape_str() const;

    std::unique_ptr<::sdfg::types::Tensor> get_sdfg_tensor(const ::sdfg::types::Scalar& element_type) const;

    ::sdfg::math::tensor::TensorLayout get_tensor_layout() const;

    std::string toStr() const;
};
 
class SDFGTranslator {
    bool builder_empty_;
    ::sdfg::builder::StructuredSDFGBuilder builder_;
 
    llvm::ScopedHashTable<Value, std::string> value_map_;
    size_t value_counter_;
 
    std::list<::sdfg::structured_control_flow::Sequence*> insertion_points_;
 
    std::unordered_map<std::string, TensorInfo> tensor_info_map_;
 
    std::unordered_map<::sdfg::structured_control_flow::Sequence*, std::list<std::string>> memory_map_;
 
    std::unordered_map<std::string, std::string> alias_map_;
 
    // Maps a linalg.fill result whose materialization was deferred to the constant scalar value it
    // fills with, so each consumer can regenerate a freshly-filled array instead of copying from a
    // single shared buffer.
    llvm::DenseMap<Value, Value> constant_fill_map_;
 
    std::vector<std::string> output_args_;
 
public:
    SDFGTranslator();
 
    ::sdfg::builder::StructuredSDFGBuilder& builder();
 
    bool builder_empty();
    void builder_empty(bool empty);
 
    llvm::ScopedHashTable<Value, std::string>& value_map();
 
    ::sdfg::structured_control_flow::Sequence& insertion_point();
    void enter_sequence(::sdfg::structured_control_flow::Sequence& sequence);
    void exit_sequence(::sdfg::structured_control_flow::Sequence& sequence);
 
    std::unordered_map<std::string, TensorInfo>& tensor_info_map();
 
    TensorInfo& get_or_create_tensor_info(const std::string& container, const TensorType& type);
 
    std::string get_or_create_container(Value val, bool argument = false);
 
    std::unique_ptr<::sdfg::types::IType> convertType(const Type mlir_type);
 
    std::string convertTypedAttr(const TypedAttr attr);
 
    void add_reference(
        const std::string& src_container,
        const std::string& dst_container,
        const ::sdfg::DebugInfo& deb_info = ::sdfg::DebugInfo()
    );
 
    void handle_malloc(
        std::string container,
        const ::sdfg::symbolic::Expression size,
        const ::sdfg::DebugInfo& deb_info = ::sdfg::DebugInfo()
    );
    void handle_frees(std::string return_container = "", const ::sdfg::DebugInfo& deb_info = ::sdfg::DebugInfo());

    std::string get_or_copy_output_container(
        Value output, const ::sdfg::DebugInfo& deb_info = ::sdfg::DebugInfo(), bool consumer_overwrites_output = false
    );

    void record_constant_fill(Value result, Value value);

    bool is_constant_fill(Value result) const;

    std::string try_inplace_reuse_container(Value input, Value result);
 
    std::string store_in_c_order(
        const std::string& container,
        const TensorInfo& tensor_info,
        const ::sdfg::types::Scalar& element_type,
        const ::sdfg::DebugInfo& deb_info = ::sdfg::DebugInfo()
    );

    void set_output_args(const std::vector<std::string>& output_args);

    const std::vector<std::string>& output_args() const;

    void copy_to_output(
        const std::string& src_container,
        const TensorInfo& tensor_info,
        const ::sdfg::types::Scalar& element_type,
        const std::string& output_container,
        const ::sdfg::DebugInfo& deb_info = ::sdfg::DebugInfo()
    );

    void copy_scalar_to_output(
        const std::string& src_container,
        const std::string& output_container,
        const ::sdfg::DebugInfo& deb_info = ::sdfg::DebugInfo()
    );

    ::sdfg::DebugInfo get_debug_info(llvm::StringLiteral operation_name, Location loc);
};
 
LogicalResult translateOp(SDFGTranslator& translator, Operation* op);
 
LogicalResult emitJSON(SDFGTranslator& translator, raw_ostream& os);
 
} // namespace sdfg
} // namespace mlir