Einsum.h

Go to the documentation of this file.

 
#ifndef KIDS_EINSUM_H
#define KIDS_EINSUM_H
 
#include <cstring>
#include <string>
#include <vector>
 
namespace PROJECT_NS {
 
//*********************************************************************************
struct EinsumIdx {
    EinsumIdx(char _label, std::size_t _cnt, std::size_t _dim, std::size_t _val)
        : label{_label}, cnt{_cnt}, dim{_dim}, val{_val} {};
 
    char        label;    
    std::size_t cnt = 0;  
    std::size_t dim = 0;  
    std::size_t val = 0;  
};
 
struct DimenHelper {
   public:
    std::size_t esshape_rank;  
    std::size_t total_esidx;   
 
    std::vector<std::size_t> dims;      
    std::vector<std::size_t> ldims;     
    std::vector<std::size_t> es_ldims;  
    std::vector<std::size_t> mapldims;  
 
    DimenHelper(){};
 
    DimenHelper(const std::string& esshape, std::vector<EinsumIdx>& idx_vec);
};
 
 
class EinsumHelper {
   public:
    std::size_t total_esidx;   
    std::size_t total_tensor;  
 
    std::vector<EinsumIdx>   einsum_idxs;  
    std::vector<std::size_t> einsum_dims;  
 
    std::vector<std::string> fixed_label_names;  
 
    std::vector<std::string> esshape_inputs;       
    std::string              esshape_output = "";  
 
    std::vector<DimenHelper> dh_inputs;  
    DimenHelper              dh_output;  
 
    std::vector<std::size_t> einsum_iposes;  
    std::vector<std::size_t> ipos_inputs;    
 
    int count1     = 0;
    int count2     = 0;
    int count3     = 0;
    int total_loop = 0;
 
    EinsumHelper(const std::string&                    einsum_expression,  //
                 std::vector<std::vector<std::size_t>> shape_inputs,       //
                 std::vector<std::size_t>              shape_output = {}   //
    );
};
 
template <typename T>
void einsum(EinsumHelper&          EH,           //
            const std::vector<T*>& data_inputs,  //
            T*                     data_output   //
) {
    auto& einsum_dims   = EH.einsum_dims;
    auto& einsum_iposes = EH.einsum_iposes;
    auto& ipos_inputs   = EH.ipos_inputs;
    // ipos_output
    auto& dh_inputs          = EH.dh_inputs;
    auto& dh_output_mapldims = EH.dh_output.mapldims;
 
    std::size_t total_loop   = EH.total_loop;
    std::size_t total_tensor = EH.total_tensor;
    std::size_t total_esidx  = EH.total_esidx;
    std::size_t imax         = EH.count3 - 1;
    std::size_t imin         = EH.count1;
 
    memset(einsum_iposes.data(), 0, total_esidx * sizeof(std::size_t));
    memset(ipos_inputs.data(), 0, total_tensor * sizeof(std::size_t));
    bool reset_zero = true;
    data_output[0]  = T(0);
    for (std::size_t iloop = 0, ipos_output = 0; iloop < total_loop; ++iloop) {
        if (reset_zero) data_output[ipos_output] = T(0);
 
        T term = T(1);
        for (int iten = 0; iten < total_tensor; ++iten) { term *= data_inputs[iten][ipos_inputs[iten]]; }
        data_output[ipos_output] += term;
 
        std::size_t i = imax;
        while (++einsum_iposes[i] == einsum_dims[i] && i > imin) { einsum_iposes[i--] = 0; }
        reset_zero = (i < EH.count2);
 
        for (int iten = 0; iten < total_tensor; ++iten)  //
            ipos_inputs[iten] += dh_inputs[iten].mapldims[i];
        ipos_output += dh_output_mapldims[i];
    }
}
 
 
template <typename T>
void einsum(const std::string&                           einsum_expression,  //
            std::vector<T*>                              data_inputs,        //
            const std::vector<std::vector<std::size_t>>& shape_inputs,       //
            T*                                           data_output,        //
            const std::vector<std::size_t>&              shape_output = {}   //
) {
    EinsumHelper EH(einsum_expression, shape_inputs, shape_output);
    einsum(EH, data_inputs, data_output);
}
 
};  // namespace PROJECT_NS
 
#endif  // KIDS_EINSUM_H