[Snippets][CPU] Support static and dynamic offsets in JIT Gemm and GemmCopyB emitters #31375

aobolensk · 2025-07-17T16:56:38Z

Details:

Add offset handling (m_memory_offsets, m_buffer_ids) to AArch64 Gemm/GemmCopyB emitters
Introduce new AArch64 utils helpers for safe GPR selection and stack‑based pointer patching

Tickets:

169427

…mmCopyB emitters

a-sidorova · 2025-07-18T07:14:20Z

@chenhu-wang May I ask you to review this PR please since you're an author of these JIT emitters?

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

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

chenhu-wang

LGTM!

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

aobolensk · 2025-07-21T09:08:33Z

@chenhu-wang had to revert gpr related change, because it crashes on some test suites.

I have investigated the root cause. The stack should be 16 bytes aligned according to ARM spec:

https://community.arm.com/arm-community-blogs/b/architectures-and-processors-blog/posts/using-the-stack-in-aarch32-and-aarch64
- "At public interfaces, the alignment of sp must be two times the pointer size. For AArch32 that's 8 bytes, and for AArch64 it's 16 bytes. A "public interface" is typically a function that is visible to some other, separately-compiled code. The exact definition depends upon the language and the toolchain, and is out of scope of this article. It's reasonable to assume that any C or C++ functions that you interact with using assembly are treated as public interfaces."
https://learn.microsoft.com/en-us/cpp/build/arm64-windows-abi-conventions
- "Following the ABI put forth by ARM, the stack must always remain 16-byte aligned"

This reverts commit 7dc7b92.

chenhu-wang · 2025-07-22T00:46:27Z

@chenhu-wang had to revert gpr related change, because it crashes on some test suites.

I have investigated the root cause. The stack should be 16 bytes aligned according to ARM spec:

https://community.arm.com/arm-community-blogs/b/architectures-and-processors-blog/posts/using-the-stack-in-aarch32-and-aarch64

"At public interfaces, the alignment of sp must be two times the pointer size. For AArch32 that's 8 bytes, and for AArch64 it's 16 bytes. A "public interface" is typically a function that is visible to some other, separately-compiled code. The exact definition depends upon the language and the toolchain, and is out of scope of this article. It's reasonable to assume that any C or C++ functions that you interact with using assembly are treated as public interfaces."

https://learn.microsoft.com/en-us/cpp/build/arm64-windows-abi-conventions

"Following the ABI put forth by ARM, the stack must always remain 16-byte aligned"

Thank you for sharing.

a-sidorova · 2025-07-18T06:58:01Z

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

