10 #ifndef EIGEN_CXX11_TENSOR_TENSOR_CONCATENATION_H
11 #define EIGEN_CXX11_TENSOR_TENSOR_CONCATENATION_H
23 template<
typename Axis,
typename LhsXprType,
typename RhsXprType>
44 template<
typename Axis,
typename LhsXprType,
typename RhsXprType>
50 template<
typename Axis,
typename LhsXprType,
typename RhsXprType>
59 template<
typename Axis,
typename LhsXprType,
typename RhsXprType>
94 template<
typename Axis,
typename LeftArgType,
typename RightArgType,
typename Device>
123 : m_leftImpl(op.lhsExpression(), device), m_rightImpl(op.rhsExpression(), device), m_axis(op.axis())
130 const Dimensions& lhs_dims = m_leftImpl.dimensions();
131 const Dimensions& rhs_dims = m_rightImpl.dimensions();
134 for (;
i < m_axis; ++
i) {
137 m_dimensions[
i] = lhs_dims[
i];
141 m_dimensions[
i] = lhs_dims[
i] + rhs_dims[
i];
142 for (++
i;
i < NumDims; ++
i) {
145 m_dimensions[
i] = lhs_dims[
i];
150 m_leftStrides[0] = 1;
151 m_rightStrides[0] = 1;
152 m_outputStrides[0] = 1;
154 for (
int j = 1;
j < NumDims; ++
j) {
155 m_leftStrides[
j] = m_leftStrides[
j-1] * lhs_dims[
j-1];
156 m_rightStrides[
j] = m_rightStrides[
j-1] * rhs_dims[
j-1];
157 m_outputStrides[
j] = m_outputStrides[
j-1] * m_dimensions[
j-1];
160 m_leftStrides[NumDims - 1] = 1;
161 m_rightStrides[NumDims - 1] = 1;
162 m_outputStrides[NumDims - 1] = 1;
164 for (
int j = NumDims - 2;
j >= 0; --
j) {
165 m_leftStrides[
j] = m_leftStrides[
j+1] * lhs_dims[
j+1];
166 m_rightStrides[
j] = m_rightStrides[
j+1] * rhs_dims[
j+1];
167 m_outputStrides[
j] = m_outputStrides[
j+1] * m_dimensions[
j+1];
177 m_leftImpl.evalSubExprsIfNeeded(
NULL);
178 m_rightImpl.evalSubExprsIfNeeded(
NULL);
184 m_leftImpl.cleanup();
185 m_rightImpl.cleanup();
195 for (
int i = NumDims - 1;
i > 0; --
i) {
196 subs[
i] = index / m_outputStrides[
i];
197 index -= subs[
i] * m_outputStrides[
i];
201 for (
int i = 0;
i < NumDims - 1; ++
i) {
202 subs[
i] = index / m_outputStrides[
i];
203 index -= subs[
i] * m_outputStrides[
i];
205 subs[NumDims - 1] = index;
208 const Dimensions& left_dims = m_leftImpl.dimensions();
209 if (subs[m_axis] < left_dims[m_axis]) {
212 left_index = subs[0];
214 for (
int i = 1;
i < NumDims; ++
i) {
215 left_index += (subs[
i] % left_dims[
i]) * m_leftStrides[
i];
218 left_index = subs[NumDims - 1];
220 for (
int i = NumDims - 2;
i >= 0; --
i) {
221 left_index += (subs[
i] % left_dims[
i]) * m_leftStrides[
i];
224 return m_leftImpl.coeff(left_index);
226 subs[m_axis] -= left_dims[m_axis];
227 const Dimensions& right_dims = m_rightImpl.dimensions();
230 right_index = subs[0];
232 for (
int i = 1;
i < NumDims; ++
i) {
233 right_index += (subs[
i] % right_dims[
i]) * m_rightStrides[
i];
236 right_index = subs[NumDims - 1];
238 for (
int i = NumDims - 2;
i >= 0; --
i) {
239 right_index += (subs[
i] % right_dims[
i]) * m_rightStrides[
i];
242 return m_rightImpl.coeff(right_index);
247 template<
int LoadMode>
256 for (
int i = 0;
i < packetSize; ++
i) {
265 const double compute_cost = NumDims * (2 * TensorOpCost::AddCost<Index>() +
266 2 * TensorOpCost::MulCost<Index>() +
267 TensorOpCost::DivCost<Index>() +
268 TensorOpCost::ModCost<Index>());
269 const double lhs_size = m_leftImpl.dimensions().TotalSize();
270 const double rhs_size = m_rightImpl.dimensions().TotalSize();
271 return (lhs_size / (lhs_size + rhs_size)) *
272 m_leftImpl.costPerCoeff(vectorized) +
273 (rhs_size / (lhs_size + rhs_size)) *
274 m_rightImpl.costPerCoeff(vectorized) +
280 #ifdef EIGEN_USE_SYCL
283 m_leftImpl.bind(cgh);
284 m_rightImpl.bind(cgh);
299 template<
typename Axis,
typename LeftArgType,
typename RightArgType,
typename Device>
301 :
public TensorEvaluator<const TensorConcatenationOp<Axis, LeftArgType, RightArgType>, Device>
336 for (
int i = Base::NumDims - 1;
i > 0; --
i) {
337 subs[
i] = index / this->m_outputStrides[
i];
338 index -= subs[
i] * this->m_outputStrides[
i];
342 const Dimensions& left_dims = this->m_leftImpl.dimensions();
343 if (subs[this->m_axis] < left_dims[this->m_axis]) {
344 Index left_index = subs[0];
345 for (
int i = 1;
i < Base::NumDims; ++
i) {
346 left_index += (subs[
i] % left_dims[
i]) * this->m_leftStrides[
i];
348 return this->m_leftImpl.coeffRef(left_index);
350 subs[this->m_axis] -= left_dims[this->m_axis];
351 const Dimensions& right_dims = this->m_rightImpl.dimensions();
352 Index right_index = subs[0];
353 for (
int i = 1;
i < Base::NumDims; ++
i) {
354 right_index += (subs[
i] % right_dims[
i]) * this->m_rightStrides[
i];
356 return this->m_rightImpl.coeffRef(right_index);
368 internal::pstore<CoeffReturnType, PacketReturnType>(
values,
x);
369 for (
int i = 0;
i < packetSize; ++
i) {
377 #endif // EIGEN_CXX11_TENSOR_TENSOR_CONCATENATION_H