[CPU] Add a sanity check to MM shape infer

maxnick · Mar 21, 2024 · eb548da · eb548da
1 parent 9072ea7
commit eb548da
Showing 1 changed file with 10 additions and 0 deletions.
diff --git 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];