| 492 | |
| 493 | template <typename T, typename Context> |
| 494 | DenseTensor PerformContraction( |
| 495 | const Context& dev_ctx, |
| 496 | const std::vector<const DenseTensor*>& operands, |
| 497 | const std::vector<std::string>& input_strs, |
| 498 | const std::vector<LabelMap>& label2perm, |
| 499 | const std::vector<char>& all_labels, |
| 500 | const LabelMap& label2type, |
| 501 | const LabelMap& label2shape, |
| 502 | const std::vector<std::vector<int64_t>>& broadcast_shapes, |
| 503 | std::vector<DenseTensor*> cache, |
| 504 | bool use_cache) { |
| 505 | auto all_valid = LabelMap(1); |
| 506 | auto recover_dim = GetShapeByType<int64_t>( |
| 507 | all_labels, label2type, all_valid, label2shape, {LabelType::Batch}); |
| 508 | auto preprocess = [&](const DenseTensor& t, |
| 509 | const LabelMap& perm, |
| 510 | const std::vector<int64_t>& broadcast, |
| 511 | int operand_idx) -> DenseTensor { |
| 512 | // reshape |
| 513 | auto frees = GetShapeByType<int64_t>(all_labels, |
| 514 | label2type, |
| 515 | perm, |
| 516 | label2shape, |
| 517 | {LabelType::AO, LabelType::BO}); |
| 518 | auto conts = GetShapeByType<int64_t>( |
| 519 | all_labels, label2type, perm, label2shape, {LabelType::Contraction}); |
| 520 | std::vector<char> reordered_all_labels = all_labels; |
| 521 | if (operand_idx == 1) { |
| 522 | reordered_all_labels = TransformLabelsOrder(all_labels, |
| 523 | label2type, |
| 524 | {LabelType::Batch, |
| 525 | LabelType::Contraction, |
| 526 | LabelType::AO, |
| 527 | LabelType::BO, |
| 528 | LabelType::Reduction}); |
| 529 | } |
| 530 | // reduction |
| 531 | DenseTensor trans_t; |
| 532 | if (use_cache && cache[operand_idx] != nullptr && |
| 533 | cache[operand_idx]->IsInitialized()) { |
| 534 | trans_t.ShareBufferWith(*(cache[operand_idx])); |
| 535 | VLOG(5) << "Cache Used!"; |
| 536 | } else { |
| 537 | auto reduct_t = |
| 538 | PerformDiagonalAndReduction<T, Context>(dev_ctx, |
| 539 | t, |
| 540 | input_strs[operand_idx], |
| 541 | perm, |
| 542 | all_labels, |
| 543 | broadcast_shapes[operand_idx], |
| 544 | label2type); |
| 545 | trans_t = PerformTranspose<T, Context>( |
| 546 | dev_ctx, reduct_t, perm, reordered_all_labels, label2type); |
| 547 | if (cache[operand_idx] != nullptr) { |
| 548 | std::vector<int64_t> broadcast_shapes_restore( |
| 549 | broadcast_shapes[operand_idx].size()); |
| 550 | |
| 551 | auto contraction_dim1 = |
nothing calls this directly
no test coverage detected