| 566 | } |
| 567 | |
| 568 | vector<Tensor> GpuRNNForwardInferenceEx(const Tensor &x, const Tensor &hx, |
| 569 | const Tensor &cx, const Tensor &W, |
| 570 | const Tensor &seq_lengths, |
| 571 | CudnnRNNHandle &h) { |
| 572 | CHECK_EQ(h.feature_size, x.shape(2)) << "feature size should not change"; |
| 573 | |
| 574 | Tensor y, hy, cy; |
| 575 | Shape yshape, states_shape; |
| 576 | |
| 577 | if (h.batch_first) { |
| 578 | LOG(FATAL) << "batch_first not implemented for GpuRNNForwardTrainingEx"; |
| 579 | } else { |
| 580 | h.seq_length = x.shape(0); |
| 581 | h.batch_size = x.shape(1); |
| 582 | yshape = Shape{h.seq_length, h.batch_size, |
| 583 | h.hidden_size * (h.bidirectional ? 2 : 1)}; |
| 584 | states_shape = Shape{h.num_layers * (h.bidirectional ? 2 : 1), h.batch_size, |
| 585 | h.hidden_size}; |
| 586 | } |
| 587 | |
| 588 | y = Tensor(yshape, x.device()); |
| 589 | hy = Tensor(states_shape, x.device()); |
| 590 | cy = Tensor(states_shape, x.device()); |
| 591 | |
| 592 | y.device()->Exec( |
| 593 | [y, hy, cy, x, seq_lengths, hx, cx, &W, &h](Context *ctx) { |
| 594 | // data descriptor |
| 595 | cudnnRNNDataDescriptor_t xDesc, yDesc; |
| 596 | init_data_desc(xDesc, h.feature_size, seq_lengths, h); |
| 597 | init_data_desc(yDesc, |
| 598 | h.bidirectional ? h.hidden_size * 2 : h.hidden_size, |
| 599 | seq_lengths, h); |
| 600 | |
| 601 | // hidden cell states descriptor |
| 602 | cudnnTensorDescriptor_t hxDesc, cxDesc, hyDesc, cyDesc; |
| 603 | init_hc_Desc(hxDesc, h); |
| 604 | init_hc_Desc(cxDesc, h); |
| 605 | init_hc_Desc(hyDesc, h); |
| 606 | init_hc_Desc(cyDesc, h); |
| 607 | |
| 608 | auto xptr = x.block()->data(); |
| 609 | auto hxptr = hx.block()->data(); |
| 610 | auto cxptr = cx.block()->data(); |
| 611 | auto Wptr = W.block()->data(); |
| 612 | auto yptr = y.block()->mutable_data(); |
| 613 | auto hyptr = hy.block()->mutable_data(); |
| 614 | auto cyptr = cy.block()->mutable_data(); |
| 615 | auto wsptr = h.workspace.block()->mutable_data(); |
| 616 | |
| 617 | /* This routine is the extended version of the cudnnRNNForwardTraining() |
| 618 | function. The cudnnRNNForwardTrainingEx() allows the user to use |
| 619 | unpacked (padded) layout for input x and output y. |
| 620 | */ |
| 621 | CUDNN_CHECK(cudnnRNNForwardInferenceEx( |
| 622 | ctx->cudnn_handle, h.rnnDesc, xDesc, xptr, hxDesc, hxptr, cxDesc, |
| 623 | cxptr, h.wDesc, Wptr, yDesc, yptr, hyDesc, hyptr, cyDesc, cyptr, |
| 624 | NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, wsptr, |
| 625 | h.workspace_size_bytes)); |