| 631 | } |
| 632 | |
| 633 | vector<Tensor> GpuRNNForwardTrainingEx(const Tensor &x, const Tensor &hx, |
| 634 | const Tensor &cx, const Tensor &W, |
| 635 | const Tensor &seq_lengths, |
| 636 | CudnnRNNHandle &h) { |
| 637 | CHECK_EQ(h.feature_size, x.shape(2)) << "feature size should not change"; |
| 638 | |
| 639 | Tensor y, hy, cy; |
| 640 | Shape yshape, states_shape; |
| 641 | |
| 642 | if (h.batch_first) { |
| 643 | LOG(FATAL) << "batch_first not implemented for GpuRNNForwardTrainingEx"; |
| 644 | } else { |
| 645 | h.seq_length = x.shape(0); |
| 646 | h.batch_size = x.shape(1); |
| 647 | yshape = Shape{h.seq_length, h.batch_size, |
| 648 | h.hidden_size * (h.bidirectional ? 2 : 1)}; |
| 649 | states_shape = Shape{h.num_layers * (h.bidirectional ? 2 : 1), h.batch_size, |
| 650 | h.hidden_size}; |
| 651 | } |
| 652 | |
| 653 | y = Tensor(yshape, x.device()); |
| 654 | hy = Tensor(states_shape, x.device()); |
| 655 | cy = Tensor(states_shape, x.device()); |
| 656 | |
| 657 | y.device()->Exec( |
| 658 | [y, hy, cy, x, seq_lengths, hx, cx, &W, &h](Context *ctx) { |
| 659 | // data descriptor |
| 660 | cudnnRNNDataDescriptor_t xDesc, yDesc; |
| 661 | init_data_desc(xDesc, h.feature_size, seq_lengths, h); |
| 662 | init_data_desc(yDesc, |
| 663 | h.bidirectional ? h.hidden_size * 2 : h.hidden_size, |
| 664 | seq_lengths, h); |
| 665 | |
| 666 | // hidden cell states descriptor |
| 667 | cudnnTensorDescriptor_t hxDesc, cxDesc, hyDesc, cyDesc; |
| 668 | init_hc_Desc(hxDesc, h); |
| 669 | init_hc_Desc(cxDesc, h); |
| 670 | init_hc_Desc(hyDesc, h); |
| 671 | init_hc_Desc(cyDesc, h); |
| 672 | |
| 673 | auto xptr = x.block()->data(); |
| 674 | auto hxptr = hx.block()->data(); |
| 675 | auto cxptr = cx.block()->data(); |
| 676 | auto Wptr = W.block()->data(); |
| 677 | auto yptr = y.block()->mutable_data(); |
| 678 | auto hyptr = hy.block()->mutable_data(); |
| 679 | auto cyptr = cy.block()->mutable_data(); |
| 680 | auto wsptr = h.workspace.block()->mutable_data(); |
| 681 | auto rsptr = h.reserve_space.block()->mutable_data(); |
| 682 | |
| 683 | /* This routine is the extended version of the cudnnRNNForwardTraining() |
| 684 | function. The cudnnRNNForwardTrainingEx() allows the user to use |
| 685 | unpacked (padded) layout for input x and output y. |
| 686 | */ |
| 687 | CUDNN_CHECK(cudnnRNNForwardTrainingEx( |
| 688 | ctx->cudnn_handle, h.rnnDesc, xDesc, xptr, hxDesc, hxptr, cxDesc, |
| 689 | cxptr, h.wDesc, Wptr, yDesc, yptr, hyDesc, hyptr, cyDesc, cyptr, |
| 690 | NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, wsptr, |