[Snippets][CPU] Introduce jit_binary_call emitter on ARM, fix large offsets issue in Gemm #31553

aobolensk · 2025-08-01T08:06:14Z

Details:

Introduce jit_binary_call emitter on ARM for calling external functions in a more unified way and switch Gemm and GemmCopyB emitters to use it
Fix large offsets issue in Gemm
Fix register allocation and splilling
Enable all MatMulBias tests

Tickets:

169427
required for 167005

a-sidorova

Reviewed everything except jit_load_store_emitter parts

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/kernel_executors/gemm.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/kernel_executors/gemm.hpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_binary_call_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_binary_call_emitter.hpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_memory_emitters.cpp

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/kernel_executors/gemm.cpp

a-sidorova · 2025-08-05T07:52:13Z

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

+        const bool is_dynamic_offset = ov::snippets::utils::is_dynamic_value(m_memory_offsets[i]);
+        const bool is_valid_buffer_id = !ov::snippets::utils::is_dynamic_value(m_buffer_ids[i]);


Please compare m_buffer_id with SIZE_MAX because this value is returned from ov::intel_cpu::utils::get_buffer_cluster_id. If we change dynamic value from size_max to any another (for example, size_max-1), this check will be invalid. So let's align check with with get_buffer_cluster_id here

Should we call assert if the offset is dynamic and id is size_max? As it's done in memory emitters

Added these checks

OPENVINO_ASSERT is still missed here. Please double-check

Returned it back

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_copy_b_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_binary_call_emitter.cpp

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/kernel_executors/gemm.cpp

a-sidorova · 2025-08-06T05:34:10Z

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

+        const bool is_dynamic_offset = ov::snippets::utils::is_dynamic_value(m_memory_offsets[i]);
+        const bool is_valid_buffer_id = !ov::snippets::utils::is_dynamic_value(m_buffer_ids[i]);


OPENVINO_ASSERT is still missed here. Please double-check

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_copy_b_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_binary_call_emitter.cpp

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp

v-Golubev · 2025-08-06T16:05:16Z

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

+        if (off_mod == 0 && off_mul_vl >= 0 && off_mul_vl <= 4095) {
+            h->ldr(QReg(dst.getIdx()), ptr(src, static_cast<uint32_t>(offset)));
+        } else {
+            h->add_imm(h->X_DEFAULT_ADDR, src, offset, h->X_TMP_0);


Just a question for better understanding: is it safe to use h->X_DEFAULT_ADDR and h->X_TMP_0 here? How do we guarantee that src or dst are not X_DEFAULT_ADDR and X_TMP_0?

YEs, we do. They are basically not used in code generation as inputs or outputs

@v-Golubev I had the same question but these registers are in blacklist for AssignRegister pass:

openvino/src/plugins/intel_cpu/src/emitters/snippets/aarch64/cpu_generator.cpp

Lines 379 to 384 in 954c63f

// Note: more details on the usage of reserved registers in aarch64/jit_kernel_emitter.cpp

if (none_of(i, Operand::SP, Operand::X18, Operand::X23, Operand::X24, Operand::X28, Operand::X29)) {

reg_pool.emplace_back(snippets::RegType::gpr, i);

}

}

return reg_pool;

Is it possible to create a static vector with reserved registers somewhere? Just to avoid such questions in the future: we will have one place (vec initialization) where we will explain why they are in a blacklist. And we will be able to use them in the places we need using this vec. @a-sidorova @aobolensk what do you think?

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/kernel_executors/gemm.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_binary_call_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.cpp

a-sidorova · 2025-08-07T05:29:34Z

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

+        if (off_mod == 0 && off_mul_vl >= 0 && off_mul_vl <= 4095) {
+            h->ldr(QReg(dst.getIdx()), ptr(src, static_cast<uint32_t>(offset)));
+        } else {
+            h->add_imm(h->X_DEFAULT_ADDR, src, offset, h->X_TMP_0);


@v-Golubev I had the same question but these registers are in blacklist for AssignRegister pass:

openvino/src/plugins/intel_cpu/src/emitters/snippets/aarch64/cpu_generator.cpp

Lines 379 to 384 in 954c63f

// Note: more details on the usage of reserved registers in aarch64/jit_kernel_emitter.cpp

if (none_of(i, Operand::SP, Operand::X18, Operand::X23, Operand::X24, Operand::X28, Operand::X29)) {

reg_pool.emplace_back(snippets::RegType::gpr, i);

}

}

return reg_pool;

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.hpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

…nary-call-emitter

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.cpp

a-sidorova

LGTM

src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.hpp

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp

v-Golubev · 2025-08-08T13:03:23Z

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

