daisytuner / sdfglib / 15656007340

nlohmann::json JSONSerializer::serialize(std::unique_ptr<sdfg::StructuredSDFG>& sdfg) {

    nlohmann::json j;

    j["name"] = sdfg->name();

    j["type"] = std::string(sdfg->type().value());

    j["structures"] = nlohmann::json::array();

    for (const auto& structure_name : sdfg->structures()) {

        const auto& structure = sdfg->structure(structure_name);

        nlohmann::json structure_json;

        structure_definition_to_json(structure_json, structure);

        j["structures"].push_back(structure_json);

    }

    j["containers"] = nlohmann::json::object();

    for (const auto& container : sdfg->containers()) {

        nlohmann::json desc;

        type_to_json(desc, sdfg->type(container));

        j["containers"][container] = desc;

    }

    j["arguments"] = nlohmann::json::array();

    for (const auto& argument : sdfg->arguments()) {

        j["arguments"].push_back(argument);

    j["externals"] = nlohmann::json::array();

    for (const auto& external : sdfg->externals()) {

        j["externals"].push_back(external);

    j["metadata"] = nlohmann::json::object();

    for (const auto& entry : sdfg->metadata()) {

        j["metadata"][entry.first] = entry.second;

    nlohmann::json root_json;

    sequence_to_json(root_json, sdfg->root());

    j["root"] = root_json;

    return j;

}

void JSONSerializer::dataflow_to_json(nlohmann::json& j, const data_flow::DataFlowGraph& dataflow) {

    j["type"] = "dataflow";

    j["nodes"] = nlohmann::json::array();

    j["edges"] = nlohmann::json::array();

    for (auto& node : dataflow.nodes()) {

        nlohmann::json node_json;

        node_json["element_id"] = node.element_id();

        node_json["debug_info"] = nlohmann::json::object();

        debug_info_to_json(node_json["debug_info"], node.debug_info());

        if (auto tasklet = dynamic_cast<const data_flow::Tasklet*>(&node)) {

            node_json["type"] = "tasklet";

            node_json["code"] = tasklet->code();

            node_json["inputs"] = nlohmann::json::array();

            for (auto& input : tasklet->inputs()) {

                nlohmann::json input_json;

                nlohmann::json type_json;

                type_to_json(type_json, input.second);

                input_json["type"] = type_json;

                input_json["name"] = input.first;

                node_json["inputs"].push_back(input_json);

            }

            node_json["output"] = nlohmann::json::object();

            node_json["output"]["name"] = tasklet->output().first;

            nlohmann::json type_json;

            type_to_json(type_json, tasklet->output().second);

            node_json["output"]["type"] = type_json;

        } else if (auto lib_node = dynamic_cast<const data_flow::LibraryNode*>(&node)) {

                LibraryNodeSerializerRegistry::instance().get_library_node_serializer(

                    lib_node->code().value());

            if (serializer_fn == nullptr) {

                throw std::runtime_error("Unknown library node code: " +

                                         std::string(lib_node->code().value()));

            auto serializer = serializer_fn();

            auto lib_node_json = serializer->serialize(*lib_node);

            node_json.merge_patch(lib_node_json);

        } else if (auto code_node = dynamic_cast<const data_flow::AccessNode*>(&node)) {

            node_json["type"] = "access_node";

            node_json["data"] = code_node->data();

        } else {

            throw std::runtime_error("Unknown node type");

        j["nodes"].push_back(node_json);

    }

    for (auto& edge : dataflow.edges()) {

        nlohmann::json edge_json;

        edge_json["element_id"] = edge.element_id();

        edge_json["debug_info"] = nlohmann::json::object();

        debug_info_to_json(edge_json["debug_info"], edge.debug_info());

        edge_json["src"] = edge.src().element_id();

        edge_json["dst"] = edge.dst().element_id();

        edge_json["src_conn"] = edge.src_conn();

        edge_json["dst_conn"] = edge.dst_conn();

        edge_json["subset"] = nlohmann::json::array();

        for (auto& subset : edge.subset()) {

            edge_json["subset"].push_back(expression(subset));

        j["edges"].push_back(edge_json);

    }

}

void JSONSerializer::block_to_json(nlohmann::json& j, const structured_control_flow::Block& block) {

    j["type"] = "block";

    j["element_id"] = block.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], block.debug_info());

    nlohmann::json dataflow_json;

    dataflow_to_json(dataflow_json, block.dataflow());

    j["dataflow"] = dataflow_json;

}

void JSONSerializer::for_to_json(nlohmann::json& j, const structured_control_flow::For& for_node) {

    j["type"] = "for";

    j["element_id"] = for_node.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], for_node.debug_info());

    j["indvar"] = expression(for_node.indvar());

    j["init"] = expression(for_node.init());

    j["condition"] = expression(for_node.condition());

    j["update"] = expression(for_node.update());

    nlohmann::json body_json;

    sequence_to_json(body_json, for_node.root());

    j["root"] = body_json;

}

void JSONSerializer::if_else_to_json(nlohmann::json& j,

    j["type"] = "if_else";

    j["element_id"] = if_else_node.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], if_else_node.debug_info());

    j["branches"] = nlohmann::json::array();

    for (size_t i = 0; i < if_else_node.size(); i++) {

        nlohmann::json branch_json;

        branch_json["condition"] = expression(if_else_node.at(i).second);

        nlohmann::json body_json;

        sequence_to_json(body_json, if_else_node.at(i).first);

        branch_json["root"] = body_json;

        j["branches"].push_back(branch_json);

    }

}

