Add Fastor library
This commit is contained in:
Bassem Girgis
293
noarch/include/Fastor/tensor_algebra/strided_contraction.h
Normal file
293
noarch/include/Fastor/tensor_algebra/strided_contraction.h
Normal file
@@ -0,0 +1,293 @@
|
||||
#ifndef REDUCIBLE_CONTRACTION_H
|
||||
#define REDUCIBLE_CONTRACTION_H
|
||||
|
||||
#include "Fastor/tensor/Tensor.h"
|
||||
#include "Fastor/tensor_algebra/indicial.h"
|
||||
|
||||
|
||||
namespace Fastor {
|
||||
|
||||
// Broadcast-vectorisable contractions (if last dimension is contracted).
|
||||
// Requres working on general strides
|
||||
|
||||
|
||||
|
||||
template<class T, class U>
|
||||
struct extractor_reducible_contract {};
|
||||
|
||||
template<size_t ... Idx0, size_t ... Idx1>
|
||||
struct extractor_reducible_contract<Index<Idx0...>, Index<Idx1...>> {
|
||||
|
||||
|
||||
#if CONTRACT_OPT==2
|
||||
|
||||
template<typename T, size_t ... Rest0, size_t ... Rest1>
|
||||
static
|
||||
typename contraction_impl<Index<Idx0...,Idx1...>, Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)+sizeof...(Rest1)>::type>::type
|
||||
contract_impl(const Tensor<T,Rest0...> &a, const Tensor<T,Rest1...> &b) {
|
||||
|
||||
static_assert(!is_pair_reduction_v<Index<Idx0...>,Index<Idx1...>>,"REDUCTION TO SCALAR REQUESTED. USE REDUCTION FUNCTION INSTEAD");
|
||||
|
||||
constexpr int total = no_of_loops_to_set<Index<Idx0...>,Index<Idx1...>,Tensor<T,Rest0...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<no_of_unique<Idx0...,Idx1...>::value>::type>::value;
|
||||
|
||||
using index_generator = contract_meta_engine<Index<Idx0...>,Index<Idx1...>,Tensor<T,Rest0...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<total>::type>;
|
||||
|
||||
using OutTensor = typename index_generator::OutTensor;
|
||||
using OutIndices = typename index_generator::OutIndices;
|
||||
|
||||
constexpr auto& index_a = index_generator::index_a;
|
||||
constexpr auto& index_b = index_generator::index_b;
|
||||
constexpr auto& index_out = index_generator::index_out;
|
||||
|
||||
OutTensor out;
|
||||
out.zeros();
|
||||
const T *a_data = a.data();
|
||||
const T *b_data = b.data();
|
||||
T *out_data = out.data();
|
||||
|
||||
// Check for reducible vectorisability
|
||||
constexpr int general_stride = general_stride_finder<Index<Idx0...>,Index<Idx1...>,
|
||||
Tensor<T,Rest0...>,Tensor<T,Rest1...>, typename std_ext::make_index_sequence<sizeof...(Rest1)>::type>::value;
|
||||
|
||||
|
||||
#ifndef FASTOR_DONT_VECTORISE
|
||||
using vectorisability = is_reducibly_vectorisable<OutIndices,OutTensor>;
|
||||
constexpr int stride = vectorisability::stride;
|
||||
using V = typename vectorisability::type;
|
||||
#else
|
||||
constexpr int stride = 1;
|
||||
using V = SIMDVector<T,sizeof(T)*8>;
|
||||
#endif
|
||||
|
||||
V _vec_a, _vec_b;
|
||||
for (int i = 0; i < total; i+=stride) {
|
||||
_vec_a.set(*(a_data+index_a[i]));
|
||||
const auto idx_b = index_b[i];
|
||||
vector_setter(_vec_b, b_data, idx_b, general_stride);
|
||||
V _vec_out = _vec_a*_vec_b + V(out_data+index_out[i]);
|
||||
_vec_out.store(out_data+index_out[i]);
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
|
||||
#elif CONTRACT_OPT==1
|
||||
|
||||
template<typename T, size_t ... Rest0, size_t ... Rest1>
|
||||
static
|
||||
typename contraction_impl<Index<Idx0...,Idx1...>, Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)+sizeof...(Rest1)>::type>::type
|
||||
contract_impl(const Tensor<T,Rest0...> &a, const Tensor<T,Rest1...> &b) {
|
||||
|
||||
static_assert(!is_pair_reduction_v<Index<Idx0...>,Index<Idx1...>>,"REDUCTION TO SCALAR REQUESTED. USE REDUCTION FUNCTION INSTEAD");
|
||||
|
||||
using OutTensor = typename contraction_impl<Index<Idx0...,Idx1...>, Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)+sizeof...(Rest1)>::type>::type;
|
||||
using OutIndices = typename contraction_impl<Index<Idx0...,Idx1...>, Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)+sizeof...(Rest1)>::type>::indices;
|
||||
|
||||
OutTensor out;
|
||||
out.zeros();
|
||||
const T *a_data = a.data();
|
||||
const T *b_data = b.data();
|
||||
T *out_data = out.data();
|
||||
|
||||
constexpr int a_dim = sizeof...(Rest0);
|
||||
constexpr int b_dim = sizeof...(Rest1);
|
||||
|
||||
constexpr auto& idx_a = IndexFirstTensor<Index<Idx0...>,Index<Idx1...>, Tensor<T,Rest0...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)>::type>::indices;
|
||||
constexpr auto& idx_b = IndexSecondTensor<Index<Idx0...>,Index<Idx1...>, Tensor<T,Rest0...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest1)>::type>::indices;
|
||||
constexpr auto& idx_out = IndexResultingTensor<Index<Idx0...>,Index<Idx1...>, Tensor<T,Rest0...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<OutTensor::Dimension>::type>::indices;
|
||||
|
||||
constexpr int total = no_of_loops_to_set<Index<Idx0...>,Index<Idx1...>,Tensor<T,Rest0...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<no_of_unique<Idx0...,Idx1...>::value>::type>::value;
|
||||
|
||||
using maxes_out_type = typename no_of_loops_to_set<Index<Idx0...>,Index<Idx1...>,Tensor<T,Rest0...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<no_of_unique<Idx0...,Idx1...>::value>::type>::type;
|
||||
|
||||
constexpr auto& as_all = cartesian_product<maxes_out_type,typename std_ext::make_index_sequence<total>::type>::values;
|
||||
|
||||
constexpr std::array<size_t,a_dim> products_a = nprods<Index<Rest0...>,typename std_ext::make_index_sequence<a_dim>::type>::values;
|
||||
constexpr std::array<size_t,b_dim> products_b = nprods<Index<Rest1...>,typename std_ext::make_index_sequence<b_dim>::type>::values;
|
||||
using Index_with_dims = typename put_dims_in_Index<OutTensor>::type;
|
||||
constexpr std::array<size_t,Index_with_dims::NoIndices> products_out = nprods<Index_with_dims,
|
||||
typename std_ext::make_index_sequence<Index_with_dims::NoIndices>::type>::values;
|
||||
|
||||
|
||||
// Check for reducible vectorisability
|
||||
constexpr int general_stride = general_stride_finder<Index<Idx0...>,Index<Idx1...>,
|
||||
Tensor<T,Rest0...>,Tensor<T,Rest1...>, typename std_ext::make_index_sequence<b_dim>::type>::value;
|
||||
|
||||
|
||||
#ifndef FASTOR_DONT_VECTORISE
|
||||
using vectorisability = is_reducibly_vectorisable<OutIndices,OutTensor>;
|
||||
constexpr int stride = vectorisability::stride;
|
||||
using V = typename vectorisability::type;
|
||||
#else
|
||||
constexpr int stride = 1;
|
||||
using V = SIMDVector<T,sizeof(T)*8>;
|
||||
#endif
|
||||
|
||||
int it;
|
||||
V _vec_a, _vec_b;
|
||||
|
||||
for (int i = 0; i < total; i+=stride) {
|
||||
int index_a = as_all[i][idx_a[a_dim-1]];
|
||||
for(it = 0; it< a_dim; it++) {
|
||||
index_a += products_a[it]*as_all[i][idx_a[it]];
|
||||
}
|
||||
|
||||
int index_b = as_all[i][idx_b[b_dim-1]];
|
||||
for(it = 0; it< b_dim; it++) {
|
||||
index_b += products_b[it]*as_all[i][idx_b[it]];
|
||||
}
|
||||
int index_out = as_all[i][idx_out[idx_out.size()-1]];
|
||||
for(it = 0; it< idx_out.size(); it++) {
|
||||
index_out += products_out[it]*as_all[i][idx_out[it]];
|
||||
}
|
||||
|
||||
_vec_a.set(*(a_data+index_a));
|
||||
// _vec_b.set(b_data[index_b+3*general_stride], b_data[index_b+2*general_stride],
|
||||
// b_data[index_b+general_stride],b_data[index_b]);
|
||||
vector_setter(_vec_b, b_data, index_b, general_stride);
|
||||
V _vec_out = _vec_a*_vec_b + V(out_data+index_out);
|
||||
_vec_out.store(out_data+index_out);
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
|
||||
#else
|
||||
|
||||
template<typename T, size_t ... Rest0, size_t ... Rest1>
|
||||
static
|
||||
typename contraction_impl<Index<Idx0...,Idx1...>, Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)+sizeof...(Rest1)>::type>::type
|
||||
contract_impl(const Tensor<T,Rest0...> &a, const Tensor<T,Rest1...> &b) {
|
||||
|
||||
static_assert(!is_pair_reduction_v<Index<Idx0...>,Index<Idx1...>>,"REDUCTION TO SCALAR REQUESTED. USE REDUCTION FUNCTION INSTEAD");
|
||||
|
||||
using OutTensor = typename contraction_impl<Index<Idx0...,Idx1...>, Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)+sizeof...(Rest1)>::type>::type;
|
||||
using OutIndice = typename contraction_impl<Index<Idx0...,Idx1...>, Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)+sizeof...(Rest1)>::type>::indices;
|
||||
|
||||
OutTensor out;
|
||||
out.zeros();
|
||||
const T *a_data = a.data();
|
||||
const T *b_data = b.data();
|
||||
T *out_data = out.data();
|
||||
|
||||
constexpr int a_dim = sizeof...(Rest0);
|
||||
constexpr int b_dim = sizeof...(Rest1);
|
||||
constexpr int out_dim = no_of_unique<Idx0...,Idx1...>::value;
|
||||
|
||||
constexpr auto& idx_a = IndexTensors<
|
||||
Index<Idx0..., Idx1...>,
|
||||
Tensor<T,Rest0...,Rest1...>,
|
||||
Index<Idx0...>,Tensor<T,Rest0...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest0)>::type>::indices;
|
||||
|
||||
constexpr auto& idx_b = IndexTensors<
|
||||
Index<Idx0..., Idx1...>,
|
||||
Tensor<T,Rest0...,Rest1...>,
|
||||
Index<Idx1...>,Tensor<T,Rest1...>,
|
||||
typename std_ext::make_index_sequence<sizeof...(Rest1)>::type>::indices;
|
||||
|
||||
constexpr auto& idx_out = IndexTensors<
|
||||
Index<Idx0..., Idx1...>,
|
||||
Tensor<T,Rest0...,Rest1...>,
|
||||
OutIndice,OutTensor,
|
||||
typename std_ext::make_index_sequence<OutTensor::Dimension>::type>::indices;
|
||||
|
||||
using nloops = loop_setter<
|
||||
Index<Idx0...,Idx1...>,
|
||||
Tensor<T,Rest0...,Rest1...>,
|
||||
typename std_ext::make_index_sequence<out_dim>::type>;
|
||||
constexpr auto& maxes_out = nloops::dims;
|
||||
constexpr int total = nloops::value;
|
||||
|
||||
constexpr std::array<size_t,a_dim> products_a = nprods<Index<Rest0...>,typename std_ext::make_index_sequence<a_dim>::type>::values;
|
||||
constexpr std::array<size_t,b_dim> products_b = nprods<Index<Rest1...>,typename std_ext::make_index_sequence<b_dim>::type>::values;
|
||||
|
||||
using Index_with_dims = typename put_dims_in_Index<OutTensor>::type;
|
||||
constexpr std::array<size_t,OutTensor::Dimension> products_out = \
|
||||
nprods<Index_with_dims,typename std_ext::make_index_sequence<OutTensor::Dimension>::type>::values;
|
||||
|
||||
// Check for reducible vectorisability
|
||||
constexpr int general_stride = general_stride_finder<Index<Idx0...>,Index<Idx1...>,
|
||||
Tensor<T,Rest0...>,Tensor<T,Rest1...>, typename std_ext::make_index_sequence<b_dim>::type>::value;
|
||||
|
||||
#ifndef FASTOR_DONT_VECTORISE
|
||||
using vectorisability = is_reducibly_vectorisable<OutIndice,OutTensor>;
|
||||
constexpr int stride = vectorisability::stride;
|
||||
using V = typename vectorisability::type;
|
||||
#else
|
||||
constexpr int stride = 1;
|
||||
using V = SIMDVector<T,simd_abi::scalar>;
|
||||
#endif
|
||||
|
||||
int as[out_dim];
|
||||
std::fill(as,as+out_dim,0);
|
||||
|
||||
int it;
|
||||
V _vec_a, _vec_b;
|
||||
|
||||
for (int i = 0; i < total; i+=stride) {
|
||||
int remaining = total;
|
||||
for (int n = 0; n < out_dim; ++n) {
|
||||
remaining /= maxes_out[n];
|
||||
as[n] = ( i / remaining ) % maxes_out[n];
|
||||
}
|
||||
|
||||
int index_a = as[idx_a[a_dim-1]];
|
||||
for(it = 0; it< a_dim; it++) {
|
||||
index_a += products_a[it]*as[idx_a[it]];
|
||||
}
|
||||
int index_b = as[idx_b[b_dim-1]];
|
||||
for(it = 0; it< b_dim; it++) {
|
||||
index_b += products_b[it]*as[idx_b[it]];
|
||||
}
|
||||
int index_out = as[idx_out[OutTensor::Dimension-1]];
|
||||
for(it = 0; it< static_cast<int>(OutTensor::Dimension); it++) {
|
||||
index_out += products_out[it]*as[idx_out[it]];
|
||||
}
|
||||
|
||||
// asm("#BEGIN");
|
||||
_vec_a.set(*(a_data+index_a));
|
||||
//_vec_a.broadcast(a_data+index_a);
|
||||
vector_setter(_vec_b, b_data, index_b, general_stride);
|
||||
V _vec_out = _vec_a*_vec_b + V(out_data+index_out);
|
||||
_vec_out.store(out_data+index_out);
|
||||
// asm("#END");
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
|
||||
#endif
|
||||
|
||||
};
|
||||
|
||||
|
||||
|
||||
template<class Index_I, class Index_J,
|
||||
typename T, size_t ... Rest0, size_t ... Rest1>
|
||||
FASTOR_INLINE auto strided_contraction(const Tensor<T,Rest0...> &a, const Tensor<T,Rest1...> &b)
|
||||
-> decltype(extractor_reducible_contract<Index_I,Index_J>::contract_impl(a,b)) {
|
||||
return extractor_reducible_contract<Index_I,Index_J>::contract_impl(a,b);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
#endif // REDUCIBLE_CONTRACTION_H
|
||||
Reference in New Issue
Block a user