folding.cpp

/** 
 * @brief This class represents a "layer" of analysis. It calculates the
 * cost of the atomic units of this layer, then folds them into a problem
 * formulation for the next layer.
 */
#include "folding.h"
#include "latency.hpp"
#include <memory>
#include <string>

#include <isl/aff.h>
#include <isl/map.h>
#include <isl/polynomial.h>
#include <isl/set.h>

#include <barvinok/isl.h>
#include <barvinok/barvinok.h>
#include <barvinok/polylib.h>

class BranchTwig
{
    public:
        /// @brief The cost formula of the folding step for this layer.
        const std::string crease_costs;
        /** 
         * @brief The folding action to a multicastable representation after
         * calculating the cost of folding. 
         */
        const std::string fold_formula;
        /** 
         * @brief Isolation function for the highest value along the axis of
         * multicast for folding.
         * @todo This is a temporary solution to the problem of folding, may
         * be able to automate this process rather than via manual input.
         */
        const std::string all_after;
        /// @brief The cost formula of the multicasting step for this layer.
        const std::string multicast_costs;
        /// @brief The src collapse formulation for the next layer.
        const std::string src_collapser;
        /// @brief The dst collapse formulation for the next layer.
        const std::string dst_collapser;
        /// @brief Whether or not this is uses virtual srcs to force a certain
        /// multicast behavior.
        const bool virtual_layer;
        /// @brief The context the layer is in.
        isl_ctx *const ctx;
    public:
        /** 
         * @brief Constructs a layer from a cost formulation and a folding
         * formulation.
         * 
         * @param crease_costs The cost formula of unmulticastable datum as 
         * an ISL string. Goes under the assumption the input is either the 
         * starting geometry architecture or the output of a previous layer.
         * @param fold_formula The isl_map in a string representation that
         * projects away the unmulticastable portions of the path of a datum to
         * a dst. This is what we refer to as "folding".
         * @param all_after The function to determine the highest value along
         * the axis of multicast for folding.
         * @param multicast_costs The cost formula for multicasting a datum
         * as an ISL string. Goes under the assumption that the input is of
         * the form of this Layer's ISL representation after folding.
         * @param collapse_formulas The collapse formulations for how to 
         * translate the srcs and dsts of this layer to the inputs that work 
         * with next layer.
         * @param virtual_layer Whether or not this is a virtual layer and thus
         * to skip the calculation of missing data, since it is all missing.
         * @param ctx The context the layer is in.
         */
        BranchTwig(
            const std::string& crease_costs, const std::string& fold_formula, 
            const std::string& all_after, const std::string& multicast_costs,
            const collapse& collapse_formulas, bool virtual_layer, isl_ctx *const ctx
        ):
        crease_costs(crease_costs), fold_formula(fold_formula), 
        all_after(all_after), multicast_costs(multicast_costs), 
        src_collapser(collapse_formulas->src_collapser), 
        dst_collapser(collapse_formulas->dst_collapser),
        virtual_layer(virtual_layer), ctx(ctx) {}

        BranchTwig(
            const std::string& crease_costs, const std::string& fold_formula, 
            const std::string& all_after, const std::string& multicast_costs,
            const collapse& collapse_formulas, isl_ctx *const ctx
        ): BranchTwig(
            crease_costs, fold_formula, all_after, multicast_costs, 
            collapse_formulas, false, ctx
        ) {}

        /** 
         * @brief Calculates the cost of the atomic units of this layer, then
         * returns a struct of the binding cost at layer and the next layer's
         * abstraction as a string. 
         * 
         * @param s_srcs The sources of the bindings at this layer as an ISL string.
         * @param s_dsts The destinations of the bindings at this layer as an ISL string.
         * 
         * @return A struct of the binding abstraction for the next layer.
         */
        const abstracted_binding evaluate(const std::string& s_srcs, const std::string& s_dsts)
        {
            // Folds the destinations onto their connected trunk.
            const fold_result fold_res = this->fold(s_dsts);
            if (islIntermediates) std::cout << "Crease Cost: " << fold_res->cost << std::endl;

            // std::cout << "Folded: " << fold_res->folded_repr << std::endl;

            // Calculates the cost to every folded node per datum.
            const long casting_cost = this->multicast(fold_res->folded_repr);
            if (islIntermediates) std::cout << "Casting Cost: " << casting_cost << std::endl;

            // Formulates the abstraction cost of this layer.
            const abstraction_cost cost = abstraction_cost(
                new abstraction_cost_struct{fold_res->cost, casting_cost}
            );

            // Calculates the requests that are not satisfied by the layer.
            ///@todo Collapse the folded destinations into the next layer.
            const binding collapsed = this->collapse(s_srcs, s_dsts);
            
            // Returns evaluation results.
            return abstracted_binding(new abstracted_binding_struct{collapsed, cost});
        }

