Skip to content

Commit 44b88e4

Browse files
nihuinihui
authored
test gemm oom (#6429)
1 parent 544230f commit 44b88e4

File tree

7 files changed

+471
-0
lines changed

7 files changed

+471
-0
lines changed

‎src/layer/arm/gemm_arm.cpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4357,6 +4357,8 @@ int Gemm_arm::create_pipeline(const Option& opt)
43574357
{
43584358
int C_elempack = constantM % 4 == 0 ? 4 : 1;
43594359
convert_packing(C_data, CT_data, C_elempack, opt);
4360+
if (CT_data.empty())
4361+
return -100;
43604362
}
43614363
#endif // __ARM_NEON
43624364

@@ -4365,6 +4367,8 @@ int Gemm_arm::create_pipeline(const Option& opt)
43654367
{
43664368
Mat C2;
43674369
C2.create_like(CT_data);
4370+
if (C2.empty())
4371+
return -100;
43684372

43694373
const int size = CT_data.total() * CT_data.elempack;
43704374
for (int i = 0; i < size; i++)
@@ -4513,6 +4517,8 @@ int Gemm_arm::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& to
45134517
{
45144518
Mat CT_data;
45154519
CT_data.create_like(C, opt.workspace_allocator);
4520+
if (CT_data.empty())
4521+
return -100;
45164522

45174523
const int size = C.total() * C.elempack;
45184524
for (int i = 0; i < size; i++)
@@ -5079,6 +5085,8 @@ int Gemm_arm::create_pipeline_bf16s(const Option& opt)
50795085
{
50805086
int C_elempack = constantM % 4 == 0 ? 4 : 1;
50815087
convert_packing(C_data, CT_data, C_elempack, opt);
5088+
if (CT_data.empty())
5089+
return -100;
50825090
}
50835091
#endif // __ARM_NEON
50845092

@@ -5087,6 +5095,8 @@ int Gemm_arm::create_pipeline_bf16s(const Option& opt)
50875095
{
50885096
Mat C2;
50895097
C2.create_like(CT_data);
5098+
if (C2.empty())
5099+
return -100;
50905100

50915101
const int size = CT_data.total() * CT_data.elempack;
50925102
for (int i = 0; i < size; i++)
@@ -5210,6 +5220,8 @@ int Gemm_arm::forward_bf16s(const std::vector<Mat>& bottom_blobs, std::vector<Ma
52105220
{
52115221
Mat CT_data;
52125222
CT_data.create_like(C, opt.workspace_allocator);
5223+
if (CT_data.empty())
5224+
return -100;
52135225

52145226
const int size = C.total() * C.elempack;
52155227
for (int i = 0; i < size; i++)

‎src/layer/arm/gemm_arm_asimdhp.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2801,13 +2801,17 @@ int Gemm_arm::create_pipeline_fp16sa(const Option& opt)
28012801
if (constantC && constant_broadcast_type_C != -1)
28022802
{
28032803
cast_float32_to_float16(C_data, CT_data, opt);
2804+
if (CT_data.empty())
2805+
return -100;
28042806

28052807
if (constant_broadcast_type_C == 3 && opt.use_packing_layout)
28062808
{
28072809
int C_elempack = constantM % 8 == 0 ? 8 : constantM % 4 == 0 ? 4 : 1;
28082810
Mat tmp;
28092811
convert_packing(CT_data, tmp, C_elempack, opt);
28102812
CT_data = tmp;
2813+
if (CT_data.empty())
2814+
return -100;
28112815
}
28122816

28132817
// pre-multiply C with beta

‎src/layer/arm/gemm_arm_vfpv4.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -498,6 +498,8 @@ int Gemm_arm::create_pipeline_fp16s(const Option& opt)
498498
{
499499
int C_elempack = constantM % 4 == 0 ? 4 : 1;
500500
convert_packing(C_data, CT_data, C_elempack, opt);
501+
if (CT_data.empty())
502+
return -100;
501503
}
502504
#endif // __ARM_NEON
503505

@@ -506,6 +508,8 @@ int Gemm_arm::create_pipeline_fp16s(const Option& opt)
506508
{
507509
Mat C2;
508510
C2.create_like(CT_data);
511+
if (C2.empty())
512+
return -100;
509513

510514
const int size = CT_data.total() * CT_data.elempack;
511515
for (int i = 0; i < size; i++)
@@ -622,13 +626,17 @@ int Gemm_arm::forward_fp16s(const std::vector<Mat>& bottom_blobs, std::vector<Ma
622626
Mat CT_data;
623627
cast_float16_to_float32(C, CT_data);
624628
C = CT_data;
629+
if (C.empty())
630+
return -100;
625631
}
626632

627633
// pre-multiply C with beta
628634
if (beta != 1.f)
629635
{
630636
Mat CT_data;
631637
CT_data.create_like(C, opt.workspace_allocator);
638+
if (CT_data.empty())
639+
return -100;
632640

633641
const int size = C.total() * C.elempack;
634642
for (int i = 0; i < size; i++)

‎src/layer/gemm.cpp‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,8 @@ int Gemm::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
305305
{
306306
// transpose A to row-major
307307
A.create((A0.dims == 3 ? A0.c : A0.h), A0.w, elemsize, opt.workspace_allocator);
308+
if (A.empty())
309+
return -100;
308310

309311
const size_t A0_hstep = A0.dims == 3 ? A0.cstep : (size_t)A0.w;
310312

@@ -323,6 +325,8 @@ int Gemm::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
323325
{
324326
// transpose B to col-major
325327
BT.create((B0.dims == 3 ? B0.c : B0.h), B0.w, elemsize, opt.workspace_allocator);
328+
if (BT.empty())
329+
return -100;
326330

327331
const size_t B0_hstep = B0.dims == 3 ? B0.cstep : (size_t)B0.w;
328332

@@ -441,6 +445,8 @@ int Gemm::forward_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& t
441445
if (A0.elemsize == 1)
442446
{
443447
A.create(A0.h, A0.w, (size_t)1u, 1, opt.workspace_allocator);
448+
if (A.empty())
449+
return -100;
444450

445451
for (int i = 0; i < A.h; i++)
446452
{
@@ -454,6 +460,8 @@ int Gemm::forward_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& t
454460
else
455461
{
456462
A.create(A0.dims == 3 ? A0.c : A0.h, A0.w, (size_t)4u, 1, opt.workspace_allocator);
463+
if (A.empty())
464+
return -100;
457465

458466
for (int i = 0; i < A.h; i++)
459467
{
@@ -472,7 +480,11 @@ int Gemm::forward_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& t
472480
if (A_int8.elemsize != 1)
473481
{
474482
A_int8.create(A.w, A.dims == 3 ? A.c : A.h, (size_t)1u, 1, opt.workspace_allocator);
483+
if (A_int8.empty())
484+
return -100;
475485
A_int8_scales.create(A_int8.h, (size_t)4u, 1, opt.workspace_allocator);
486+
if (A_int8_scales.empty())
487+
return -100;
476488

477489
for (int i = 0; i < A_int8.h; i++)
478490
{
@@ -503,6 +515,8 @@ int Gemm::forward_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& t
503515
if (B0_int8.elemsize != 1)
504516
{
505517
B0_int8.create(B0.w, B0.dims == 3 ? B0.c : B0.h, (size_t)1u, 1, opt.workspace_allocator);
518+
if (B0_int8.empty())
519+
return -100;
506520

507521
float absmax = 0.f;
508522
for (int i = 0; i < B0_int8.h; i++)
@@ -537,6 +551,8 @@ int Gemm::forward_int8(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& t
537551
{
538552
// transpose B to col-major
539553
BT_int8.create(B0_int8.h, B0_int8.w, (size_t)1u, 1, opt.workspace_allocator);
554+
if (BT_int8.empty())
555+
return -100;
540556

541557
for (int i = 0; i < BT_int8.h; i++)
542558
{

‎src/layer/riscv/gemm_riscv.cpp‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1503,7 +1503,11 @@ static int gemm_riscv(const Mat& A, const Mat& B, const Mat& C, Mat& top_blob, i
15031503
int nn_K = (K + TILE_K - 1) / TILE_K;
15041504

15051505
Mat ATX(TILE_K * TILE_M, (K + TILE_K - 1) / TILE_K, nT, 4u, opt.workspace_allocator);
1506+
if (ATX.empty())
1507+
return -100;
15061508
Mat BT(TILE_K * TILE_N, (K + TILE_K - 1) / TILE_K, (N + TILE_N - 1) / TILE_N, 4u, opt.workspace_allocator);
1509+
if (BT.empty())
1510+
return -100;
15071511

15081512
const int nn_NK = nn_N * nn_K;
15091513

@@ -1537,7 +1541,11 @@ static int gemm_riscv(const Mat& A, const Mat& B, const Mat& C, Mat& top_blob, i
15371541

15381542
Mat topT;
15391543
if (K > TILE_K || broadcast_type_C == 3 || output_transpose)
1544+
{
15401545
topT.create(TILE_N * TILE_M, 1, nT, 4u, opt.workspace_allocator);
1546+
if (topT.empty())
1547+
return -100;
1548+
}
15411549

15421550
#pragma omp parallel for num_threads(nT)
15431551
for (int ppi = 0; ppi < nn_M; ppi++)
@@ -1617,6 +1625,8 @@ static int gemm_AT_riscv(const Mat& AT, const Mat& B, const Mat& C, Mat& top_blo
16171625
int nn_K = (K + TILE_K - 1) / TILE_K;
16181626

16191627
Mat BT(TILE_K * TILE_N, (K + TILE_K - 1) / TILE_K, (N + TILE_N - 1) / TILE_N, 4u, opt.workspace_allocator);
1628+
if (BT.empty())
1629+
return -100;
16201630

16211631
const int nn_NK = nn_N * nn_K;
16221632

@@ -1647,7 +1657,11 @@ static int gemm_AT_riscv(const Mat& AT, const Mat& B, const Mat& C, Mat& top_blo
16471657

16481658
Mat topT;
16491659
if (K > TILE_K || broadcast_type_C == 3 || output_transpose)
1660+
{
16501661
topT.create(TILE_N * TILE_M, 1, nT, 4u, opt.workspace_allocator);
1662+
if (topT.empty())
1663+
return -100;
1664+
}
16511665

16521666
#pragma omp parallel for num_threads(nT)
16531667
for (int ppi = 0; ppi < nn_M; ppi++)
@@ -1708,10 +1722,16 @@ static int gemm_BT_riscv(const Mat& A, const Mat& BT, const Mat& C, Mat& top_blo
17081722
// int nn_N = (N + TILE_N - 1) / TILE_N;
17091723

17101724
Mat ATX(TILE_K * TILE_M, (K + TILE_K - 1) / TILE_K, nT, 4u, opt.workspace_allocator);
1725+
if (ATX.empty())
1726+
return -100;
17111727

17121728
Mat topT;
17131729
if (K > TILE_K || broadcast_type_C == 3 || output_transpose)
1730+
{
17141731
topT.create(TILE_N * TILE_M, 1, nT, 4u, opt.workspace_allocator);
1732+
if (topT.empty())
1733+
return -100;
1734+
}
17151735

17161736
#pragma omp parallel for num_threads(nT)
17171737
for (int ppi = 0; ppi < nn_M; ppi++)
@@ -1790,7 +1810,11 @@ static int gemm_AT_BT_riscv(const Mat& AT, const Mat& BT, const Mat& C, Mat& top
17901810

17911811
Mat topT;
17921812
if (K > TILE_K || broadcast_type_C == 3 || output_transpose)
1813+
{
17931814
topT.create(TILE_N * TILE_M, 1, nT, 4u, opt.workspace_allocator);
1815+
if (topT.empty())
1816+
return -100;
1817+
}
17941818

17951819
#pragma omp parallel for num_threads(nT)
17961820
for (int ppi = 0; ppi < nn_M; ppi++)
@@ -1951,6 +1975,8 @@ int Gemm_riscv::create_pipeline(const Option& opt)
19511975
{
19521976
int C_elempack = constantM % packn == 0 ? packn : 1;
19531977
convert_packing(C_data, CT_data, C_elempack, opt);
1978+
if (CT_data.empty())
1979+
return -100;
19541980
}
19551981
#endif // __riscv_vector
19561982

@@ -1959,6 +1985,8 @@ int Gemm_riscv::create_pipeline(const Option& opt)
19591985
{
19601986
Mat C2;
19611987
C2.create_like(CT_data);
1988+
if (C2.empty())
1989+
return -100;
19621990

19631991
const int size = CT_data.total() * CT_data.elempack;
19641992
for (int i = 0; i < size; i++)
@@ -2082,6 +2110,8 @@ int Gemm_riscv::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
20822110
{
20832111
Mat CT_data;
20842112
CT_data.create_like(C, opt.workspace_allocator);
2113+
if (CT_data.empty())
2114+
return -100;
20852115

20862116
const int size = C.total() * C.elempack;
20872117
for (int i = 0; i < size; i++)

‎src/layer/x86/gemm_x86.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7500,6 +7500,8 @@ int Gemm_x86::create_pipeline(const Option& opt)
75007500
int C_elempack = constantM % 4 == 0 ? 4 : 1;
75017501
#endif
75027502
convert_packing(C_data, CT_data, C_elempack, opt);
7503+
if (CT_data.empty())
7504+
return -100;
75037505
}
75047506
#endif // __SSE2__
75057507

@@ -7508,6 +7510,8 @@ int Gemm_x86::create_pipeline(const Option& opt)
75087510
{
75097511
Mat C2;
75107512
C2.create_like(CT_data);
7513+
if (C2.empty())
7514+
return -100;
75117515

75127516
const int size = CT_data.total() * CT_data.elempack;
75137517
for (int i = 0; i < size; i++)

0 commit comments

Comments
 (0)