| 638 | template <typename Scalar> |
| 639 | struct ParallelMatMulKernelSYCL { |
| 640 | static void Run(const OpKernelContext* context, const Tensor& in_x, |
| 641 | const Tensor& in_y, bool adj_x, bool adj_y, bool trans_x, |
| 642 | bool trans_y, const MatMulBCast& bcast, Tensor* out, |
| 643 | int start, int limit) { |
| 644 | auto Tx = in_x.tensor<Scalar, 3>(); |
| 645 | auto Ty = in_y.tensor<Scalar, 3>(); |
| 646 | auto Tz = out->tensor<Scalar, 3>(); |
| 647 | Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> contract_pairs; |
| 648 | contract_pairs[0] = ContractionDims(adj_x || trans_x, adj_y || trans_y); |
| 649 | auto d = context->eigen_sycl_device(); |
| 650 | |
| 651 | const bool should_bcast = bcast.IsBroadcastingRequired(); |
| 652 | const auto& x_batch_indices = bcast.x_batch_indices(); |
| 653 | const auto& y_batch_indices = bcast.y_batch_indices(); |
| 654 | for (int64 i = start; i < limit; ++i) { |
| 655 | const int64 x_batch_index = should_bcast ? x_batch_indices[i] : i; |
| 656 | const int64 y_batch_index = should_bcast ? y_batch_indices[i] : i; |
| 657 | |
| 658 | auto x = Tx.template chip<0>(x_batch_index); |
| 659 | auto y = Ty.template chip<0>(y_batch_index); |
| 660 | auto z = Tz.template chip<0>(i); |
| 661 | z.device(d) = x.contract(y, contract_pairs); |
| 662 | } |
| 663 | } |
| 664 | }; |
| 665 | |
| 666 | template <typename Scalar> |
nothing calls this directly
no test coverage detected