        /// @brief Wraps evaluate by accepting the binding as a struct.
        const abstracted_binding inline evaluate(const binding& b)
        {
            return this->evaluate(b->srcs, b->dsts);
        }

        /// @brief Wraps evaluate by accepting an abstracted binding as a struct.
        const abstracted_binding inline evaluate(const abstracted_binding& b)
        {
            return this->evaluate(b->abstraction);
        }
    private:
        /** 
         * @brief Folds the destinations onto their connected trunk. 
         * 
         * @param s_dsts The destinations to fold as an ISL string.
         * @return A unique_ptr to a struct holding costs of the folding step and
         * the folded destinations as an ISL string.
         */
        fold_result fold(const std::string& dsts)
        {
            // Reads the dsts into isl format.
            isl_map *p_dsts = isl_map_read_from_str(ctx, dsts.c_str());
            DUMP(p_dsts);
            
            /// @note Gets the total cost of the folded dsts.
            // Returns { [id, x, y] -> number_of_data}
            isl_pw_qpolynomial *p_card = isl_map_card(isl_map_copy(p_dsts));
            DUMP(p_card);
            // Calculates the cost per datum per dst cast from the trunk.
            isl_pw_qpolynomial *p_fold_cost = isl_pw_qpolynomial_read_from_str(ctx, this->crease_costs.c_str());
            DUMP(p_fold_cost);
            // Calculates the cost per dst cast from the trunk.
            isl_pw_qpolynomial *p_cost_at_dst = isl_pw_qpolynomial_mul(p_card, p_fold_cost);
            DUMP(p_cost_at_dst);
            // Calculates the cost to cast all data from the trunk.
            isl_pw_qpolynomial *p_total_cost = isl_pw_qpolynomial_sum(p_cost_at_dst);
            DUMP(p_total_cost);
            // Reads the value from p_total_cost.
            isl_val *v_total_cost = isl_pw_qpolynomial_eval(p_total_cost, isl_point_zero(isl_pw_qpolynomial_get_domain_space(p_total_cost)));
            // Initializes the variable storing the cost of folding.
            long fold_cost = isl_val_get_num_si(v_total_cost);
            // Frees the values.
            isl_val_free(v_total_cost);

            /// @note Folds p_dsts onto the trunk according to the fold formula.
            // Reads the fold formula into isl format.
            isl_map *p_fold = isl_map_read_from_str(ctx, this->fold_formula.c_str());
            // Converts dsts->data to data->dsts
            isl_map *p_data_to_dsts = isl_map_reverse(p_dsts);
            DUMP(p_data_to_dsts);
            // Folds the dsts onto the trunk.
            isl_map *p_folded = isl_map_apply_range(p_data_to_dsts, p_fold);
            p_folded = isl_map_reverse(p_folded);
            DUMP(p_folded);
            // Gets the largest y value per datum.
            isl_map *p_all_after = isl_map_read_from_str(ctx, this->all_after.c_str());
            DUMP(p_all_after);
            isl_map *p_max_y = isl_map_apply_range(p_all_after, isl_map_copy(p_folded));
            DUMP(p_max_y);
            isl_map *p_folded_condensed = isl_map_subtract(p_folded, p_max_y);
            p_folded_condensed = isl_map_reverse(p_folded_condensed);
            // Converts p_folded_condensed to a string.
            std::string s_folded_condensed = isl_map_to_str(p_folded_condensed);
            // Frees the maps.
            isl_map_free(p_folded_condensed);

            // Initializes the struct ptr to return.
            fold_result result = fold_result(new fold_struct{fold_cost, s_folded_condensed});

            return result;
        }