+
+namespace ov::intel_cpu::aarch64::utils {
+
+size_t get_buffer_cluster_id(const ov::snippets::lowered::ExpressionPort& port) {


If this helper for aarch64 and helper for x64 are the same, what's about to move them into one common file in emitters/snippets/utils/util.* for example?
These helpers aren't depended on arch so can be reused on all platforms

a-sidorova · 2025-07-18T07:01:32Z

src/plugins/intel_cpu/tests/functional/shared_tests_instances/skip_tests_config.cpp

+    // TODO: support for long offsets is required for MatMulBias tests with bigger shapes
+    retVector.emplace_back(R"(smoke_Snippets_MatMulBias.*1023.*)");


Is it functional problem or limitation? Can we handle it in this PR?

As we discussed offline, this problem should be resolved in this PR. Let's do our best to find root cause and fix it 😊

Unfortunately, no

[ RUN ] smoke_Snippets_MatMulBias/MatMulBias.CompareWithRefImpl/IS[0]=[]_([1.2.95.1023])_IS[1]=[]_([1.2.1023.255])_IS[2]=[]_([1.2.95.255])_T[0]=f32_T[1]=f32_MatMul_#N=1_#S=1_targetDevice=CPU MEM_USAGE=5129428KB bad err=15 in Xbyak::Error src/tests/functional/base_func_tests/src/base/ov_subgraph.cpp:97: Failure Exception from src/inference/src/cpp/core.cpp:112: Exception from src/inference/src/dev/plugin.cpp:53: illegal immediate parameter (range error) unknown file: Failure C++ exception with description "Exception from src/inference/src/cpp/compiled_model.cpp:56: CompiledModel was not initialized. " thrown in the test body.

a-sidorova · 2025-07-22T07:01:16Z

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

+    }
+
+    // For non-zero offsets, apply the offset and then store
+    Xbyak_aarch64::XReg temp_reg(4);


What's is the magic number 4? What's happened if ptr_reg is XReg(4) too? We can corrupt initial value in XReg(4).

To solve this problem, can we pass temp_reg as the argument of the helper as it's done in push_ptr_with_runtime_offset_on_stack?

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

a-sidorova · 2025-07-22T07:06:32Z

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

+
+    // For large offsets, use a register to hold the offset value
+    if (ptr_offset > 4095) {  // 12-bit immediate limit for add instruction
+        Xbyak_aarch64::XReg offset_reg(6);


The same comment: we should avoid initialization of register by explicit numbers in any helpers since they might be called in any code parts. This is unsafe way if somewhere XReg(6) is already live register during push_ptr_with_static_offset_on_stack call

Please avoid any initialization of the register by explicit number in helpers

Or can we just use TMP_X_0 for example here?

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

a-sidorova · 2025-07-22T07:50:45Z

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

+                                           size_t runtime_offset);
+
+/**
+ * @brief Push data pointer on stack adding static offset `ptr_offset`


Please add the comment that these helpers don't allocate stack pointer - the user should guarantee allocated space on stack

I think this comment should be in the both helpers: push_ptr_with_runtime_offset_on_stack and push_ptr_with_static_offset_on_stack :)

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

aobolensk

Applied comments

aobolensk · 2025-07-22T13:43:58Z

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

+    }
+
+    // For non-zero offsets, apply the offset and then store
+    Xbyak_aarch64::XReg temp_reg(4);


aobolensk · 2025-07-22T13:44:01Z

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

+
+    // For large offsets, use a register to hold the offset value
+    if (ptr_offset > 4095) {  // 12-bit immediate limit for add instruction
+        Xbyak_aarch64::XReg offset_reg(6);


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

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

a-sidorova · 2025-07-23T06:22:41Z

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

+    default:
+        OV_CPU_JIT_EMITTER_THROW("Uknown type of expression port!");
+    }
+    OV_CPU_JIT_EMITTER_ASSERT(IMPLICATION(ov::snippets::utils::is_dynamic_value(offset), id != SIZE_MAX),


Minor:
we can reuse existing implication helper here.

a-sidorova · 2025-07-23T06:23:32Z

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

+                                           size_t runtime_offset);
+
+/**
+ * @brief Push data pointer on stack adding static offset `ptr_offset`


I think this comment should be in the both helpers: push_ptr_with_runtime_offset_on_stack and push_ptr_with_static_offset_on_stack :)

a-sidorova · 2025-07-23T06:25:43Z

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

+
+    // Load the runtime offset from abi_param1 (X0) and add it to the pointer
+    Xbyak_aarch64::XReg abi_param1(0);
+    Xbyak_aarch64::XReg offset_reg(4);


I still see magic XReg(4) here. Please avoid any initializations of registers by const numbers

a-sidorova · 2025-07-23T06:26:10Z

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

+    h->mov(temp_reg, ptr_reg);
+    h->add_imm(temp_reg, temp_reg, ptr_offset, Xbyak_aarch64::XReg(h->X_TMP_1));


Suggested change

h->mov(temp_reg, ptr_reg);

h->add_imm(temp_reg, temp_reg, ptr_offset, Xbyak_aarch64::XReg(h->X_TMP_1));

h->add_imm(temp_reg, ptr_reg, ptr_offset, Xbyak_aarch64::XReg(h->X_TMP_1));

a-sidorova · 2025-07-23T06:29:06Z

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

+
+    // Load the offset value from the runtime parameter location
+    h->add_imm(temp_reg, abi_param1, runtime_offset, Xbyak_aarch64::XReg(h->X_TMP_1));
+    h->ldr(offset_reg, Xbyak_aarch64::ptr(temp_reg));


Is it possible to replace offset_reg with temp_reg here?

a-sidorova · 2025-07-23T06:33:09Z

src/plugins/intel_cpu/tests/functional/shared_tests_instances/skip_tests_config.cpp

+    // TODO: support for long offsets is required for MatMulBias tests with bigger shapes
+    retVector.emplace_back(R"(smoke_Snippets_MatMulBias.*1023.*)");


As we discussed offline, this problem should be resolved in this PR. Let's do our best to find root cause and fix it 😊

aobolensk assigned a-sidorova Jul 17, 2025

aobolensk requested review from a team as code owners July 17, 2025 16:56

aobolensk added the platform: arm OpenVINO on ARM / ARM64 label Jul 17, 2025

github-actions bot added the category: CPU OpenVINO CPU plugin label Jul 17, 2025

aobolensk force-pushed the snippets-arm-matmul-offsets branch from cc4e8fb to c1ddb94 Compare July 17, 2025 17:01

[Snippets][CPU] Support static and dynamic offsets in JIT Gemm and Ge…

c1ddb94

…mmCopyB emitters

a-sidorova assigned chenhu-wang Jul 18, 2025

chenhu-wang reviewed Jul 21, 2025

View reviewed changes

aobolensk force-pushed the snippets-arm-matmul-offsets branch from 744ad13 to 36c75d7 Compare July 21, 2025 06:16

chenhu-wang approved these changes Jul 21, 2025

View reviewed changes

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

aobolensk added 3 commits July 21, 2025 09:14

Address review comments

36c75d7

gpr length

7dc7b92

Remove redundant comments

8031b12

Revert "gpr length"

343629e

This reverts commit 7dc7b92.

a-sidorova reviewed Jul 22, 2025

View reviewed changes

Move utils

2d3d563

aobolensk commented Jul 22, 2025

View reviewed changes

aobolensk force-pushed the snippets-arm-matmul-offsets branch 3 times, most recently from 1637662 to 9cef4e0 Compare July 22, 2025 14:22

aobolensk added 4 commits July 22, 2025 16:43

Address comments

6ce8366

fmt & tidy

9cef4e0

tidy 2

f35a155

x64 tidy

6aac927

a-sidorova reviewed Jul 23, 2025

View reviewed changes

Address utils related review comments

71c987e


		namespace ov::intel_cpu::aarch64::utils {

		size_t get_buffer_cluster_id(const ov::snippets::lowered::ExpressionPort& port) {

		// TODO: support for long offsets is required for MatMulBias tests with bigger shapes
		retVector.emplace_back(R"(smoke_Snippets_MatMulBias.1023.)");

		h->mov(temp_reg, ptr_reg);
		h->add_imm(temp_reg, temp_reg, ptr_offset, Xbyak_aarch64::XReg(h->X_TMP_1));

	h->mov(temp_reg, ptr_reg);
	h->add_imm(temp_reg, temp_reg, ptr_offset, Xbyak_aarch64::XReg(h->X_TMP_1));
	h->add_imm(temp_reg, ptr_reg, ptr_offset, Xbyak_aarch64::XReg(h->X_TMP_1));

[Snippets][CPU] Support static and dynamic offsets in JIT Gemm and GemmCopyB emitters #31375

Are you sure you want to change the base?

[Snippets][CPU] Support static and dynamic offsets in JIT Gemm and GemmCopyB emitters #31375

Conversation

aobolensk commented Jul 17, 2025

Details:

Tickets:

Uh oh!

a-sidorova commented Jul 18, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chenhu-wang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

aobolensk commented Jul 21, 2025

Uh oh!

chenhu-wang commented Jul 22, 2025

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!

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!

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!

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!

aobolensk left a comment

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!

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!

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!

Choose a reason for hiding this comment