| 680 | class GeometryLSTMBlockCell : public GeometryComputer { |
| 681 | public: |
| 682 | virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, |
| 683 | Context& context, CommandBuffer& res) const override { |
| 684 | /* |
| 685 | shapes: |
| 686 | x: [batchSize, inputSize] |
| 687 | cs_prev, i, cs, f, o, ci, co, h: [batchSize, hiddenSize] |
| 688 | wci, wcf, wco: [hiddenSize] |
| 689 | w: [inputSize + hiddenSize, 4 * hiddenSize] |
| 690 | b: [4 * hiddenSize] |
| 691 | */ |
| 692 | // inputs |
| 693 | auto x = inputs[0]; |
| 694 | auto cs_prev = inputs[1]; |
| 695 | auto h_prev = inputs[2]; |
| 696 | auto w = inputs[3]; |
| 697 | auto wci = inputs[4]; |
| 698 | auto wcf = inputs[5]; |
| 699 | auto wco = inputs[6]; |
| 700 | auto b = inputs[7]; |
| 701 | // outputs |
| 702 | auto i = outputs[0]; |
| 703 | auto cs = outputs[1]; |
| 704 | auto f = outputs[2]; |
| 705 | auto o = outputs[3]; |
| 706 | auto ci = outputs[4]; |
| 707 | auto co = outputs[5]; |
| 708 | auto h = outputs[6]; |
| 709 | int batchSize = x->length(0); |
| 710 | int inputSize = x->length(1); |
| 711 | int hiddenSize = h_prev->length(1); |
| 712 | // params |
| 713 | auto param = op->main_as_LSTMBlockCell(); |
| 714 | float cell_clip = param->cell_clip(); |
| 715 | float forget_bias = param->forget_bias(); |
| 716 | bool use_peephole = param->use_peephole(); |
| 717 | // xh = [x, h_prev] |
| 718 | std::shared_ptr<Tensor> xh(Tensor::createDevice<float>({batchSize, inputSize + hiddenSize})); |
| 719 | { |
| 720 | auto xhDes = TensorUtils::getDescribe(xh.get()); |
| 721 | xhDes->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL; |
| 722 | xhDes->regions.resize(2); |
| 723 | xhDes->regions[0].origin = x; |
| 724 | xhDes->regions[0].size[0] = batchSize; |
| 725 | xhDes->regions[0].size[1] = inputSize; |
| 726 | xhDes->regions[0].src.stride[0] = inputSize; |
| 727 | xhDes->regions[0].dst.stride[0] = inputSize + hiddenSize; |
| 728 | xhDes->regions[1].origin = h_prev; |
| 729 | xhDes->regions[1].size[0] = batchSize; |
| 730 | xhDes->regions[1].size[1] = hiddenSize; |
| 731 | xhDes->regions[1].dst.offset = inputSize; |
| 732 | xhDes->regions[1].src.stride[0] = hiddenSize; |
| 733 | xhDes->regions[1].dst.stride[0] = inputSize + hiddenSize; |
| 734 | res.extras.emplace_back(xh); |
| 735 | } |
| 736 | // icfo = xh * w + b |
| 737 | std::shared_ptr<Tensor> icfo(Tensor::createDevice<float>({batchSize, 4 * hiddenSize})); |
| 738 | { |
| 739 | res.command.emplace_back(GeometryComputerUtils::makeMatMul(xh.get(), w, icfo.get(), b, false, false)); |