m53297601
/
Paddle
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

init gru jitcode and fix lstm jitcode

Browse Source 
test=develop
local_add_cudnn_lstm
tensor-tang 6 years ago
parent 0f25446574
commit e3b61cf52b
3 changed files with 170 additions and 42 deletions
Show Stats Download Patch File Download Diff File

36
paddle/fluid/operators/math/jit_code.cc
Unescape View File

@ -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_t>(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_t>(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_t>(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_t>(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_t>(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

140
paddle/fluid/operators/math/jit_code.h
Unescape View File

@ -302,6 +302,34 @@ class VActJitCode : public JitCode {
pop(reg_ptr_global);
}
template <typename JMM>
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<JMM>(dst, src, zero);
break;
case operand_type::exp:
exp_jmm<JMM>(dst, src, 2, 3, 4, 5);
break;
case operand_type::sigmoid:
sigmoid_jmm<JMM>(dst, src, 2, 3, 4, 5);
break;
case operand_type::tanh:
tanh_jmm<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 <typename JMM>
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<JMM>(dst, src, zero);
break;
case operand_type::exp:
exp_jmm<JMM>(dst, src, 2, 3, 4, 5);
break;
case operand_type::sigmoid:
sigmoid_jmm<JMM>(dst, src, 2, 3, 4, 5);
break;
case operand_type::tanh:
tanh_jmm<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

36
paddle/fluid/operators/math/jit_kernel_rnn.cc
Unescape View File

@ -40,7 +40,7 @@ class LSTMKernelImpl : public LSTMKernel<T> {
explicit LSTMKernelImpl(const lstm_attr_t& attr) : LSTMKernel<T>() {
#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<void (*)(lstm_t*, const lstm_attr_t*)>();
@ -66,7 +66,7 @@ class LSTMKernelImpl : public LSTMKernel<T> {
#ifdef PADDLE_WITH_XBYAK
template <>
bool LSTMKernelImpl<float>::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<T> {
explicit PeepholeKernelImpl(const lstm_attr_t& attr) : LSTMKernel<T>() {
#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<void (*)(lstm_t*, const lstm_attr_t*)>();
@ -175,12 +175,42 @@ class GRUKernelImpl : public GRUKernel<T> {
static inline bool useJIT(int d) { return false; }
static inline bool useMKL(int d) { return false; }
explicit GRUKernelImpl(const gru_attr_t& attr) : GRUKernel<T>() {
#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<void (*)(gru_t*, const gru_attr_t*)>();
jitcode1_.reset(new gen::GRUJitCode(1, attr, sz > 4096 ? sz : 4096));
this->ComputeHtPart1 =
jitcode1_->getCode<void (*)(gru_t*, const gru_attr_t*)>();
jitcode2_.reset(new gen::GRUJitCode(2, attr, sz > 4096 ? sz : 4096));
this->ComputeHtPart2 =
jitcode1_->getCode<void (*)(gru_t*, const gru_attr_t*)>();
return;
}
#endif
this->ComputeH1 = refer::GRUH1<T>;
this->ComputeHtPart1 = refer::GRUHtPart1<T>;
this->ComputeHtPart2 = refer::GRUHtPart2<T>;
}
#ifdef PADDLE_WITH_XBYAK
private:
std::unique_ptr<gen::GRUJitCode> jitcode0_{nullptr}, jitcode1_{nullptr},
jitcode2_{nullptr};
#endif
};
#ifdef PADDLE_WITH_XBYAK
template <>
bool GRUKernelImpl<float>::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<float>::name(const gru_attr_t& attr) { \

Write Preview
Loading…
Cancel
Save