From eb548dab4d14c5249b0ad3910524cd469ea7e581 Mon Sep 17 00:00:00 2001 From: Maksim Kutakov Date: Thu, 7 Mar 2024 16:22:53 +0100 Subject: [PATCH] [CPU] Add a sanity check to MM shape infer --- .../intel_cpu/src/shape_inference/custom/matmul.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/plugins/intel_cpu/src/shape_inference/custom/matmul.cpp b/src/plugins/intel_cpu/src/shape_inference/custom/matmul.cpp index c1959933a4e5f0..ac1aa0319386bb 100644 --- a/src/plugins/intel_cpu/src/shape_inference/custom/matmul.cpp +++ b/src/plugins/intel_cpu/src/shape_inference/custom/matmul.cpp @@ -29,6 +29,16 @@ Result MMShapeInfer::infer( return {{m_shapeY}, ShapeInferStatus::success}; } OPENVINO_ASSERT(m_out_rank >= 2, "The output rank should be greater or euqal to 2."); + const size_t k_lhs = m_transpose_a ? shapeA[rankA-2] : shapeA[rankA-1]; + const size_t k_rhs = m_transpose_b ? shapeB[rankB-1] : shapeB[rankB-2]; + OPENVINO_ASSERT(k_lhs == k_rhs, + "Matmul input shapes are incompatible shape A: ", + vec2str(shapeA), + m_transpose_a ? "T " : " ", + "shape B: ", + vec2str(shapeB), + m_transpose_b ? "T" : ""); + m_shapeY[m_out_rank-2] = m_transpose_a ? shapeA[rankA-1] : shapeA[rankA-2]; m_shapeY[m_out_rank-1] = m_transpose_b ? shapeB[rankB-2] : shapeB[rankB-1];