Skip to content

Commit 8465e78

Browse files
committed
auto grow the size and fix test
test=develop
1 parent 9255119 commit 8465e78

File tree

2 files changed

+4
-3
lines changed

2 files changed

+4
-3
lines changed

paddle/fluid/operators/math/jit_kernel_blas.cc

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,8 +65,9 @@ class VMulKernelImpl : public VMulKernel<T> {
6565

6666
explicit VMulKernelImpl(int d) : VMulKernel<T>() {
6767
if (useJIT(d)) {
68-
constexpr size_t sz = 256 * 1024; // TODO(TJ): should be related with d
69-
jitcode_.reset(new gen::VMulJitCode(d, sz));
68+
// roughly estimate the size of code
69+
size_t sz = 96 + d / AVX_FLOAT_BLOCK * 4 * 8;
70+
jitcode_.reset(new gen::VMulJitCode(d, sz > 4096 ? sz : 4096));
7071
this->Compute =
7172
jitcode_->getCode<void (*)(const T*, const T*, T*, int)>();
7273
return;

paddle/fluid/operators/math/jit_kernel_test.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -800,7 +800,7 @@ TEST(JitKernel, pool) {
800800
EXPECT_TRUE(std::dynamic_pointer_cast<const jit::Kernel>(pvmul_f) !=
801801
std::dynamic_pointer_cast<const jit::Kernel>(pvmul_d));
802802

803-
const auto& pvmul_from_key = jit::KernelPool::Instance().Get("vmulfany");
803+
const auto& pvmul_from_key = jit::KernelPool::Instance().Get("vmulfjit4");
804804
EXPECT_EQ(pvmul_f, pvmul_from_key);
805805
const auto& pvmul_from_key2 = jit::KernelPool::Instance().Get("vmulfjit");
806806
EXPECT_TRUE(pvmul_from_key2 == nullptr);

0 commit comments

Comments
 (0)