Fixes for brgemm

plaidml · Oct 14, 2024 · ff762d9 · ff762d9
1 parent 53b2c85
commit ff762d9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/lib/TPP/Dialect/Xsmm/XsmmUtils.cpp b/lib/TPP/Dialect/Xsmm/XsmmUtils.cpp
@@ -309,6 +309,7 @@ FailureOr<BrgemmInfo> isMappableToBrgemm(PatternRewriter &rewriter,
   SmallVector<unsigned, 2> kVector;
   std::optional<unsigned> batch;
   if (contractionDims->k.size() >= 2) {
+    batch = contractionDims->k[0];
     for (size_t i = 1; i < contractionDims->k.size(); i++)
       kVector.push_back(contractionDims->k[i]);
   } else {

diff --git a/lib/TPP/Transforms/Vectorization.cpp b/lib/TPP/Transforms/Vectorization.cpp
@@ -111,7 +111,7 @@ struct VectorizationPass
     patterns.add<
         LinalgToVector<linalg::BatchReduceMatmulOp>,
         LinalgToVector<linalg::FillOp>, LinalgToVector<linalg::TransposeOp>,
-        LinalgToVector<linalg::BroadcastOp>, LinalgToVector<linalg::CopyOp>>(
+        LinalgToVector<linalg::BroadcastOp>, LinalgToVector<linalg::MatmulOp>, LinalgToVector<linalg::CopyOp>>(
         patterns.getContext());
     patterns.add<LinalgGenericToVector>(patterns.getContext());
   }