10 #ifndef EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_H
11 #define EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_H
32 template<
typename Scalar,
typename Index,
int side,
34 typename nocontract_t,
typename contract_t,
35 int packet_size,
bool inner_dim_contiguous>
36 class BaseTensorContractionMapper {
39 BaseTensorContractionMapper(
const Tensor& tensor,
40 const nocontract_t& nocontract_strides,
41 const nocontract_t& ij_strides,
42 const contract_t& contract_strides,
43 const contract_t& k_strides) :
45 m_nocontract_strides(nocontract_strides),
46 m_ij_strides(ij_strides),
47 m_contract_strides(contract_strides),
48 m_k_strides(k_strides) { }
51 EIGEN_STRONG_INLINE
void prefetch(Index ) { }
54 EIGEN_STRONG_INLINE Scalar operator()(Index row)
const {
56 return operator()(row, 0);
60 EIGEN_STRONG_INLINE Scalar operator()(Index row, Index col)
const {
61 return m_tensor.coeff(computeIndex(row, col));
65 EIGEN_STRONG_INLINE Index computeIndex(Index row, Index col)
const {
66 const bool left = (side == Lhs);
67 Index nocontract_val = left ? row : col;
69 for (
int i = static_cast<int>(array_size<nocontract_t>::value) - 1; i > 0; i--) {
70 const Index idx = nocontract_val / m_ij_strides[i];
71 linidx += idx * m_nocontract_strides[i];
72 nocontract_val -= idx * m_ij_strides[i];
74 if (array_size<typename Tensor::Dimensions>::value > array_size<contract_t>::value) {
75 if (side == Lhs && inner_dim_contiguous) {
76 eigen_assert(m_nocontract_strides[0] == 1);
77 linidx += nocontract_val;
79 linidx += nocontract_val * m_nocontract_strides[0];
83 Index contract_val = left ? col : row;
84 for (
int i = static_cast<int>(array_size<contract_t>::value) - 1; i > 0; i--) {
85 const Index idx = contract_val / m_k_strides[i];
86 linidx += idx * m_contract_strides[i];
87 contract_val -= idx * m_k_strides[i];
90 if(array_size<contract_t>::value > 0) {
91 if (side == Rhs && inner_dim_contiguous) {
92 eigen_assert(m_contract_strides[0] == 1);
93 linidx += contract_val;
95 linidx += contract_val * m_contract_strides[0];
103 EIGEN_STRONG_INLINE IndexPair<Index> computeIndexPair(Index row, Index col,
const Index distance)
const {
104 const bool left = (side == Lhs);
105 Index nocontract_val[2] = {left ? row : col, left ? row + distance : col};
106 Index linidx[2] = {0, 0};
107 for (
int i = static_cast<int>(array_size<nocontract_t>::value) - 1; i > 0; i--) {
108 const Index idx0 = nocontract_val[0] / m_ij_strides[i];
109 const Index idx1 = nocontract_val[1] / m_ij_strides[i];
110 linidx[0] += idx0 * m_nocontract_strides[i];
111 linidx[1] += idx1 * m_nocontract_strides[i];
112 nocontract_val[0] -= idx0 * m_ij_strides[i];
113 nocontract_val[1] -= idx1 * m_ij_strides[i];
115 if (array_size<typename Tensor::Dimensions>::value > array_size<contract_t>::value) {
116 if (side == Lhs && inner_dim_contiguous) {
117 eigen_assert(m_nocontract_strides[0] == 1);
118 linidx[0] += nocontract_val[0];
119 linidx[1] += nocontract_val[1];
121 linidx[0] += nocontract_val[0] * m_nocontract_strides[0];
122 linidx[1] += nocontract_val[1] * m_nocontract_strides[0];
126 Index contract_val[2] = {left ? col : row, left ? col : row + distance};
127 for (
int i = static_cast<int>(array_size<contract_t>::value) - 1; i > 0; i--) {
128 const Index idx0 = contract_val[0] / m_k_strides[i];
129 const Index idx1 = contract_val[1] / m_k_strides[i];
130 linidx[0] += idx0 * m_contract_strides[i];
131 linidx[1] += idx1 * m_contract_strides[i];
132 contract_val[0] -= idx0 * m_k_strides[i];
133 contract_val[1] -= idx1 * m_k_strides[i];
136 if (side == Rhs && inner_dim_contiguous) {
137 eigen_assert(m_contract_strides[0] == 1);
138 linidx[0] += contract_val[0];
139 linidx[1] += contract_val[1];
141 linidx[0] += contract_val[0] * m_contract_strides[0];
142 linidx[1] += contract_val[1] * m_contract_strides[0];
144 return IndexPair<Index>(linidx[0], linidx[1]);
147 Index firstAligned(Index size)
const {
150 Index stride()
const {
155 const Tensor m_tensor;
156 const nocontract_t m_nocontract_strides;
157 const nocontract_t m_ij_strides;
158 const contract_t m_contract_strides;
159 const contract_t m_k_strides;
164 template<
typename Scalar,
typename Index,
int side,
166 typename nocontract_t,
typename contract_t,
168 bool inner_dim_contiguous,
bool inner_dim_reordered,
int Alignment>
169 class TensorContractionInputMapper;
171 template<
typename Scalar,
typename Index,
int side,
173 typename nocontract_t,
typename contract_t,
175 bool inner_dim_contiguous,
bool inner_dim_reordered,
int Alignment>
176 class TensorContractionSubMapper {
178 typedef typename packet_traits<Scalar>::type Packet;
179 typedef typename packet_traits<Scalar>::half HalfPacket;
181 typedef TensorContractionInputMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, packet_size, inner_dim_contiguous, inner_dim_reordered, Alignment> ParentMapper;
182 typedef TensorContractionSubMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, packet_size, inner_dim_contiguous, inner_dim_reordered, Alignment> Self;
183 typedef Self LinearMapper;
185 EIGEN_DEVICE_FUNC TensorContractionSubMapper(
const ParentMapper& base_mapper, Index vert_offset, Index horiz_offset)
186 : m_base_mapper(base_mapper), m_vert_offset(vert_offset), m_horiz_offset(horiz_offset) { }
188 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Scalar operator()(Index i)
const {
189 return m_base_mapper(i + m_vert_offset, m_horiz_offset);
191 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Scalar operator()(Index i, Index j)
const {
192 return m_base_mapper(i + m_vert_offset, j + m_horiz_offset);
195 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Packet loadPacket(Index i)
const {
196 return m_base_mapper.loadPacket(i + m_vert_offset, m_horiz_offset);
198 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Packet loadPacket(Index i, Index j)
const {
199 return m_base_mapper.loadPacket(i + m_vert_offset, j + m_horiz_offset);
202 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE HalfPacket loadHalfPacket(Index i)
const {
203 return m_base_mapper.loadHalfPacket(i + m_vert_offset, m_horiz_offset);
206 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE
void storePacket(Index i, Packet p)
const {
207 m_base_mapper.storePacket(i + m_vert_offset, m_horiz_offset, p);
210 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE LinearMapper getLinearMapper(Index i, Index j)
const {
211 return LinearMapper(m_base_mapper, i + m_vert_offset, j + m_horiz_offset);
214 template <
typename PacketT,
int AlignmentType>
215 EIGEN_ALWAYS_INLINE PacketT load(Index i)
const {
216 EIGEN_STATIC_ASSERT((internal::is_same<PacketT, Packet>::value), YOU_MADE_A_PROGRAMMING_MISTAKE);
217 EIGEN_STATIC_ASSERT((AlignmentType == Aligned || Alignment == Unaligned), YOU_MADE_A_PROGRAMMING_MISTAKE);
218 return loadPacket(i);
221 template <
typename Packet>
222 bool aligned(Index )
const {
227 const ParentMapper& m_base_mapper;
228 const Index m_vert_offset;
229 const Index m_horiz_offset;
233 template<
typename Scalar,
typename Index,
int side,
235 typename nocontract_t,
typename contract_t,
236 int packet_size = (Tensor::PacketAccess ? packet_traits<Scalar>::size : 1),
237 bool inner_dim_contiguous =
false,
bool inner_dim_reordered = (side != Lhs),
int Alignment=Unaligned>
238 class TensorContractionInputMapper
239 :
public BaseTensorContractionMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, packet_size, inner_dim_contiguous> {
242 typedef BaseTensorContractionMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, packet_size, inner_dim_contiguous> Base;
243 typedef TensorContractionSubMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, packet_size, inner_dim_contiguous, inner_dim_reordered, Alignment> SubMapper;
244 typedef SubMapper VectorMapper;
246 TensorContractionInputMapper(
const Tensor& tensor,
247 const nocontract_t& nocontract_strides,
248 const nocontract_t& ij_strides,
249 const contract_t& contract_strides,
250 const contract_t& k_strides)
251 : Base(tensor, nocontract_strides, ij_strides, contract_strides, k_strides) { }
254 EIGEN_STRONG_INLINE SubMapper getSubMapper(Index i, Index j)
const {
255 return SubMapper(*
this, i, j);
258 EIGEN_ALWAYS_INLINE VectorMapper getVectorMapper(Index i, Index j)
const {
259 return VectorMapper(*
this, i, j);
262 typedef typename packet_traits<Scalar>::type Packet;
263 typedef typename packet_traits<Scalar>::half HalfPacket;
266 EIGEN_STRONG_INLINE Packet loadPacket(Index i, Index j)
const {
271 EIGEN_STATIC_ASSERT(packet_size % 2 == 0, YOU_MADE_A_PROGRAMMING_MISTAKE);
273 if (Tensor::PacketAccess && inner_dim_contiguous && !inner_dim_reordered) {
274 const Index index = this->computeIndex(i, j);
275 eigen_assert(this->computeIndex(i+packet_size-1, j) == index + packet_size-1);
276 return this->m_tensor.template packet<Alignment>(index);
279 const IndexPair<Index> indexPair = this->computeIndexPair(i, j, packet_size - 1);
280 const Index first = indexPair.first;
281 const Index last = indexPair.second;
287 if (Tensor::PacketAccess &&
288 (side == Lhs || internal::array_size<contract_t>::value <= 1 || !inner_dim_reordered) &&
289 (last - first) == (packet_size - 1)) {
291 return this->m_tensor.template packet<Alignment>(first);
294 EIGEN_ALIGN_MAX Scalar data[packet_size];
296 data[0] = this->m_tensor.coeff(first);
297 for (Index k = 1; k < packet_size - 1; k += 2) {
298 const IndexPair<Index> internal_pair = this->computeIndexPair(i + k, j, 1);
299 data[k] = this->m_tensor.coeff(internal_pair.first);
300 data[k + 1] = this->m_tensor.coeff(internal_pair.second);
302 data[packet_size - 1] = this->m_tensor.coeff(last);
304 return pload<Packet>(data);
308 EIGEN_STRONG_INLINE HalfPacket loadHalfPacket(Index i, Index j)
const {
312 const Index half_packet_size = unpacket_traits<HalfPacket>::size;
313 if (half_packet_size == packet_size) {
314 return loadPacket(i, j);
316 EIGEN_ALIGN_MAX Scalar data[half_packet_size];
317 for (Index k = 0; k < half_packet_size; k++) {
318 data[k] = operator()(i + k, j);
320 return pload<HalfPacket>(data);
327 template<
typename Scalar,
typename Index,
int side,
329 typename nocontract_t,
typename contract_t,
330 bool inner_dim_contiguous,
bool inner_dim_reordered,
int Alignment>
331 class TensorContractionInputMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, 1, inner_dim_contiguous, inner_dim_reordered, Alignment>
332 :
public BaseTensorContractionMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, 1, inner_dim_contiguous> {
335 typedef BaseTensorContractionMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, 1, inner_dim_contiguous> Base;
336 typedef TensorContractionSubMapper<Scalar, Index, side, Tensor, nocontract_t, contract_t, 1, inner_dim_contiguous, inner_dim_reordered, Alignment> SubMapper;
337 typedef SubMapper VectorMapper;
339 TensorContractionInputMapper(
const Tensor& tensor,
340 const nocontract_t& nocontract_strides,
341 const nocontract_t& ij_strides,
342 const contract_t& contract_strides,
343 const contract_t& k_strides)
344 : Base(tensor, nocontract_strides, ij_strides, contract_strides, k_strides) { }
347 EIGEN_STRONG_INLINE SubMapper getSubMapper(Index i, Index j)
const {
348 return SubMapper(*
this, i, j);
351 EIGEN_ALWAYS_INLINE VectorMapper getVectorMapper(Index i, Index j)
const {
352 return VectorMapper(*
this, i, j);
355 typedef typename packet_traits<Scalar>::type Packet;
357 EIGEN_STRONG_INLINE Packet loadPacket(Index i, Index j)
const {
358 EIGEN_ALIGN_MAX Scalar data[1];
359 data[0] = this->m_tensor.coeff(this->computeIndex(i, j));
360 return pload<typename packet_traits<Scalar>::type>(data);
363 EIGEN_STRONG_INLINE Packet loadHalfPacket(Index i, Index j)
const {
364 return loadPacket(i, j);
369 template<
typename Dimensions,
typename LhsXprType,
typename RhsXprType>
370 struct traits<TensorContractionOp<Dimensions, LhsXprType, RhsXprType> >
373 typedef typename internal::promote_storage_type<
typename LhsXprType::Scalar,
374 typename RhsXprType::Scalar>::ret Scalar;
375 typedef typename internal::packet_traits<Scalar>::type Packet;
376 typedef typename promote_storage_type<typename traits<LhsXprType>::StorageKind,
377 typename traits<RhsXprType>::StorageKind>::ret StorageKind;
378 typedef typename promote_index_type<typename traits<LhsXprType>::Index,
379 typename traits<RhsXprType>::Index>::type Index;
380 typedef typename LhsXprType::Nested LhsNested;
381 typedef typename RhsXprType::Nested RhsNested;
382 typedef typename remove_reference<LhsNested>::type _LhsNested;
383 typedef typename remove_reference<RhsNested>::type _RhsNested;
386 static const int NumDimensions = max_n_1<traits<RhsXprType>::NumDimensions + traits<RhsXprType>::NumDimensions - 2 * array_size<Dimensions>::value>::size;
387 static const int Layout = traits<LhsXprType>::Layout;
394 template<
typename Dimensions,
typename LhsXprType,
typename RhsXprType>
395 struct eval<TensorContractionOp<Dimensions, LhsXprType, RhsXprType>,
Eigen::Dense>
397 typedef const TensorContractionOp<Dimensions, LhsXprType, RhsXprType>& type;
400 template<
typename Dimensions,
typename LhsXprType,
typename RhsXprType>
401 struct nested<TensorContractionOp<Dimensions, LhsXprType, RhsXprType>, 1, typename eval<TensorContractionOp<Dimensions, LhsXprType, RhsXprType> >::type>
403 typedef TensorContractionOp<Dimensions, LhsXprType, RhsXprType> type;
406 template<
typename Indices_,
typename LeftArgType_,
typename RightArgType_,
typename Device_>
407 struct traits<TensorEvaluator<const TensorContractionOp<Indices_, LeftArgType_, RightArgType_>, Device_> > {
408 typedef Indices_ Indices;
409 typedef LeftArgType_ LeftArgType;
410 typedef RightArgType_ RightArgType;
411 typedef Device_ Device;
414 static const int NumDimensions = max_n_1<traits<LeftArgType_>::NumDimensions + traits<RightArgType_>::NumDimensions - 2 * array_size<Indices_>::value>::size;
419 template<
typename Indices,
typename LhsXprType,
typename RhsXprType>
420 class TensorContractionOp :
public TensorBase<TensorContractionOp<Indices, LhsXprType, RhsXprType>, ReadOnlyAccessors>
423 typedef typename Eigen::internal::traits<TensorContractionOp>::Scalar Scalar;
424 typedef typename Eigen::internal::traits<TensorContractionOp>::Packet Packet;
425 typedef typename internal::promote_storage_type<
typename LhsXprType::CoeffReturnType,
426 typename RhsXprType::CoeffReturnType>::ret CoeffReturnType;
427 typedef typename internal::promote_storage_type<
typename LhsXprType::PacketReturnType,
428 typename RhsXprType::PacketReturnType>::ret PacketReturnType;
429 typedef typename Eigen::internal::nested<TensorContractionOp>::type Nested;
430 typedef typename Eigen::internal::traits<TensorContractionOp>::StorageKind StorageKind;
431 typedef typename Eigen::internal::traits<TensorContractionOp>::Index Index;
433 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorContractionOp(
434 const LhsXprType& lhs,
const RhsXprType& rhs,
const Indices& dims)
435 : m_lhs_xpr(lhs), m_rhs_xpr(rhs), m_indices(dims) {}
438 const Indices& indices()
const {
return m_indices; }
442 const typename internal::remove_all<typename LhsXprType::Nested>::type&
443 lhsExpression()
const {
return m_lhs_xpr; }
446 const typename internal::remove_all<typename RhsXprType::Nested>::type&
447 rhsExpression()
const {
return m_rhs_xpr; }
450 typename LhsXprType::Nested m_lhs_xpr;
451 typename RhsXprType::Nested m_rhs_xpr;
452 const Indices m_indices;
456 template<
typename Derived>
457 struct TensorContractionEvaluatorBase
459 typedef typename internal::traits<Derived>::Indices Indices;
460 typedef typename internal::traits<Derived>::LeftArgType LeftArgType;
461 typedef typename internal::traits<Derived>::RightArgType RightArgType;
462 typedef typename internal::traits<Derived>::Device Device;
464 typedef TensorContractionOp<Indices, LeftArgType, RightArgType> XprType;
465 typedef typename internal::remove_const<typename XprType::Scalar>::type Scalar;
466 typedef typename XprType::Packet Packet;
467 typedef typename XprType::Index Index;
468 typedef typename XprType::CoeffReturnType CoeffReturnType;
469 typedef typename XprType::PacketReturnType PacketReturnType;
473 PacketAccess = (internal::packet_traits<Scalar>::size > 1),
474 Layout = TensorEvaluator<LeftArgType, Device>::Layout,
482 typedef typename internal::conditional<
483 static_cast<int>(Layout) == static_cast<int>(ColMajor), LeftArgType, RightArgType>::type EvalLeftArgType;
484 typedef typename internal::conditional<
485 static_cast<int>(Layout) == static_cast<int>(ColMajor), RightArgType, LeftArgType>::type EvalRightArgType;
487 static const int LDims =
488 internal::array_size<typename TensorEvaluator<EvalLeftArgType, Device>::Dimensions>::value;
489 static const int RDims =
490 internal::array_size<typename TensorEvaluator<EvalRightArgType, Device>::Dimensions>::value;
491 static const unsigned int ContractDims = internal::array_size<Indices>::value;
492 static const int NumDims = max_n_1<LDims + RDims - 2 * ContractDims>::size;
494 typedef array<Index, LDims> left_dim_mapper_t;
495 typedef array<Index, RDims> right_dim_mapper_t;
496 typedef array<Index, ContractDims> contract_t;
497 typedef array<Index, max_n_1<LDims - ContractDims>::size> left_nocontract_t;
498 typedef array<Index, max_n_1<RDims - ContractDims>::size> right_nocontract_t;
500 typedef DSizes<Index, NumDims> Dimensions;
502 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
503 TensorContractionEvaluatorBase(
const XprType& op,
const Device& device)
504 : m_leftImpl(choose(Cond<static_cast<int>(Layout) == static_cast<int>(ColMajor)>(),
505 op.lhsExpression(), op.rhsExpression()), device),
506 m_rightImpl(choose(Cond<static_cast<int>(Layout) == static_cast<int>(ColMajor)>(),
507 op.rhsExpression(), op.lhsExpression()), device),
510 EIGEN_STATIC_ASSERT((static_cast<int>(TensorEvaluator<LeftArgType, Device>::Layout) ==
511 static_cast<int>(TensorEvaluator<RightArgType, Device>::Layout)),
512 YOU_MADE_A_PROGRAMMING_MISTAKE);
515 DSizes<Index, LDims> eval_left_dims;
516 DSizes<Index, RDims> eval_right_dims;
517 array<IndexPair<Index>, ContractDims> eval_op_indices;
518 if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
520 for (
int i = 0; i < LDims; i++) {
521 eval_left_dims[i] = m_leftImpl.dimensions()[i];
523 for (
int i = 0; i < RDims; i++) {
524 eval_right_dims[i] = m_rightImpl.dimensions()[i];
527 for (
unsigned int i = 0; i < ContractDims; i++) {
528 eval_op_indices[i].first = op.indices()[i].first;
529 eval_op_indices[i].second = op.indices()[i].second;
533 for (
int i = 0; i < LDims; i++) {
534 eval_left_dims[i] = m_leftImpl.dimensions()[LDims - i - 1];
536 for (
int i = 0; i < RDims; i++) {
537 eval_right_dims[i] = m_rightImpl.dimensions()[RDims - i - 1];
541 for (
unsigned int i = 0; i < ContractDims; i++) {
542 eval_op_indices[i].first = LDims - 1 - op.indices()[i].second;
543 eval_op_indices[i].second = RDims - 1 - op.indices()[i].first;
547 array<Index, LDims> lhs_strides;
549 for (
int i = 0; i < LDims-1; ++i) {
550 lhs_strides[i+1] = lhs_strides[i] * eval_left_dims[i];
553 array<Index, RDims> rhs_strides;
555 for (
int i = 0; i < RDims-1; ++i) {
556 rhs_strides[i+1] = rhs_strides[i] * eval_right_dims[i];
573 m_lhs_inner_dim_contiguous =
true;
575 unsigned int nocontract_idx = 0;
577 for (
int i = 0; i < LDims; i++) {
579 bool contracting =
false;
580 for (
unsigned int j = 0; j < ContractDims; j++) {
581 if (eval_op_indices[j].first == i) {
588 m_dimensions[dim_idx] = eval_left_dims[i];
589 m_left_nocontract_strides[nocontract_idx] = lhs_strides[i];
591 m_lhs_inner_dim_contiguous =
false;
593 if (nocontract_idx+1 < internal::array_size<left_nocontract_t>::value) {
594 m_i_strides[nocontract_idx+1] =
595 m_i_strides[nocontract_idx] * eval_left_dims[i];
597 m_i_size = m_i_strides[nocontract_idx] * eval_left_dims[i];
605 for (
int i = 0; i < RDims; i++) {
606 bool contracting =
false;
608 for (
unsigned int j = 0; j < ContractDims; j++) {
609 if (eval_op_indices[j].second == i) {
615 m_dimensions[dim_idx] = eval_right_dims[i];
616 if (nocontract_idx+1 < internal::array_size<right_nocontract_t>::value) {
617 m_j_strides[nocontract_idx+1] =
618 m_j_strides[nocontract_idx] * eval_right_dims[i];
620 m_j_size = m_j_strides[nocontract_idx] * eval_right_dims[i];
622 m_right_nocontract_strides[nocontract_idx] = rhs_strides[i];
633 m_rhs_inner_dim_contiguous =
true;
634 m_rhs_inner_dim_reordered =
false;
635 for (
unsigned int i = 0; i < ContractDims; i++) {
636 Index left = eval_op_indices[i].first;
637 Index right = eval_op_indices[i].second;
639 Index size = eval_left_dims[left];
640 eigen_assert(size == eval_right_dims[right] &&
641 "Contraction axes must be same size");
643 if (i+1 < internal::array_size<contract_t>::value) {
644 m_k_strides[i+1] = m_k_strides[i] * size;
646 m_k_size = m_k_strides[i] * size;
648 m_left_contracting_strides[i] = lhs_strides[left];
649 m_right_contracting_strides[i] = rhs_strides[right];
651 if (i > 0 && right < eval_op_indices[i-1].second) {
652 m_rhs_inner_dim_reordered =
true;
655 m_rhs_inner_dim_contiguous =
false;
660 if (LDims + RDims == 2 * ContractDims) {
665 if (static_cast<int>(Layout) == static_cast<int>(RowMajor)) {
666 for (
int i = 0, j = NumDims - 1; i < j; i++, j--) {
667 numext::swap(m_dimensions[i], m_dimensions[j]);
672 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
const Dimensions& dimensions()
const {
return m_dimensions; }
674 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
bool evalSubExprsIfNeeded(Scalar* data) {
675 m_leftImpl.evalSubExprsIfNeeded(NULL);
676 m_rightImpl.evalSubExprsIfNeeded(NULL);
681 m_result =
static_cast<Scalar *
>(m_device.allocate(dimensions().TotalSize() *
sizeof(Scalar)));
687 EIGEN_DEVICE_FUNC
void evalTo(Scalar* buffer)
const {
688 if (this->m_lhs_inner_dim_contiguous) {
689 if (this->m_rhs_inner_dim_contiguous) {
690 if (this->m_rhs_inner_dim_reordered) {
691 static_cast<const Derived*
>(
this)->
template evalProduct<true, true, true, Unaligned>(buffer);
694 static_cast<const Derived*
>(
this)->
template evalProduct<true, true, false, Unaligned>(buffer);
698 if (this->m_rhs_inner_dim_reordered) {
699 static_cast<const Derived*
>(
this)->
template evalProduct<true, false, true, Unaligned>(buffer);
702 static_cast<const Derived*
>(
this)->
template evalProduct<true, false, false, Unaligned>(buffer);
707 if (this->m_rhs_inner_dim_contiguous) {
708 if (this->m_rhs_inner_dim_reordered) {
709 static_cast<const Derived*
>(
this)->
template evalProduct<false, true, true, Unaligned>(buffer);
712 static_cast<const Derived*
>(
this)->
template evalProduct<false, true, false, Unaligned>(buffer);
716 if (this->m_rhs_inner_dim_reordered) {
717 static_cast<const Derived*
>(
this)->
template evalProduct<false, false, true, Unaligned>(buffer);
720 static_cast<const Derived*
>(
this)->
template evalProduct<false, false, false, Unaligned>(buffer);
726 template <
bool lhs_inner_dim_contiguous,
bool rhs_inner_dim_contiguous,
bool rhs_inner_dim_reordered,
int Alignment>
727 void evalGemv(Scalar* buffer)
const {
728 const Index rows = m_i_size;
729 const Index cols = m_k_size;
731 typedef typename internal::remove_const<typename EvalLeftArgType::Scalar>::type LhsScalar;
732 typedef typename internal::remove_const<typename EvalRightArgType::Scalar>::type RhsScalar;
733 typedef TensorEvaluator<EvalLeftArgType, Device> LeftEvaluator;
734 typedef TensorEvaluator<EvalRightArgType, Device> RightEvaluator;
735 const Index lhs_packet_size = internal::packet_traits<LhsScalar>::size;
736 const Index rhs_packet_size = internal::packet_traits<RhsScalar>::size;
737 typedef internal::TensorContractionInputMapper<LhsScalar, Index, internal::Lhs,
738 LeftEvaluator, left_nocontract_t,
739 contract_t, lhs_packet_size,
740 lhs_inner_dim_contiguous,
741 false, Unaligned> LhsMapper;
743 typedef internal::TensorContractionInputMapper<RhsScalar, Index, internal::Rhs,
744 RightEvaluator, right_nocontract_t,
745 contract_t, rhs_packet_size,
746 rhs_inner_dim_contiguous,
747 rhs_inner_dim_reordered, Unaligned> RhsMapper;
749 LhsMapper lhs(m_leftImpl, m_left_nocontract_strides, m_i_strides,
750 m_left_contracting_strides, m_k_strides);
751 RhsMapper rhs(m_rightImpl, m_right_nocontract_strides, m_j_strides,
752 m_right_contracting_strides, m_k_strides);
754 const Scalar alpha(1);
755 const Index resIncr(1);
758 m_device.memset(buffer, 0, rows *
sizeof(Scalar));
760 internal::general_matrix_vector_product<Index,LhsScalar,LhsMapper,ColMajor,false,RhsScalar,RhsMapper,false>::run(
761 rows, cols, lhs, rhs,
762 buffer, resIncr, alpha);
765 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
void cleanup() {
766 m_leftImpl.cleanup();
767 m_rightImpl.cleanup();
769 if (m_result != NULL) {
770 m_device.deallocate(m_result);
775 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index)
const {
776 return m_result[index];
779 template<
int LoadMode>
780 EIGEN_DEVICE_FUNC PacketReturnType packet(Index index)
const {
781 return internal::ploadt<Packet, LoadMode>(m_result + index);
784 EIGEN_DEVICE_FUNC Scalar* data()
const {
return NULL; }
788 TensorContractionEvaluatorBase& operator = (
const TensorContractionEvaluatorBase&);
789 Dimensions m_dimensions;
791 contract_t m_k_strides;
792 contract_t m_left_contracting_strides;
793 contract_t m_right_contracting_strides;
795 bool m_lhs_inner_dim_contiguous;
796 bool m_rhs_inner_dim_contiguous;
797 bool m_rhs_inner_dim_reordered;
799 left_nocontract_t m_i_strides;
800 right_nocontract_t m_j_strides;
801 left_nocontract_t m_left_nocontract_strides;
802 right_nocontract_t m_right_nocontract_strides;
808 TensorEvaluator<EvalLeftArgType, Device> m_leftImpl;
809 TensorEvaluator<EvalRightArgType, Device> m_rightImpl;
810 const Device& m_device;
816 template<
typename Indices,
typename LeftArgType,
typename RightArgType,
typename Device>
817 struct TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgType>, Device> :
818 public TensorContractionEvaluatorBase<
819 TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgType>, Device> > {
820 typedef TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgType>, Device> Self;
821 typedef TensorContractionEvaluatorBase<Self> Base;
823 typedef TensorContractionOp<Indices, LeftArgType, RightArgType> XprType;
824 typedef typename internal::remove_const<typename XprType::Scalar>::type Scalar;
825 typedef typename XprType::Packet Packet;
826 typedef typename XprType::Index Index;
827 typedef typename XprType::CoeffReturnType CoeffReturnType;
828 typedef typename XprType::PacketReturnType PacketReturnType;
831 Layout = TensorEvaluator<LeftArgType, Device>::Layout,
838 typedef typename internal::conditional<
839 static_cast<int>(Layout) == static_cast<int>(ColMajor), LeftArgType, RightArgType>::type EvalLeftArgType;
840 typedef typename internal::conditional<
841 static_cast<int>(Layout) == static_cast<int>(ColMajor), RightArgType, LeftArgType>::type EvalRightArgType;
843 static const int LDims =
844 internal::array_size<typename TensorEvaluator<EvalLeftArgType, Device>::Dimensions>::value;
845 static const int RDims =
846 internal::array_size<typename TensorEvaluator<EvalRightArgType, Device>::Dimensions>::value;
847 static const int ContractDims = internal::array_size<Indices>::value;
849 typedef array<Index, LDims> left_dim_mapper_t;
850 typedef array<Index, RDims> right_dim_mapper_t;
852 typedef array<Index, ContractDims> contract_t;
853 typedef array<Index, max_n_1<LDims - ContractDims>::size> left_nocontract_t;
854 typedef array<Index, max_n_1<RDims - ContractDims>::size> right_nocontract_t;
856 static const int NumDims = max_n_1<LDims + RDims - 2 * ContractDims>::size;
859 typedef DSizes<Index, NumDims> Dimensions;
862 EIGEN_DEVICE_FUNC TensorEvaluator(
const XprType& op,
const Device& device) :
865 template <
bool lhs_inner_dim_contiguous,
bool rhs_inner_dim_contiguous,
bool rhs_inner_dim_reordered,
int Alignment>
866 void evalProduct(Scalar* buffer)
const {
867 if (this->m_j_size == 1) {
868 this->
template evalGemv<lhs_inner_dim_contiguous, rhs_inner_dim_contiguous, rhs_inner_dim_reordered, Alignment>(buffer);
872 evalGemm<lhs_inner_dim_contiguous, rhs_inner_dim_contiguous, rhs_inner_dim_reordered, Alignment>(buffer);
875 template <
bool lhs_inner_dim_contiguous,
bool rhs_inner_dim_contiguous,
bool rhs_inner_dim_reordered,
int Alignment>
876 EIGEN_DEVICE_FUNC
void evalGemm(Scalar* buffer)
const {
878 const Index k = this->m_k_size;
881 const Index m = this->m_i_size;
884 const Index n = this->m_j_size;
887 this->m_device.memset(buffer, 0, m * n *
sizeof(Scalar));
890 typedef typename internal::remove_const<typename EvalLeftArgType::Scalar>::type LhsScalar;
891 typedef typename internal::remove_const<typename EvalRightArgType::Scalar>::type RhsScalar;
892 typedef typename internal::gebp_traits<LhsScalar, RhsScalar> Traits;
894 const Index nr = Traits::nr;
895 const Index mr = Traits::mr;
897 typedef TensorEvaluator<EvalLeftArgType, Device> LeftEvaluator;
898 typedef TensorEvaluator<EvalRightArgType, Device> RightEvaluator;
900 const Index lhs_packet_size = internal::packet_traits<LhsScalar>::size;
901 const Index rhs_packet_size = internal::packet_traits<RhsScalar>::size;
903 typedef internal::TensorContractionInputMapper<LhsScalar, Index, internal::Lhs,
904 LeftEvaluator, left_nocontract_t,
905 contract_t, lhs_packet_size,
906 lhs_inner_dim_contiguous,
907 false, Unaligned> LhsMapper;
909 typedef internal::TensorContractionInputMapper<RhsScalar, Index, internal::Rhs,
910 RightEvaluator, right_nocontract_t,
911 contract_t, rhs_packet_size,
912 rhs_inner_dim_contiguous,
913 rhs_inner_dim_reordered, Unaligned> RhsMapper;
915 typedef internal::blas_data_mapper<Scalar, Index, ColMajor> OutputMapper;
918 internal::gemm_pack_lhs<LhsScalar, Index, typename LhsMapper::SubMapper, mr, Traits::LhsProgress, ColMajor> pack_lhs;
919 internal::gemm_pack_rhs<RhsScalar, Index, typename RhsMapper::SubMapper, nr, ColMajor> pack_rhs;
921 internal::gebp_kernel<LhsScalar, RhsScalar, Index, OutputMapper, mr, nr, false, false> gebp;
924 LhsMapper lhs(this->m_leftImpl, this->m_left_nocontract_strides, this->m_i_strides,
925 this->m_left_contracting_strides, this->m_k_strides);
927 RhsMapper rhs(this->m_rightImpl, this->m_right_nocontract_strides, this->m_j_strides,
928 this->m_right_contracting_strides, this->m_k_strides);
930 OutputMapper output(buffer, m);
932 typedef typename internal::gemm_blocking_space<ColMajor, LhsScalar, RhsScalar, Dynamic, Dynamic, Dynamic> BlockingType;
935 BlockingType blocking(m, n, k, 1,
true);
936 const Index kc = blocking.kc();
937 const Index mc = numext::mini(m, blocking.mc());
938 const Index nc = numext::mini(n, blocking.nc());
939 const Index sizeA = mc * kc;
940 const Index sizeB = kc * nc;
942 LhsScalar* blockA =
static_cast<LhsScalar *
>(this->m_device.allocate(sizeA *
sizeof(LhsScalar)));
943 RhsScalar* blockB =
static_cast<RhsScalar *
>(this->m_device.allocate(sizeB *
sizeof(RhsScalar)));
945 for(Index i2=0; i2<m; i2+=mc)
947 const Index actual_mc = numext::mini(i2+mc,m)-i2;
948 for (Index k2 = 0; k2 < k; k2 += kc) {
950 const Index actual_kc = numext::mini(k2 + kc, k) - k2;
951 pack_lhs(blockA, lhs.getSubMapper(i2, k2), actual_kc, actual_mc, 0, 0);
954 for (Index j2 = 0; j2 < n; j2 += nc) {
956 const Index actual_nc = numext::mini(j2 + nc, n) - j2;
957 pack_rhs(blockB, rhs.getSubMapper(k2, j2), actual_kc, actual_nc, 0, 0);
961 gebp(output.getSubMapper(i2, j2), blockA, blockB, actual_mc, actual_kc, actual_nc, 1.0, -1, -1, 0, 0);
966 this->m_device.deallocate(blockA);
967 this->m_device.deallocate(blockB);
973 #endif // EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_H
Namespace containing all symbols from the Eigen library.
Definition: CXX11Meta.h:13