From d176b84506c47b5500230a16b862c062fd522af1 Mon Sep 17 00:00:00 2001 From: Doru Bercea Date: Fri, 10 Jan 2020 11:22:41 -0500 Subject: [PATCH] Add support for broadcasting right matrix. --- src/dialect/onnx/onnx_ops.cpp | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/src/dialect/onnx/onnx_ops.cpp b/src/dialect/onnx/onnx_ops.cpp index 2380f5a..f19242b 100644 --- a/src/dialect/onnx/onnx_ops.cpp +++ b/src/dialect/onnx/onnx_ops.cpp @@ -329,14 +329,31 @@ void ONNXMatMulOp::inferShapes() { // Special case when both arrays are 1-dimensional and according to // numpy rules the types need to be extended to 1xN and Nx1. Helper sizes // need to be removed after the multiplication but cannot be removed if all - // remaining sizes are 1. + // sizes are 1. if (lhsShape[0] != -1 && rhsShape[0] != -1 && lhsShape[0] != rhsShape[0]) emitError("Attempt to multiply incompatible matrices."); dims.emplace_back(1); } else { - dims.emplace_back(lhsShape.size() == 1 ? 1 : lhsShape[0]); - dims.emplace_back(rhsShape.size() == 1 ? 1 : rhsShape[1]); + // Special cases for when at least one matrix has more than two dimensions. + if (lhsShape.size() > 2 && rhsShape.size() == 2) { + // (s1 x s2 x... x sKx x M x N) MATMUL (N x P) + // => + // (s1 x s2 x... x sKx x M x P) + + // Check legality of matrix multiplication. + unsigned leftDims = lhsShape.size(); + if (lhsShape[leftDims - 2] != -1 && rhsShape[0] != -1 && + lhsShape[leftDims - 2] != rhsShape[0]) + emitError("Attempt to multiply incompatible matrices."); + + for (int i = 0; i < leftDims - 1; ++i) + dims.emplace_back(lhsShape[i]); + dims.emplace_back(rhsShape[1]); + } else { + dims.emplace_back(lhsShape.size() == 1 ? 1 : lhsShape[0]); + dims.emplace_back(rhsShape.size() == 1 ? 1 : rhsShape[1]); + } } getResult()->setType(RankedTensorType::get(dims, lhsTy.getElementType()));