+        if (off_mod == 0 && off_mul_vl >= 0 && off_mul_vl <= 4095) {
+            h->ldr(QReg(dst.getIdx()), ptr(src, static_cast<uint32_t>(offset)));
+        } else {
+            h->add_imm(h->X_DEFAULT_ADDR, src, offset, h->X_TMP_0);


Is it possible to create a static vector with reserved registers somewhere? Just to avoid such questions in the future: we will have one place (vec initialization) where we will explain why they are in a blacklist. And we will be able to use them in the places we need using this vec. @a-sidorova @aobolensk what do you think?

v-Golubev

LGTM

aobolensk requested review from a team as code owners August 1, 2025 08:06

aobolensk added the platform: arm OpenVINO on ARM / ARM64 label Aug 1, 2025

github-actions bot added the category: CPU OpenVINO CPU plugin label Aug 1, 2025

aobolensk assigned a-sidorova Aug 1, 2025

aobolensk force-pushed the snippets-arm-jit-binary-call-emitter branch 2 times, most recently from 5d584fa to b0328ac Compare August 1, 2025 09:53

a-sidorova assigned v-Golubev Aug 1, 2025

a-sidorova modified the milestones: 2025.2, 2025.3 Aug 1, 2025

a-sidorova reviewed Aug 1, 2025

View reviewed changes

aobolensk force-pushed the snippets-arm-jit-binary-call-emitter branch 3 times, most recently from 1b05f56 to e163aef Compare August 4, 2025 09:23

aobolensk requested a review from a-sidorova August 4, 2025 11:30

v-Golubev reviewed Aug 4, 2025

View reviewed changes

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp Outdated Show resolved Hide resolved

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_memory_emitters.cpp Show resolved Hide resolved

aobolensk requested a review from v-Golubev August 5, 2025 06:18

a-sidorova reviewed Aug 5, 2025

View reviewed changes

aobolensk force-pushed the snippets-arm-jit-binary-call-emitter branch 3 times, most recently from 4bb0d29 to e02ae65 Compare August 5, 2025 13:00

aobolensk requested a review from a-sidorova August 5, 2025 13:09

aobolensk added 7 commits August 5, 2025 15:57

no segfault

79378a2

fix stack

db5fdc0

WA uni_ldr issue

3a42edb

fmt

8ff18d8

tidy

6b53e85

Partially address Alexandra's review comments

2eedfd8

Adjust jit_gemm_emitter::emit_call

f5b04ed

Merge branch 'master' into snippets-arm-jit-binary-call-emitter

08c58cf

a-sidorova reviewed Aug 6, 2025

View reviewed changes

aobolensk added 4 commits August 6, 2025 08:58

Address Alexandra's review comments 2

2f26e95

tidy

4a15cad

Merge branch 'master' into snippets-arm-jit-binary-call-emitter

5fcb045

store offsets

954c63f

v-Golubev reviewed Aug 6, 2025

View reviewed changes

aobolensk force-pushed the snippets-arm-jit-binary-call-emitter branch from 4b9f771 to a20f240 Compare August 6, 2025 18:36

Apply Vladislav's comments

a20f240

a-sidorova reviewed Aug 7, 2025

View reviewed changes

aobolensk force-pushed the snippets-arm-jit-binary-call-emitter branch from d0e26c7 to a8fec5c Compare August 7, 2025 05:47

aobolensk added 2 commits August 7, 2025 08:19

Merge remote-tracking branch 'origin/master' into snippets-arm-jit-bi…

f9579c4

…nary-call-emitter

tidy

a8fec5c

aobolensk force-pushed the snippets-arm-jit-binary-call-emitter branch 2 times, most recently from e7031da to 34483bb Compare August 7, 2025 07:54

aobolensk added 2 commits August 7, 2025 10:51

Address Alexandra's review comments 3

34483bb

tidy 2

534480d

a-sidorova reviewed Aug 8, 2025

View reviewed changes

a-sidorova approved these changes Aug 8, 2025

View reviewed changes

a-sidorova requested a review from v-Golubev August 8, 2025 11:56

Address Alexandra's comments 4

c80fc13

v-Golubev reviewed Aug 8, 2025

View reviewed changes

aobolensk requested a review from v-Golubev August 8, 2025 14:12

v-Golubev approved these changes Aug 8, 2025

View reviewed changes

Address Vladislav's comments 2

44b125f

aobolensk enabled auto-merge August 8, 2025 16:56

aobolensk added this pull request to the merge queue Aug 8, 2025

Merged via the queue into openvinotoolkit:master with commit 3195c6f Aug 8, 2025
229 of 232 checks passed

aobolensk deleted the snippets-arm-jit-binary-call-emitter branch August 8, 2025 19:31

		const bool is_dynamic_offset = ov::snippets::utils::is_dynamic_value(m_memory_offsets[i]);
		const bool is_valid_buffer_id = !ov::snippets::utils::is_dynamic_value(m_buffer_ids[i]);

	// Note: more details on the usage of reserved registers in aarch64/jit_kernel_emitter.cpp
	if (none_of(i, Operand::SP, Operand::X18, Operand::X23, Operand::X24, Operand::X28, Operand::X29)) {
	reg_pool.emplace_back(snippets::RegType::gpr, i);
	}
	}
	return reg_pool;

[Snippets][CPU] Introduce jit_binary_call emitter on ARM, fix large offsets issue in Gemm #31553

[Snippets][CPU] Introduce jit_binary_call emitter on ARM, fix large offsets issue in Gemm #31553

Uh oh!

Conversation

aobolensk commented Aug 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Details:

Tickets:

Uh oh!

a-sidorova left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aobolensk commented Aug 1, 2025 •

edited

Loading