void JSONSerializer::while_node_to_json(nlohmann::json& j,

    j["type"] = "while";

    j["element_id"] = while_node.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], while_node.debug_info());

    nlohmann::json body_json;

    sequence_to_json(body_json, while_node.root());

    j["root"] = body_json;

}

void JSONSerializer::break_node_to_json(nlohmann::json& j,

    j["type"] = "break";

    j["element_id"] = break_node.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], break_node.debug_info());

}

void JSONSerializer::continue_node_to_json(nlohmann::json& j,

    j["type"] = "continue";

    j["element_id"] = continue_node.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], continue_node.debug_info());

}

void JSONSerializer::map_to_json(nlohmann::json& j, const structured_control_flow::Map& map_node) {

    j["type"] = "map";

    j["element_id"] = map_node.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], map_node.debug_info());

    j["indvar"] = expression(map_node.indvar());

    j["num_iterations"] = expression(map_node.num_iterations());

    j["schedule_type"] = std::string(map_node.schedule_type().value());

    nlohmann::json body_json;

    sequence_to_json(body_json, map_node.root());

    j["root"] = body_json;

}

void JSONSerializer::return_node_to_json(nlohmann::json& j,

    j["type"] = "return";

    j["element_id"] = return_node.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], return_node.debug_info());

}

