Fix error by checking for m.size()

Signed-off-by: Bruce Lai <[email protected]>
iree-org · Nov 20, 2024 · 846ff1d · 846ff1d
1 parent f74ee00
commit 846ff1d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -1334,6 +1334,10 @@ getMatmulRISCVVectorSizes(mlir::FunctionOpInterface entryPointFn,
     // TODO: support int data type
     return;
   }
+  FailureOr<linalg::ContractionDimensions> cDims =
+      linalg::inferContractionDims(op);
+  if (failed(cDims) || cDims->m.size() != 1)
+    return;
   // Use 7 x lmul4 to fully utilize vector registers.
   sizes[0] = 7;
   // Calculate tile size for the main vector dimension (N).
@@ -1342,10 +1346,6 @@ getMatmulRISCVVectorSizes(mlir::FunctionOpInterface entryPointFn,
       (nativeVectorSize * 2 * kByteSizeInBits) / elementSize;
   sizes[1] = maxNumberElementsForLMUL4;
   sizes[2] = 1;
-  FailureOr<linalg::ContractionDimensions> cDims =
-      linalg::inferContractionDims(op);
-  if (failed(cDims))
-    return;
   ArrayRef<int64_t> lhsShape = op.getShape(op.getDpsInputOperand(0));
   // If m = 1, set tile size to 1 x lmul8
   if (lhsShape[cDims->m[0]] == 1) {