diff --git a/paddle/fluid/operators/math/jit_code.cc b/paddle/fluid/operators/math/jit_code.cc index ccc9206f5cda888ea11be9f5492570340dec87db..03b67238fe7eb5404a6d68841017bfe6d8665f89 100644 --- a/paddle/fluid/operators/math/jit_code.cc +++ b/paddle/fluid/operators/math/jit_code.cc @@ -214,6 +214,9 @@ void VActJitCode::generate() { bool LSTMJitCode::init(int d) { return MayIUse(avx) && d % 8 == 0; } void LSTMJitCode::generate() { + if (use_peephole_) { + preCode(); + } reg64_t reg_ptr_gates = rax; reg64_t reg_ptr_ct_1 = r9; reg64_t reg_ptr_ct = r10; @@ -224,18 +227,19 @@ void LSTMJitCode::generate() { mov(reg_ptr_ht, ptr[param1 + offsetof(lstm_t, ht)]); int offset = 0; + int d = num_ * sizeof(float); for (int i = 0; i < num_ / YMM_FLOAT_BLOCK; ++i) { /* C_t = C_t-1 * fgated + cand_gated * igated*/ // c vmovups(ymm_src, ptr[reg_ptr_gates + offset]); act(ymm_c, ymm_src, act_cand_); // i - vmovups(ymm_src, ptr[reg_ptr_gates + offset + num_]); + vmovups(ymm_src, ptr[reg_ptr_gates + offset + d]); act(ymm_i, ymm_src, act_gate_); vmulps(ymm_c, ymm_c, ymm_i); if (!compute_c1h1_) { // f - vmovups(ymm_src, ptr[reg_ptr_gates + offset + 2 * num_]); + vmovups(ymm_src, ptr[reg_ptr_gates + offset + 2 * d]); act(ymm_f, ymm_src, act_gate_); vmovups(ymm_i, ptr[reg_ptr_ct_1 + offset]); vmulps(ymm_f, ymm_f, ymm_i); @@ -245,20 +249,36 @@ void LSTMJitCode::generate() { ymm_t ymm_ct = compute_c1h1_ ? ymm_c : ymm_f; ymm_t ymm_o = compute_c1h1_ ? ymm_f : ymm_c; ymm_t ymm_tmp = ymm_i; + vmovups(ptr[reg_ptr_ct + offset], ymm_ct); // save ct act(ymm_tmp, ymm_ct, act_cell_); - vmovups(ymm_src, ptr[reg_ptr_gates + offset + 3 * num_]); + vmovups(ymm_src, ptr[reg_ptr_gates + offset + 3 * d]); act(ymm_o, ymm_src, act_gate_); vmulps(ymm_o, ymm_tmp, ymm_o); - // save ct and ht - vmovups(ptr[reg_ptr_ct + offset], ymm_ct); - vmovups(ptr[reg_ptr_ht + offset], ymm_o); - + vmovups(ptr[reg_ptr_ht + offset], ymm_o); // save ht offset += sizeof(float) * YMM_FLOAT_BLOCK; } - ret(); + if (use_peephole_) { + postCode(); + } else { + ret(); + } } +bool GRUJitCode::init(int d) { return MayIUse(avx) && d % 8 == 0; } + +void GRUJitCode::generate() { + reg64_t reg_ptr_gates = rax; + reg64_t reg_ptr_ct_1 = r9; + reg64_t reg_ptr_ct = r10; + reg64_t reg_ptr_ht = r11; + mov(reg_ptr_gates, ptr[param1 + offsetof(lstm_t, gates)]); + mov(reg_ptr_ct_1, ptr[param1 + offsetof(lstm_t, ct_1)]); + mov(reg_ptr_ct, ptr[param1 + offsetof(lstm_t, ct)]); + mov(reg_ptr_ht, ptr[param1 + offsetof(lstm_t, ht)]); + + ret(); +} } // namespace gen } // namespace jitkernel } // namespace math diff --git a/paddle/fluid/operators/math/jit_code.h b/paddle/fluid/operators/math/jit_code.h index bf28d444b7712a5bac32985df31b5d24b94837a1..403cea399102c04328567ca119e65ed56d57bc35 100644 --- a/paddle/fluid/operators/math/jit_code.h +++ b/paddle/fluid/operators/math/jit_code.h @@ -302,6 +302,34 @@ class VActJitCode : public JitCode { pop(reg_ptr_global); } + template + void act(JMM& dst, JMM& src, operand_type type) { // NOLINT + // use 15 + JMM zero = JMM(15); + if (type_ == operand_type::relu) { + vxorps(zero, zero, zero); + } + switch (type) { + case operand_type::relu: + relu_jmm(dst, src, zero); + break; + case operand_type::exp: + exp_jmm(dst, src, 2, 3, 4, 5); + break; + case operand_type::sigmoid: + sigmoid_jmm(dst, src, 2, 3, 4, 5); + break; + case operand_type::tanh: + tanh_jmm(dst, src, 2, 3, 4, 5); + break; + case operand_type::identity: + break; + default: + // throw error + break; + } + } + protected: int num_; operand_type type_; @@ -386,44 +414,94 @@ class LSTMJitCode : public VActJitCode { operand_type act_cand_; operand_type act_cell_; reg64_t param1{abi_param1}; - xmm_t xmm_src = xmm_t(0); xmm_t xmm_c = xmm_t(1); - xmm_t xmm_i = xmm_t(2); - xmm_t xmm_f = xmm_t(3); + xmm_t xmm_i = xmm_t(6); + xmm_t xmm_f = xmm_t(7); ymm_t ymm_src = ymm_t(0); - ymm_t ymm_c = ymm_t(1); - ymm_t ymm_i = ymm_t(2); - ymm_t ymm_f = ymm_t(3); + ymm_t ymm_c = ymm_t(1); // 2~5 for act + ymm_t ymm_i = ymm_t(6); + ymm_t ymm_f = ymm_t(7); +}; - template - void act(JMM& dst, JMM& src, operand_type type) { // NOLINT - // use 15 - JMM zero = JMM(15); - if (type_ == operand_type::relu) { - vxorps(zero, zero, zero); - } - switch (type) { - case operand_type::relu: - relu_jmm(dst, src, zero); - break; - case operand_type::exp: - exp_jmm(dst, src, 2, 3, 4, 5); - break; - case operand_type::sigmoid: - sigmoid_jmm(dst, src, 2, 3, 4, 5); - break; - case operand_type::tanh: - tanh_jmm(dst, src, 2, 3, 4, 5); - break; - case operand_type::identity: - break; - default: - // throw error - break; +class GRUJitCode : public VActJitCode { + public: + const char* name() const override { + std::string base = "GRUJitCode"; + if (id_ == 0) { + base += "_H1"; + } else if (id_ == 1) { + base += "_HtPart1"; + } else if (id_ == 2) { + base += "_HtPart2"; } + auto AddTypeStr = [&](operand_type type) { + switch (type) { + case operand_type::relu: + base += "_Relu"; + break; + case operand_type::exp: + base += "_Exp"; + break; + case operand_type::sigmoid: + base += "_Sigmoid"; + break; + case operand_type::tanh: + base += "_Tanh"; + break; + case operand_type::identity: + base += "_Identity"; + break; + default: + break; + } + }; + AddTypeStr(act_gate_); + AddTypeStr(act_cand_); + return base.c_str(); } + + explicit GRUJitCode(int id, const gru_attr_t& attr, + size_t code_size = 256 * 1024, void* code_ptr = nullptr) + : VActJitCode(attr.d, operand_type::sigmoid /* this is bugy*/, code_size, + code_ptr), + id_(id) { + auto typeExchange = [](const std::string& type) -> gen::operand_type { + if (type == "sigmoid") { + return operand_type::sigmoid; + } else if (type == "relu") { + return operand_type::relu; + } else if (type == "tanh") { + return operand_type::tanh; + } else if (type == "identity" || type == "") { + return operand_type::identity; + } // else throw error + return operand_type::identity; + }; + num_ = attr.d; + act_gate_ = typeExchange(attr.act_gate); + act_cand_ = typeExchange(attr.act_cand); + } + static bool init(int d); + void generate() override; + + protected: + int id_; + int num_; + operand_type act_gate_; + operand_type act_cand_; + reg64_t param1{abi_param1}; + + xmm_t xmm_src = xmm_t(0); + xmm_t xmm_c = xmm_t(1); + xmm_t xmm_i = xmm_t(6); + xmm_t xmm_f = xmm_t(7); + + ymm_t ymm_src = ymm_t(0); + ymm_t ymm_c = ymm_t(1); + ymm_t ymm_i = ymm_t(6); + ymm_t ymm_f = ymm_t(7); }; #ifdef PADDLE_WITH_MKLDNN diff --git a/paddle/fluid/operators/math/jit_kernel_rnn.cc b/paddle/fluid/operators/math/jit_kernel_rnn.cc index dbfd212e6e7b2b1e6f26c9442ab20ac05d8c3a30..e571d8adf4af681405fcd30a01edc895c66c34b6 100644 --- a/paddle/fluid/operators/math/jit_kernel_rnn.cc +++ b/paddle/fluid/operators/math/jit_kernel_rnn.cc @@ -40,7 +40,7 @@ class LSTMKernelImpl : public LSTMKernel { explicit LSTMKernelImpl(const lstm_attr_t& attr) : LSTMKernel() { #ifdef PADDLE_WITH_XBYAK if (useJIT(attr.d)) { - size_t sz = 96 + attr.d / YMM_FLOAT_BLOCK * 84 * 8; // should change + size_t sz = 96 + attr.d / YMM_FLOAT_BLOCK * 90 * 4 * 8; jitcode0_.reset(new gen::LSTMJitCode(false, attr, sz > 4096 ? sz : 4096)); this->ComputeCtHt = jitcode0_->getCode(); @@ -66,7 +66,7 @@ class LSTMKernelImpl : public LSTMKernel { #ifdef PADDLE_WITH_XBYAK template <> bool LSTMKernelImpl::useJIT(int d) { - return false; // not ready yet gen::LSTMJitCode::init(d); + return gen::LSTMJitCode::init(d); } #endif @@ -82,7 +82,7 @@ class PeepholeKernelImpl : public LSTMKernel { explicit PeepholeKernelImpl(const lstm_attr_t& attr) : LSTMKernel() { #ifdef PADDLE_WITH_XBYAK if (useJIT(attr.d)) { - size_t sz = 96 + attr.d / YMM_FLOAT_BLOCK * 84 * 8; // should change + size_t sz = 96 + attr.d / YMM_FLOAT_BLOCK * 96 * 4 * 8; jitcode0_.reset(new gen::LSTMJitCode(false, attr, sz > 4096 ? sz : 4096)); this->ComputeCtHt = jitcode0_->getCode(); @@ -175,12 +175,42 @@ class GRUKernelImpl : public GRUKernel { static inline bool useJIT(int d) { return false; } static inline bool useMKL(int d) { return false; } explicit GRUKernelImpl(const gru_attr_t& attr) : GRUKernel() { +#ifdef PADDLE_WITH_XBYAK + if (useJIT(attr.d)) { + size_t sz = 96 + attr.d / YMM_FLOAT_BLOCK * 84 * 8; // should change + jitcode0_.reset(new gen::GRUJitCode(0, attr, sz > 4096 ? sz : 4096)); + this->ComputeH1 = + jitcode0_->getCode(); + + jitcode1_.reset(new gen::GRUJitCode(1, attr, sz > 4096 ? sz : 4096)); + this->ComputeHtPart1 = + jitcode1_->getCode(); + + jitcode2_.reset(new gen::GRUJitCode(2, attr, sz > 4096 ? sz : 4096)); + this->ComputeHtPart2 = + jitcode1_->getCode(); + return; + } +#endif this->ComputeH1 = refer::GRUH1; this->ComputeHtPart1 = refer::GRUHtPart1; this->ComputeHtPart2 = refer::GRUHtPart2; } +#ifdef PADDLE_WITH_XBYAK + + private: + std::unique_ptr jitcode0_{nullptr}, jitcode1_{nullptr}, + jitcode2_{nullptr}; +#endif }; +#ifdef PADDLE_WITH_XBYAK +template <> +bool GRUKernelImpl::useJIT(int d) { + return false; // jitcode not ready yet +} +#endif + #define JITKERNEL_DEFINE_NAME_GRU(ker_key, ker_class) \ template <> \ std::string ker_class##Impl::name(const gru_attr_t& attr) { \