@@ -191,7 +191,7 @@ class MHABF16AMXBufferAllocationTest : public BufferAllocationCPUTest {
191
191
std::shared_ptr<ov::Model> GetModel (const std::vector<ov::PartialShape>& shapes) const override {
192
192
const auto subtensor_scalar = std::vector<size_t >{1 };
193
193
const auto subtensor_power = std::vector<size_t >{1 , ov::snippets::utils::get_full_dim_value ()};
194
- const auto subtensor_full = std::vector<size_t >(2 , ov::snippets::utils::get_full_dim_value ());
194
+ const auto subtensor_flat = std::vector<size_t >(1 , ov::snippets::utils::get_full_dim_value ());
195
195
196
196
OPENVINO_ASSERT (shapes.size () == 3 , " Incorrect count of input shapes" );
197
197
const auto parameter0 = std::make_shared<ov::op::v0::Parameter>(ov::element::bf16, shapes[0 ]);
@@ -238,12 +238,12 @@ class MHABF16AMXBufferAllocationTest : public BufferAllocationCPUTest {
238
238
MarkOp (store, subtensor_scalar);
239
239
MarkOp (power, subtensor_power);
240
240
241
- MarkOp (brgemm_cpu0, subtensor_full );
242
- MarkOp (brgemm_cpu1, subtensor_full );
243
- MarkOp (brgemm_copyb0, subtensor_full );
244
- MarkOp (brgemm_copyb1, subtensor_full );
245
- MarkOp (scratch0, subtensor_full );
246
- MarkOp (scratch1, subtensor_full );
241
+ MarkOp (brgemm_cpu0, subtensor_flat );
242
+ MarkOp (brgemm_cpu1, subtensor_flat );
243
+ MarkOp (brgemm_copyb0, subtensor_flat );
244
+ MarkOp (brgemm_copyb1, subtensor_flat );
245
+ MarkOp (scratch0, subtensor_flat );
246
+ MarkOp (scratch1, subtensor_flat );
247
247
248
248
ov::snippets::lowered::PortDescriptorUtils::get_port_descriptor_ptr (load_reshape->input (0 ))->set_layout (order);
249
249
0 commit comments