        /** 
         * @brief Calculates the the cost to every folded node per datum.
         * 
         * @param folded_geometry The folded destinations from this->fold(*).
         * @return The cost of multicasting to the folded destinations.
         */
        long multicast(const std::string& s_folded_bindings)
        {
            // Reads the folded destinations into isl format.
            isl_map *p_folded_bindings = isl_map_read_from_str(ctx, s_folded_bindings.c_str());
            DUMP(p_folded_bindings);
            // std::cout << s_folded_bindings << std::endl;
            // Reads the multicast cost formulation into isl format.
            isl_pw_qpolynomial *p_cast_cost = isl_pw_qpolynomial_read_from_str(ctx, this->multicast_costs.c_str());
            DUMP(p_cast_cost);

            /** 
             * @note Calculates the cost of multicasting to the folded dsts
             * according to architecture spec. */
            // Applies the cost formulation to the folded dsts.
            isl_pw_qpolynomial *p_cost_applied = isl_map_apply_pw_qpolynomial(p_folded_bindings, p_cast_cost);
            DUMP(p_cost_applied);
            // Sums all the costs.
            isl_pw_qpolynomial *p_total_cost = isl_pw_qpolynomial_sum(p_cost_applied);
            DUMP(p_total_cost);
            // Evaluates the cost.
            isl_val *v_total_cost = isl_pw_qpolynomial_eval(p_total_cost, isl_point_zero(isl_pw_qpolynomial_get_domain_space(p_total_cost)));
            // Initializes the return value to the cost of multicasting.
            long cost = isl_val_get_d(v_total_cost);
            // Frees the values.
            isl_val_free(v_total_cost);

            return cost;
        }

        /**
         * @brief Identifies the requests that are not satisfied by the layer 
         * and passes them as dsts to the next layer. Also calculates the cost 
         * of getting all srcs to a position accessible by the next layer.
         * 
         * @param s_srcs The sources of the bindings at this layer as an ISL string.
         * @param s_dsts The destinations of the bindings at this layer as an ISL string.
         * 
         * @return The collapsed binding abstraction for the next layer.
         */
        binding collapse(const std::string& s_srcs, const std::string& s_dsts)
        {
            // Reads the srcs and dsts into isl format.
            isl_map *p_srcs = isl_map_read_from_str(ctx, s_srcs.c_str());
            isl_map *p_dsts = isl_map_read_from_str(ctx, s_dsts.c_str());
            // Reads the collapse formulas into isl format.
            isl_map *p_dsts_collapser = isl_map_read_from_str(
                ctx, this->dst_collapser.c_str()
            );
            isl_map *p_srcs_collapser = isl_map_read_from_str(
                ctx, this->src_collapser.c_str()
            );

            // Collapses all src requests.
            isl_map *p_collapsed_srcs = isl_map_apply_range(p_srcs_collapser, isl_map_copy(p_srcs));
            DUMP(p_collapsed_srcs);
            // Converts p_collapsed_srcs to a string.
            std::string s_collapsed_srcs = isl_map_to_str(p_collapsed_srcs);
            // Collapses all dst requests to the same format as their SRCs.
            isl_map *p_collapsed_dsts = isl_map_apply_range(p_dsts_collapser, p_dsts);
            DUMP(p_collapsed_dsts);

            // Calculates the requests that are not satisfied by the layer.
            isl_map *p_missing_data = NULL;
            if (!this->virtual_layer) {
                p_missing_data = isl_map_subtract(p_collapsed_dsts, p_srcs);
            } else {
                if (islIntermediates) std::cout << "Virtual layer" << std::endl;
                p_missing_data = p_collapsed_dsts;
            }
            DUMP(p_missing_data);
            // Converts p_missing to a string.
            std::string s_missing_data = isl_map_to_str(p_missing_data);
            // Frees the maps.
            isl_map_free(p_missing_data);

            // Initializes the collapsed binding abstraction for the next layer.
            binding collapsed = binding(new binding_struct{s_collapsed_srcs, s_missing_data});
            return collapsed;
        }
};

