198 parseEquation(equation);
202 void operator()(std::string &equation,
const std::vector<Shape> &input_shapes,
211 const int num_inputs = input_shapes.size();
212 std::vector<InputTensor<float>> inputs(num_inputs);
213 for (
int i = 0; i < num_inputs; i++)
215 inputs[i].shape.ReplaceWith(input_shapes[i].DimensionsCount(), input_shapes[i].DimsData());
216 inputs[i].buffer = input_data[i];
220 Labels output_labels(_output_labels);
221 std::vector<DimensionType> label_types(_label_types);
223 LabelCounts output_label_counts(_output_label_counts);
226 processDimensions(inputs, &input_labels, &output_labels, &label_types, &input_label_counts,
227 &output_label_counts, &label_to_dim_sizes);
236 std::vector<Tensor> inputs_reduced(num_inputs);
237 std::vector<bool> swap_free_and_contract(num_inputs);
238 for (
int i = 0; i < num_inputs; ++i)
240 bool temp_swap_free_and_contract =
false;
241 reduceOperand<float>(inputs[i], label_types, input_label_counts[i], &input_labels[i],
242 &free_labels[i], &temp_swap_free_and_contract, &inputs_reduced[i]);
243 swap_free_and_contract[i] = temp_swap_free_and_contract;
249 Tensor contraction_output_reshaped;
250 contractOperands(inputs_reduced, swap_free_and_contract, &contraction_output_reshaped);
254 std::vector<int32_t> result_shape_dims(contraction_output_reshaped.
shape.
DimensionsCount() - 2);
256 for (
size_t i = 0; i < result_shape_dims.size(); i++)
258 result_shape_dims[i] = contraction_output_reshaped.
shape.
Dims(i);
261 int num_labels = label_types.size();
265 for (
int label = 0; label < num_labels; ++label)
268 result_labels.push_back(label);
270 for (
int label = 0; label < num_labels; ++label)
272 if (label_types[label] ==
kBatch)
273 result_labels.push_back(label);
275 for (
int i = 0; i < num_inputs; ++i)
277 for (
auto &&label : free_labels[i])
279 result_labels.push_back(label);
280 result_shape_dims.push_back(label_to_dim_sizes[label]);
284 Shape result_shape(result_shape_dims.size(), result_shape_dims.data());
288 Tensor contraction_output;
289 copyFrom(contraction_output_reshaped, result_shape, &contraction_output);
296 strideOrInflate<float>(contraction_output, result_labels, output_label_counts,
297 true , &output_inflated);
303 for (
auto &&label : result_labels)
305 inflated_labels.insert(inflated_labels.end(), output_label_counts[label], label);
307 result_labels.swap(inflated_labels);
316 std::vector<int32_t> output_permutation(output_labels.size());
317 std::vector<int32_t> label_to_position(num_labels, -1);
318 for (
size_t i = 0; i < result_labels.size(); ++i)
321 if (label_to_position[result_labels[i]] == -1)
323 label_to_position[result_labels[i]] = i;
326 for (
size_t i = 0; i < output_labels.size(); ++i)
328 output_permutation[i] = label_to_position[output_labels[i]];
330 label_to_position[output_labels[i]] += 1;
336 temp_inflated.
buffer = (
reinterpret_cast<const float *
>(output_inflated.
buffer));
340 transposeOperand<float>(temp_inflated, output_permutation, &output);
342 memcpy(output_data, output.buffer,
output_shape.FlatSize() *
sizeof(
float));
344 temp_operand.clear();
348 void parseEquation(std::string &equation)
350 std::vector<std::string> input_str;
351 std::string output_str;
353 parseEinsumEquation(equation, input_str, output_str);
357 std::map<char, int> label_mapping;
358 int num_inputs = input_str.size();
359 _input_labels.resize(num_inputs);
362 for (
int i = 0; i < num_inputs; ++i)
364 mapToLabels(input_str[i], _input_labels.at(i), label_mapping);
366 mapToLabels(output_str, _output_labels, label_mapping);
369 int num_labels = label_mapping.size();
370 _input_label_counts.resize(num_inputs);
371 _input_has_ellipsis.resize(num_inputs);
372 for (
int i = 0; i < num_inputs; ++i)
374 _input_label_counts.at(i).resize(num_labels);
375 for (
const int label : _input_labels.at(i))
377 if (label != kEllipsisLabel)
378 _input_label_counts.at(i)[label] += 1;
380 _input_has_ellipsis.at(i) =
true;
383 _output_label_counts.resize(num_labels);
384 for (
const int label : _output_labels)
386 if (label != kEllipsisLabel)
387 _output_label_counts.at(label) += 1;
389 _output_has_ellipsis =
true;
393 _label_types.resize(num_labels);
394 for (
int label = 0; label < num_labels; ++label)
396 bool removed = (_output_label_counts[label] == 0);
398 num_inputs == 1 || _input_label_counts[0][label] == 0 || _input_label_counts[1][label] == 0;
399 _label_types[label] = getDimensionType(removed, unique);
403 void parseEinsumEquation(
const std::string &equation, std::vector<std::string> &input_subscripts,
404 std::string &output_subscript)
406 std::vector<std::string> inputs_and_output_subscripts = strSplit(equation,
"->");
407 if (inputs_and_output_subscripts.size() != 2)
409 throw std::runtime_error{
"Einsum: Expecting exactly one '->' in einsum equation: " +
413 output_subscript = inputs_and_output_subscripts[1];
414 input_subscripts = strSplit(inputs_and_output_subscripts[0],
",");
415 if (input_subscripts.size() != 1 && input_subscripts.size() != 2)
417 throw std::runtime_error{
"Einsum: Expecting 1 or 2 input subscripts in equation '" +
418 equation +
"' but got: " + std::to_string(input_subscripts.size())};
423 void mapToLabels(
const std::string &subscript,
Labels &labels, std::map<char, int> &label_mapping)
425 for (
size_t i = 0; i < subscript.size(); ++i)
427 const char label_char = subscript[i];
428 if (label_char ==
'.')
430 labels.push_back(kEllipsisLabel);
434 if (label_mapping.find(label_char) == label_mapping.end())
436 const int next_label = label_mapping.size();
437 label_mapping[label_char] = next_label;
439 const int mapped_label = label_mapping[label_char];
440 labels.push_back(mapped_label);
444 template <
typename T>
445 void processDimensions(
const std::vector<InputTensor<T>> &inputs,
OperandLabels *input_labels,
446 Labels *output_labels, std::vector<DimensionType> *label_types,
450 if (
inputs.size() != input_labels->size())
452 throw std::runtime_error{
"Expected " + std::to_string(input_labels->size()) +
453 " inputs but got: " + std::to_string(
inputs.size())};
459 int max_bcast_dims = 0;
460 const int num_named_labels = label_types->size();
461 label_to_dim_sizes->resize(num_named_labels);
464 Labels *labels = &(*input_labels)[i];
466 if (!_input_has_ellipsis[i])
468 if (inputs[i].shape.DimensionsCount() != ((int32_t)labels->size()))
470 throw std::runtime_error{
"Expected input " + std::to_string(i) +
" to have rank " +
471 std::to_string(labels->size()) +
" but got: " +
472 std::to_string(inputs[i].shape.DimensionsCount())};
474 for (
size_t label_idx = 0; label_idx < labels->size(); ++label_idx)
476 const int label = (*labels)[label_idx];
477 recordLabelToDimension(label, label_idx, inputs[i].shape, label_to_dim_sizes);
483 if (inputs[i].shape.DimensionsCount() + 1 < (int32_t)labels->size())
485 throw std::runtime_error{
"Expected input " + std::to_string(i) +
" to have rank at least " +
486 std::to_string(labels->size() - 1) +
487 " but got: " + std::to_string(inputs[i].shape.DimensionsCount())};
489 int ellipsis_axis = -1;
490 const int num_bcast_dims =
inputs[i].shape.DimensionsCount() - labels->size() + 1;
491 for (
size_t label_idx = 0; label_idx < labels->size(); ++label_idx)
493 const int label = (*labels)[label_idx];
494 if (label == kEllipsisLabel)
496 ellipsis_axis = label_idx;
500 const int axis = label_idx + (ellipsis_axis == -1 ? 0 : num_bcast_dims - 1);
501 recordLabelToDimension(label, axis, inputs[i].shape, label_to_dim_sizes);
505 if (ellipsis_axis != -1)
507 insertBroadcastLabels(num_bcast_dims, num_named_labels, ellipsis_axis, labels,
508 &input_label_counts->at(i));
509 max_bcast_dims = std::max(max_bcast_dims, num_bcast_dims);
513 std::vector<bool>::iterator it_input =
514 std::find(_input_has_ellipsis.begin(), _input_has_ellipsis.end(),
true);
515 if (it_input == _input_has_ellipsis.end() && !_output_has_ellipsis)
520 auto it = std::find(output_labels->begin(), output_labels->end(), kEllipsisLabel);
521 if (it != output_labels->end())
523 const int ellipsis_axis = it - output_labels->begin();
524 insertBroadcastLabels(max_bcast_dims, num_named_labels, ellipsis_axis, output_labels,
525 output_label_counts);
527 else if (max_bcast_dims > 0)
529 std::runtime_error{
"Output contains " + std::to_string(max_bcast_dims) +
530 " broadcasting dimension(s) but no ellipsis " +
531 "(...) was found in the output subscripts."};
534 label_types->resize(num_named_labels + max_bcast_dims,
kBroadcasting);
537 void recordLabelToDimension(
const int32_t label,
const int axis,
const Shape &input_shape,
540 const int32_t input_dim = input_shape.Dims(axis);
542 if (label_to_dim_sizes->at(label) != 0 && label_to_dim_sizes->at(label) != input_dim)
544 std::runtime_error{
"Expected dimension " + std::to_string(label_to_dim_sizes->at(label)) +
545 " at axis " + std::to_string(axis) +
546 " of the input shaped but got dimension " + std::to_string(input_dim)};
548 (*label_to_dim_sizes)[label] = input_dim;
551 void insertBroadcastLabels(
int num_bcast_dims,
int num_named_labels,
int ellipsis_axis,
554 labels->erase(labels->begin() + ellipsis_axis);
555 labels->insert(labels->begin() + ellipsis_axis, num_bcast_dims, 0);
556 std::iota(labels->begin() + ellipsis_axis, labels->begin() + ellipsis_axis + num_bcast_dims,
560 label_counts->resize(num_named_labels + num_bcast_dims, 1);
563 template <
typename T>
564 void reduceOperand(
const InputTensor<T> &input,
const std::vector<DimensionType> &label_types,
566 bool *swap_free_and_contract,
Tensor *output)
571 std::vector<int32_t> permutation(
input.shape.DimensionsCount());
572 std::iota(permutation.begin(), permutation.end(), 0);
578 if (shouldSwapFreeAndContract(*labels, label_types))
580 *swap_free_and_contract =
true;
584 std::sort(permutation.begin(), permutation.end(), [&](
int i,
int j) {
585 int label_i = (*labels)[i];
586 int label_j = (*labels)[j];
587 return std::tie(label_types[label_i], label_i) < std::tie(label_types[label_j], label_j);
591 transposeOperand<T>(input, permutation, &input_transposed);
593 permuteLabels(permutation, labels);
597 labels->erase(std::unique(labels->begin(), labels->end()), labels->end());
598 strideOrInflate<T>(input_transposed, *labels, label_counts,
false ,
603 std::vector<int32_t>
reshape(5, 1);
610 std::vector<int32_t> output_shape_dims;
611 for (
size_t label_idx = 0; label_idx < labels->size(); ++label_idx)
613 const int label = labels->at(label_idx);
614 int32_t dim = input_deduped.shape.Dims(label_idx);
617 output_shape_dims.push_back(dim);
619 else if (label_types[label] ==
kFree)
621 free_labels->push_back(label);
623 reshape[label_types[label]] *= dim;
626 if (*swap_free_and_contract)
629 output_shape_dims.push_back(reshape[
kFree]);
630 output_shape_dims.push_back(reshape[
kContract]);
632 output_shape.ReplaceWith(output_shape_dims.size(), output_shape_dims.data());
641 using Reducer = Eigen::internal::SumReducer<T>;
642 using Index =
typename TTypes<T>::Tensor::Index;
647 const int32_t output_size =
650 device,
output->shaped<T, 1>({output_size}),
651 input_deduped.shaped<T, 2>({output_size, reshape[kReduce]}), Eigen::array<Index, 1>({1}),
655 bool shouldSwapFreeAndContract(
const Labels &labels,
656 const std::vector<DimensionType> &label_types)
660 std::vector<int> remap = {0, 1, 3, 2, 4};
661 for (
size_t i = 0; i + 1 < labels.size(); ++i)
663 const int dimtype_a = remap[label_types[labels[i]]];
664 const int dimtype_b = remap[label_types[labels[i + 1]]];
665 if (dimtype_a > dimtype_b || (dimtype_a == dimtype_b && labels[i] > labels[i + 1]))
673 template <
typename T>
674 void transposeOperand(
const InputTensor<T> &input,
const std::vector<int32_t> &permutation,
677 if (!shouldTranspose(
input.shape, permutation))
679 copyFrom(input,
input.shape, output);
682 Shape transposed_shape(
input.shape.DimensionsCount());
683 for (
int i = 0; i <
input.shape.DimensionsCount(); ++i)
685 transposed_shape.SetDim(i,
input.shape.Dims(permutation[i]));
689 if (
input.shape.FlatSize() == 0)
691 copyFrom(input, transposed_shape, output);
695 temp_operand.emplace_back(std::make_unique<T[]>(transposed_shape.FlatSize()));
696 T *new_buffer = temp_operand.back().get();
698 TransposeParams transpose_params;
699 transpose_params.perm_count = permutation.size();
700 for (
size_t i = 0; i < permutation.size(); i++)
702 transpose_params.perm[i] = permutation[i];
705 Transpose<T>(transpose_params,
input.shape,
input.buffer, transposed_shape, new_buffer);
707 output->shape.ReplaceWith(transposed_shape.DimensionsCount(), transposed_shape.DimsData());
708 output->buffer = new_buffer;
711 bool shouldTranspose(
const Shape &input_shape,
const std::vector<int32_t> &permutation)
713 if (input_shape.DimensionsCount() < 2)
715 for (
size_t i = 0; i < permutation.size(); ++i)
717 if (permutation[i] != (int32_t)i)
723 template <
typename T>
724 void copyFrom(
const InputTensor<T> &input,
const Shape &shape,
Tensor *output)
727 temp_tensor.shape.ReplaceWith(
input.shape.DimensionsCount(),
input.shape.DimsData());
728 temp_operand.emplace_back(std::make_unique<
float[]>(
input.shape.FlatSize()));
729 temp_tensor.buffer = temp_operand.back().get();
730 memcpy(temp_tensor.buffer,
input.buffer,
input.shape.FlatSize() *
sizeof(
float));
732 copyFrom(temp_tensor, shape, output);
737 if (
output->copyFrom(input, shape))
740 throw std::runtime_error{
"Einsum: Encountered error while reshaping a Tensor"};
744 void permuteLabels(
const std::vector<int32_t> &permutation,
Labels *labels)
746 Labels permuted_labels(labels->size());
747 for (
size_t i = 0; i < labels->size(); ++i)
749 permuted_labels[i] = (*labels)[permutation[i]];
751 labels->swap(permuted_labels);
756 template <
typename T>
758 const bool should_inflate,
Tensor *output)
761 if (std::all_of(label_counts.begin(), label_counts.end(), [](
int c) { return c <= 1; }))
763 return copyFrom(input,
input.shape, output);
770 std::vector<int32_t> strides;
775 Shape inflated_shape;
776 std::vector<int32_t> strided_shape_dims;
777 std::vector<int32_t> inflated_shape_dims;
778 for (
auto &&label : labels)
780 const int32_t count = label_counts[label];
781 const int current_axis =
782 should_inflate ? strided_shape_dims.size() : inflated_shape_dims.size();
783 const int32_t dim =
input.shape.Dims(current_axis);
784 strided_shape_dims.push_back(dim);
785 inflated_shape_dims.insert(inflated_shape_dims.end(), count, dim);
786 const int32_t reshape_dim = std::pow(dim, count);
787 reshape.push_back(reshape_dim);
791 const int32_t stride = (dim > 1 && count > 1) ? (reshape_dim - 1) / (dim - 1) : 1;
792 strides.push_back(stride);
795 strided_shape.ReplaceWith(strided_shape_dims.size(), strided_shape_dims.data());
796 inflated_shape.ReplaceWith(inflated_shape_dims.size(), inflated_shape_dims.data());
801 temp_operand.emplace_back(std::make_unique<
float[]>(
output_shape.FlatSize()));
802 output->buffer = temp_operand.back().get();
808#define NDIMS_CASE(N) \
811 if (should_inflate) \
813 auto output_map = output->shaped<T, N>(reshape); \
814 auto input_map = input.shaped<T, N>(strided_shape_dims); \
815 functor::InflateFunctor<Eigen::ThreadPoolDevice, T, N>()(device, input_map, strides, \
820 auto input_map = input.shaped<T, N>(reshape); \
821 auto output_map = output->shaped<T, N>(strided_shape_dims); \
822 functor::StrideFunctor<Eigen::ThreadPoolDevice, T, N>()(device, input_map, strides, \
834 throw std::runtime_error{
"Unsupported rank: " + std::to_string(
reshape.size()) +
835 " while handling repeated indices. Up to rank 6 is supported."};
840 void allocateTemp(
const Shape &shape,
Tensor *output)
842 output->shape.ReplaceWith(shape.DimensionsCount(), shape.DimsData());
843 temp_operand.emplace_back(std::make_unique<
float[]>(shape.FlatSize()));
844 output->buffer = temp_operand.back().get();
856 void contractOperands(std::vector<Tensor> &inputs, std::vector<bool> &swap_free_and_contract,
860 return copyFrom(inputs[0], inputs[0].shape, output);
862 MatMulBCast bcast(inputs[0].shape, inputs[1].shape);
863 if (!bcast.IsValid())
865 throw std::runtime_error{
"Einsum: Invalid broadcasting dimensions"};
869 reshapeToRank3(inputs[0], bcast.x_batch_size(), &lhs);
871 reshapeToRank3(inputs[1], bcast.y_batch_size(), &rhs);
872 Shape old_output_shape = bcast.output_batch_shape();
874 for (
int i = 0; i < old_output_shape.DimensionsCount(); i++)
879 for (
size_t i = 0; i <
inputs.size(); ++i)
881 const int32_t free_axis =
882 inputs[i].shape.DimensionsCount() - (swap_free_and_contract[i] ? 1 : 2);
883 output_shape.SetDim(i + old_output_shape.DimensionsCount(),
inputs[i].shape.Dims(free_axis));
885 bool adj_x = swap_free_and_contract[0];
886 bool adj_y = !swap_free_and_contract[1];
892 if (lhs.shape.FlatSize() == 0 || rhs.shape.FlatSize() == 0)
894 functor::SetZeroFunctor<Eigen::ThreadPoolDevice, float> set_zero;
901 reshapeToRank3(*output, bcast.output_batch_size(), &output_reshaped);
905 batchMatMul.prepare(lhs.shape, rhs.shape, adj_x, adj_y);
906 batchMatMul(lhs.shape, lhs.base<
float>(), rhs.shape, rhs.base<
float>(), adj_x, adj_y,
907 output_reshaped.shape, output_reshaped.base<
float>());
910 void reshapeToRank3(
const Tensor &input,
int batch_size,
Tensor *output)
912 const int rank =
input.shape.DimensionsCount();
922 std::vector<DimensionType> _label_types;
925 std::vector<bool> _input_has_ellipsis;
926 bool _output_has_ellipsis =
false;
928 std::vector<std::unique_ptr<float[]>> temp_operand;