void JSONSerializer::sequence_to_json(nlohmann::json& j,

    j["type"] = "sequence";

    j["element_id"] = sequence.element_id();

    j["debug_info"] = nlohmann::json::object();

    debug_info_to_json(j["debug_info"], sequence.debug_info());

    j["children"] = nlohmann::json::array();

    j["transitions"] = nlohmann::json::array();

    for (size_t i = 0; i < sequence.size(); i++) {

        nlohmann::json child_json;

        auto& child = sequence.at(i).first;

        auto& transition = sequence.at(i).second;

        if (auto block = dynamic_cast<const structured_control_flow::Block*>(&child)) {

            block_to_json(child_json, *block);

        } else if (auto for_node = dynamic_cast<const structured_control_flow::For*>(&child)) {

            for_to_json(child_json, *for_node);

        } else if (auto sequence_node =

                       dynamic_cast<const structured_control_flow::Sequence*>(&child)) {

            sequence_to_json(child_json, *sequence_node);

        } else if (auto condition_node =

                       dynamic_cast<const structured_control_flow::IfElse*>(&child)) {

            if_else_to_json(child_json, *condition_node);

        } else if (auto while_node = dynamic_cast<const structured_control_flow::While*>(&child)) {

            while_node_to_json(child_json, *while_node);

        } else if (auto return_node =

                       dynamic_cast<const structured_control_flow::Return*>(&child)) {

            return_node_to_json(child_json, *return_node);

        } else if (auto break_node = dynamic_cast<const structured_control_flow::Break*>(&child)) {

            break_node_to_json(child_json, *break_node);

        } else if (auto continue_node =

                       dynamic_cast<const structured_control_flow::Continue*>(&child)) {

            continue_node_to_json(child_json, *continue_node);

        } else if (auto map_node = dynamic_cast<const structured_control_flow::Map*>(&child)) {

            map_to_json(child_json, *map_node);

        } else {

            throw std::runtime_error("Unknown child type");

        j["children"].push_back(child_json);

        nlohmann::json transition_json;

        transition_json["type"] = "transition";

        transition_json["element_id"] = transition.element_id();

        transition_json["debug_info"] = nlohmann::json::object();

        debug_info_to_json(transition_json["debug_info"], transition.debug_info());

        transition_json["assignments"] = nlohmann::json::array();

        for (const auto& assignment : transition.assignments()) {

            nlohmann::json assignment_json;

            assignment_json["symbol"] = expression(assignment.first);

            assignment_json["expression"] = expression(assignment.second);

            transition_json["assignments"].push_back(assignment_json);

        }

        j["transitions"].push_back(transition_json);

    }

}

void JSONSerializer::type_to_json(nlohmann::json& j, const types::IType& type) {

    if (auto scalar_type = dynamic_cast<const types::Scalar*>(&type)) {

        j["type"] = "scalar";

        j["primitive_type"] = scalar_type->primitive_type();

        j["storage_type"] = std::string(scalar_type->storage_type().value());

        j["initializer"] = scalar_type->initializer();

        j["alignment"] = scalar_type->alignment();

    } else if (auto array_type = dynamic_cast<const types::Array*>(&type)) {

        j["type"] = "array";

        nlohmann::json element_type_json;

        type_to_json(element_type_json, array_type->element_type());

        j["element_type"] = element_type_json;

        j["num_elements"] = expression(array_type->num_elements());

        j["storage_type"] = std::string(array_type->storage_type().value());

        j["initializer"] = array_type->initializer();

        j["alignment"] = array_type->alignment();

    } else if (auto pointer_type = dynamic_cast<const types::Pointer*>(&type)) {

        j["type"] = "pointer";

        nlohmann::json pointee_type_json;

        type_to_json(pointee_type_json, pointer_type->pointee_type());

        j["pointee_type"] = pointee_type_json;

        j["storage_type"] = std::string(pointer_type->storage_type().value());

        j["initializer"] = pointer_type->initializer();

        j["alignment"] = pointer_type->alignment();

    } else if (auto structure_type = dynamic_cast<const types::Structure*>(&type)) {

        j["type"] = "structure";

        j["name"] = structure_type->name();

        j["storage_type"] = std::string(structure_type->storage_type().value());

        j["initializer"] = structure_type->initializer();

        j["alignment"] = structure_type->alignment();

    } else if (auto function_type = dynamic_cast<const types::Function*>(&type)) {

        j["type"] = "function";

        nlohmann::json return_type_json;

        type_to_json(return_type_json, function_type->return_type());

        j["return_type"] = return_type_json;

        j["params"] = nlohmann::json::array();

        for (size_t i = 0; i < function_type->num_params(); i++) {

            nlohmann::json param_json;

            type_to_json(param_json, function_type->param_type(symbolic::integer(i)));

            j["params"].push_back(param_json);

        }

        j["is_var_arg"] = function_type->is_var_arg();

        j["storage_type"] = std::string(function_type->storage_type().value());

        j["initializer"] = function_type->initializer();

        j["alignment"] = function_type->alignment();

    } else {

        throw std::runtime_error("Unknown type");

}

void JSONSerializer::structure_definition_to_json(nlohmann::json& j,

    j["name"] = definition.name();

    j["members"] = nlohmann::json::array();

    for (size_t i = 0; i < definition.num_members(); i++) {

        nlohmann::json member_json;

        type_to_json(member_json, definition.member_type(symbolic::integer(i)));

        j["members"].push_back(member_json);

    }

    j["is_packed"] = definition.is_packed();

}

void JSONSerializer::debug_info_to_json(nlohmann::json& j, const DebugInfo& debug_info) {

    j["has"] = debug_info.has();

    j["filename"] = debug_info.filename();

    j["start_line"] = debug_info.start_line();

    j["start_column"] = debug_info.start_column();

    j["end_line"] = debug_info.end_line();

    j["end_column"] = debug_info.end_column();

}

std::unique_ptr<StructuredSDFG> JSONSerializer::deserialize(nlohmann::json& j) {

    assert(j.contains("name"));

    assert(j["name"].is_string());

    assert(j.contains("type"));

    assert(j["type"].is_string());

    FunctionType type{"Unknown"};

    if (j["type"] == FunctionType_CPU.value()) {

        type = FunctionType_CPU;

    } else if (j["type"] == FunctionType_NV_GLOBAL.value()) {

        type = FunctionType_NV_GLOBAL;

    } else {

        throw std::runtime_error("Unknown function type");

    builder::StructuredSDFGBuilder builder(j["name"], type);

    assert(j.contains("structures"));

    assert(j["structures"].is_array());

    for (const auto& structure : j["structures"]) {

        assert(structure.contains("name"));

        assert(structure["name"].is_string());

        json_to_structure_definition(structure, builder);

    nlohmann::json& containers = j["containers"];

    for (const auto& name : j["externals"]) {

        auto& type_desc = containers.at(name.get<std::string>());

        auto type = json_to_type(type_desc);

        builder.add_container(name, *type, false, true);

    }

    for (const auto& name : j["arguments"]) {

        auto& type_desc = containers.at(name.get<std::string>());

        auto type = json_to_type(type_desc);

        builder.add_container(name, *type, true, false);

    }

    for (const auto& entry : containers.items()) {

        if (builder.subject().is_argument(entry.key())) {

            continue;

        if (builder.subject().is_external(entry.key())) {

            continue;

        auto type = json_to_type(entry.value());

        builder.add_container(entry.key(), *type, false, false);

    }

    assert(j.contains("root"));

    auto& root = builder.subject().root();

    json_to_sequence(j["root"], builder, root);

    assert(j.contains("metadata"));

    assert(j["metadata"].is_object());

    for (const auto& entry : j["metadata"].items()) {

        builder.subject().add_metadata(entry.key(), entry.value());

    return builder.move();

}

void JSONSerializer::json_to_structure_definition(const nlohmann::json& j,

    assert(j.contains("name"));

    assert(j["name"].is_string());

    assert(j.contains("members"));

    assert(j["members"].is_array());

    assert(j.contains("is_packed"));

    assert(j["is_packed"].is_boolean());

    auto is_packed = j["is_packed"];

    auto& definition = builder.add_structure(j["name"], is_packed);

    for (const auto& member : j["members"]) {

        nlohmann::json member_json;

        auto member_type = json_to_type(member);

        definition.add_member(*member_type);

    }

}

std::vector<std::pair<std::string, types::Scalar>> JSONSerializer::json_to_arguments(

    std::vector<std::pair<std::string, types::Scalar>> arguments;

    for (const auto& argument : j) {

        assert(argument.contains("name"));

        assert(argument["name"].is_string());

        assert(argument.contains("type"));

        assert(argument["type"].is_object());

        std::string name = argument["name"];

        auto type = json_to_type(argument["type"]);

        arguments.emplace_back(name, *dynamic_cast<types::Scalar*>(type.get()));

    }

    return arguments;

}

void JSONSerializer::json_to_dataflow(const nlohmann::json& j,

    std::unordered_map<std::string, data_flow::DataFlowNode&> nodes_map;

    assert(j.contains("nodes"));

    assert(j["nodes"].is_array());

    for (const auto& node : j["nodes"]) {

        assert(node.contains("type"));

        assert(node["type"].is_string());

        assert(node.contains("element_id"));

        assert(node["element_id"].is_string());

        std::string type = node["type"];

        if (type == "tasklet") {

            assert(node.contains("code"));

            assert(node["code"].is_number_integer());

            assert(node.contains("inputs"));

            assert(node["inputs"].is_array());

            assert(node.contains("output"));

            assert(node["output"].is_object());

            assert(node["output"].contains("name"));

            assert(node["output"].contains("type"));

            auto inputs = json_to_arguments(node["inputs"]);

            auto output_name = node["output"]["name"];

            auto output_type = json_to_type(node["output"]["type"]);

            auto& output_type_scalar = dynamic_cast<types::Scalar&>(*output_type);

            auto& tasklet =

                builder.add_tasklet(parent, node["code"], {output_name, output_type_scalar}, inputs,

                                    json_to_debug_info(node["debug_info"]));

            tasklet.element_id_ = node["element_id"];

            nodes_map.insert({node["element_id"], tasklet});

        } else if (type == "library_node") {

            assert(node.contains("code"));

            assert(node.contains("inputs"));

            assert(node["inputs"].is_array());

            assert(node.contains("outputs"));

            assert(node["outputs"].is_array());

            data_flow::LibraryNodeCode code(node["code"].get<std::string_view>());

                LibraryNodeSerializerRegistry::instance().get_library_node_serializer(code.value());

            if (serializer_fn == nullptr) {

                throw std::runtime_error("Unknown library node code: " + std::string(code.value()));

            auto serializer = serializer_fn();

            auto& lib_node = serializer->deserialize(node, builder, parent);

            lib_node.element_id_ = node["element_id"];

            nodes_map.insert({node["element_id"], lib_node});

        } else if (type == "access_node") {

            assert(node.contains("data"));

            auto& access_node =

                builder.add_access(parent, node["data"], json_to_debug_info(node["debug_info"]));

            access_node.element_id_ = node["element_id"];

            nodes_map.insert({node["element_id"], access_node});

        } else {

            throw std::runtime_error("Unknown node type");

    }

    assert(j.contains("edges"));

    assert(j["edges"].is_array());

    for (const auto& edge : j["edges"]) {

        assert(edge.contains("src"));

        assert(edge["src"].is_string());

        assert(edge.contains("dst"));

        assert(edge["dst"].is_string());

        assert(edge.contains("src_conn"));

        assert(edge["src_conn"].is_string());

        assert(edge.contains("dst_conn"));

        assert(edge["dst_conn"].is_string());

        assert(edge.contains("subset"));

        assert(edge["subset"].is_array());

        assert(nodes_map.find(edge["src"]) != nodes_map.end());

        assert(nodes_map.find(edge["dst"]) != nodes_map.end());

        auto& source = nodes_map.at(edge["src"]);

        auto& target = nodes_map.at(edge["dst"]);

        assert(edge.contains("subset"));

        assert(edge["subset"].is_array());

        std::vector<symbolic::Expression> subset;

        for (const auto& subset_str : edge["subset"]) {

            assert(subset_str.is_string());

            SymEngine::Expression subset_expr(subset_str);

            subset.push_back(subset_expr);

        }

        auto& memlet =

            builder.add_memlet(parent, source, edge["src_conn"], target, edge["dst_conn"], subset,

                               json_to_debug_info(edge["debug_info"]));

        memlet.element_id_ = edge["element_id"];

    }

}

void JSONSerializer::json_to_sequence(const nlohmann::json& j,

    assert(j.contains("type"));

    assert(j["type"].is_string());

    assert(j.contains("children"));

    assert(j["children"].is_array());

    assert(j.contains("transitions"));

    assert(j["transitions"].is_array());

    assert(j["transitions"].size() == j["children"].size());

    sequence.element_id_ = j["element_id"];

    sequence.debug_info_ = json_to_debug_info(j["debug_info"]);

    std::string type = j["type"];

    if (type == "sequence") {

        for (size_t i = 0; i < j["children"].size(); i++) {

            auto& child = j["children"][i];

            auto& transition = j["transitions"][i];

            assert(child.contains("type"));

            assert(child["type"].is_string());

            assert(transition.contains("type"));

            assert(transition["type"].is_string());

            assert(transition.contains("assignments"));

            assert(transition["assignments"].is_array());

            control_flow::Assignments assignments;

            for (const auto& assignment : transition["assignments"]) {

                assert(assignment.contains("symbol"));

                assert(assignment["symbol"].is_string());

                assert(assignment.contains("expression"));

                assert(assignment["expression"].is_string());

                SymEngine::Expression expr(assignment["expression"]);

                assignments.insert({symbolic::symbol(assignment["symbol"]), expr});

            }

            if (child["type"] == "block") {

                json_to_block_node(child, builder, sequence, assignments);

            } else if (child["type"] == "for") {

                json_to_for_node(child, builder, sequence, assignments);

            } else if (child["type"] == "if_else") {

                json_to_if_else_node(child, builder, sequence, assignments);

            } else if (child["type"] == "while") {

                json_to_while_node(child, builder, sequence, assignments);

            } else if (child["type"] == "break") {

                json_to_break_node(child, builder, sequence, assignments);

            } else if (child["type"] == "continue") {

                json_to_continue_node(child, builder, sequence, assignments);

            } else if (child["type"] == "return") {

                json_to_return_node(child, builder, sequence, assignments);

            } else if (child["type"] == "map") {

                json_to_map_node(child, builder, sequence, assignments);

            } else if (child["type"] == "sequence") {

                auto& subseq = builder.add_sequence(sequence, assignments,

                                                    json_to_debug_info(child["debug_info"]));

                json_to_sequence(child, builder, subseq);

            } else {

                throw std::runtime_error("Unknown child type");

            sequence.at(i).second.debug_info_ = json_to_debug_info(transition["debug_info"]);

            sequence.at(i).second.element_id_ = transition["element_id"];

        }

    } else {

        throw std::runtime_error("expected sequence type");

}

void JSONSerializer::json_to_block_node(const nlohmann::json& j,

    assert(j.contains("type"));

    assert(j["type"].is_string());

    assert(j.contains("dataflow"));

    assert(j["dataflow"].is_object());

    auto& block = builder.add_block(parent, assignments, json_to_debug_info(j["debug_info"]));

    block.element_id_ = j["element_id"];

    assert(j["dataflow"].contains("type"));

    assert(j["dataflow"]["type"].is_string());

    std::string type = j["dataflow"]["type"];

    if (type == "dataflow") {

        json_to_dataflow(j["dataflow"], builder, block);

    } else {

        throw std::runtime_error("Unknown dataflow type");

}

void JSONSerializer::json_to_for_node(const nlohmann::json& j,

    assert(j.contains("type"));

    assert(j["type"].is_string());

    assert(j.contains("indvar"));

    assert(j["indvar"].is_string());

    assert(j.contains("init"));

    assert(j["init"].is_string());

    assert(j.contains("condition"));

    assert(j["condition"].is_string());

    assert(j.contains("update"));

    assert(j["update"].is_string());

    assert(j.contains("root"));

    assert(j["root"].is_object());

    symbolic::Symbol indvar = symbolic::symbol(j["indvar"]);

    SymEngine::Expression init(j["init"]);

    SymEngine::Expression condition_expr(j["condition"]);

    assert(!SymEngine::rcp_static_cast<const SymEngine::Boolean>(condition_expr.get_basic())

        SymEngine::rcp_static_cast<const SymEngine::Boolean>(condition_expr.get_basic());

    SymEngine::Expression update(j["update"]);

    auto& for_node = builder.add_for(parent, indvar, condition, init, update, assignments,

                                     json_to_debug_info(j["debug_info"]));

    for_node.element_id_ = j["element_id"];

    assert(j["root"].contains("type"));

    assert(j["root"]["type"].is_string());

    assert(j["root"]["type"] == "sequence");

    json_to_sequence(j["root"], builder, for_node.root());

}

void JSONSerializer::json_to_if_else_node(const nlohmann::json& j,

    assert(j.contains("type"));

    assert(j["type"].is_string());

    assert(j["type"] == "if_else");