int main(int argc, char* argv[])
{
    // Creates an isl context.
    isl_ctx *ctx = isl_ctx_alloc();

    // Creates the binding abstraction for the first layer.
    int M_int = 1024;
    int N_int = 1024;
    std::string M = std::to_string(M_int);
    std::string N = std::to_string(N_int);
    std::vector<int> D_vals({1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024});
    clock_t start, end;
    double cpu_time_used;
    for (int D_int : D_vals) {
        start = clock();
        std::string D = std::to_string(D_int);
        std::string srcs = R"SRC(
            {trunk[id, y] -> data[a, b] : id = 0}
        )SRC";
        // Defines the src occupancy map as a string.
        std::string dsts =  "{dst[id, x, y] -> data[a, b] : ("+D+"*x)%"+M+" <= a <= ("+
                            D+"*x+"+D+"-1)%"+M+" and b=y and 0 <= x < "+M+
                            " and 0 <= y < "+N+" and 0 <= a < "+M+" and 0 <= b < "
                            +N+" and id = 0 }";
        binding test_case = binding(new binding_struct{srcs, dsts});

        // Calculates the cost formulas of the first layer.
        /// @note Read right to left like function composition.
        std::string leaf_multicast_costs = "{ dst[id, x, y] -> x }";
        std::string leaf_crease_costs = "{ dst[id, x, y] -> 0 }";

        // Calculates the collapse formulas of the leaf layer.
        std::string leaf_fold_formula = "{ dst[id, x, y] -> dst[id, x, y] }";
        std::string leaf_all_after = "{ dst[id, x, y] -> dst[id, x', y] : x' > x }";
        std::string leaf_dst_collapse_formula = "{ trunk[id, y] -> dst[id, x, y] }";
        std::string leaf_src_collapse_formula = "{ off[id] -> trunk[id, y] }";
        collapse leaf_collapse_formulas = collapse(new collapse_struct{leaf_src_collapse_formula, leaf_dst_collapse_formula});
        BranchTwig leaves = BranchTwig(
            leaf_crease_costs, leaf_fold_formula, 
            leaf_all_after, leaf_multicast_costs, 
            leaf_collapse_formulas, true, ctx
        );
        abstracted_binding collapsed = leaves.evaluate(test_case);

        if (islIntermediates) std::cout << "Collapsed: " << collapsed->abstraction->srcs << std::endl;
        if (islIntermediates) std::cout << "Missing: " << collapsed->abstraction->dsts << std::endl;
        std::cout << "Leaf Costs: " << collapsed->cost->multicast_cost + collapsed->cost->fold_cost << std::endl;

        // Calculates the collapse formulas of the first layer.
        std::string fold_formula = "{ trunk[id, y] -> trunk[id, y]  }";
        std::string all_after = "{ trunk[id, y] -> trunk[id, y'] : y' > y }";
        std::string multicast_costs = "{ trunk[id, y] -> y+1 }";
        std::string crease_costs = "{ trunk[id, y] -> 0 }";
        std::string dst_collapse_formula = "{ off[id] -> trunk[id, y] }";
        std::string src_collapse_formula = "{ null[0] -> off[id] }";
        collapse collapse_formulas = collapse(new collapse_struct{src_collapse_formula, dst_collapse_formula});

        BranchTwig test = BranchTwig(
            crease_costs, fold_formula, 
            all_after, multicast_costs, 
            collapse_formulas, ctx
        );
        if (islIntermediates) std::cout << "####Evaluating...####" << std::endl;
        abstracted_binding final = test.evaluate(collapsed);
        std::cout << "Trunk Costs: " << final->cost->multicast_cost + final->cost->fold_cost << std::endl;

        end = clock();
        cpu_time_used = ((double) (end - start)) / CLOCKS_PER_SEC;
        if (islIntermediates) std::cout << "time: " << cpu_time_used << " | D: " << D << std::endl;
    }

    ///@note Frees ctx to check for memory leaks through ISL.
    isl_ctx_free(ctx);

    return 0;
}