From 85b3a55316df050ff17c1ca417ea988cc5ba9c4e Mon Sep 17 00:00:00 2001
From: Polykarpos Thomadakis <polykarpos.thomadakis@pnnl.gov>
Date: Wed, 16 Apr 2025 17:34:16 -0700
Subject: [PATCH 1/2] linalg.matvec will now lower to parallel loops

---
 frontends/comet_dsl/comet.cpp                 |  2 +
 .../Transforms/LinalgTransforms.cpp           | 60 +++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/frontends/comet_dsl/comet.cpp b/frontends/comet_dsl/comet.cpp
index b7d4d392..7ad5da23 100644
--- a/frontends/comet_dsl/comet.cpp
+++ b/frontends/comet_dsl/comet.cpp
@@ -586,6 +586,8 @@ int loadAndProcessMLIR(mlir::MLIRContext &context,
   if (OptCallToMatMulMicroKernel)
   {
     pm.addNestedPass<mlir::func::FuncOp>(mlir::comet::createLinAlgMatmulMicroKernelPass());
+    pm.addNestedPass<mlir::func::FuncOp>(mlir::comet::createMatvecToParallelLoopsPass());
+
   }
   pm.addNestedPass<mlir::func::FuncOp>(mlir::createConvertVectorToSCFPass());
   /// Blanket-convert any remaining linalg ops to loops if any remain.
diff --git a/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp b/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp
index b233d2d1..e9853379 100644
--- a/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp
+++ b/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp
@@ -36,6 +36,7 @@
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -51,6 +52,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/SmallVector.h"
+#include <memory>
 #include <optional>
 
 // suppress all warnings coming from inclusion of blis.h in source tree
@@ -810,6 +812,57 @@ namespace
     uint64_t tile_size;
     bool seperate_tiles;
   };
+} /// end anonym
+// ous namespace
+namespace
+{
+  struct MatvecToParallelLoops : public ConversionPattern
+  {
+    MatvecToParallelLoops(MLIRContext *ctx)
+        : ConversionPattern(linalg::MatvecOp::getOperationName(), 1, ctx)
+           {}
+    
+    LogicalResult
+    matchAndRewrite(Operation *input_op, ArrayRef<Value> operands,
+                    ConversionPatternRewriter &rewriter) const final
+    {
+
+      auto op = dyn_cast<linalg::MatvecOp>(input_op);
+
+      if(failed(mlir::linalg::linalgOpToParallelLoops(rewriter, op)))
+      {
+        return mlir::failure();
+      }
+      else{
+        rewriter.eraseOp(input_op);
+        return success();
+      }
+    }
+  };
+  
+  class MatvecToParallelLoopsPass : public PassWrapper<MatvecToParallelLoopsPass, OperationPass<func::FuncOp>>
+  {
+  public:
+    MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MatvecToParallelLoopsPass)
+    MatvecToParallelLoopsPass(){};
+    void runOnOperation() override
+    {
+      comet_debug() << "MatvecToParallelLoopsPass : public PassWrapper<MatvecToParallelLoopsPass, FunctionPass>\n";
+      func::FuncOp func = getOperation();
+      ConversionTarget target(getContext());
+      target.addLegalDialect<TADialect, ArithDialect, scf::SCFDialect, AffineDialect, memref::MemRefDialect>();
+      RewritePatternSet patterns(&getContext());
+      patterns.insert<MatvecToParallelLoops>(&getContext()  );
+
+      if (failed(applyPartialConversion(func, target, std::move(patterns))))
+      {
+        llvm::errs() << "Failed to Lower dense transpose operation\n";
+        signalPassFailure();
+      }
+      comet_debug() << "MatvecToParallelLoopsPass done\n";
+    }
+
+  };
 } /// end anonymous namespace
 
 /// Create a pass to optimize LinAlg Matmul Op with tiling
@@ -833,3 +886,10 @@ std::unique_ptr<mlir::Pass> mlir::comet::createOptDenseTransposePass(uint64_t ti
   comet_debug() << "LinAlgTransforms createOptDenseTransposePass\n";
   return std::make_unique<OptDenseTransposePass>(tile_size, seperate_tiles);
 }
+
+
+std::unique_ptr<mlir::Pass> mlir::comet::createMatvecToParallelLoopsPass()
+{
+  comet_debug() << "LinAlgTransforms createMatvecToParallelLoopsPass\n";
+  return std::make_unique<MatvecToParallelLoopsPass>();
+}
\ No newline at end of file

From e05a5e176ead8d8243313d19296d3d9682355d59 Mon Sep 17 00:00:00 2001
From: Polykarpos Thomadakis <polykarpos.thomadakis@pnnl.gov>
Date: Wed, 16 Apr 2025 17:44:32 -0700
Subject: [PATCH 2/2] Added missing change

---
 include/comet/Dialect/TensorAlgebra/Passes.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/include/comet/Dialect/TensorAlgebra/Passes.h b/include/comet/Dialect/TensorAlgebra/Passes.h
index 06d828bd..f21562e2 100644
--- a/include/comet/Dialect/TensorAlgebra/Passes.h
+++ b/include/comet/Dialect/TensorAlgebra/Passes.h
@@ -70,6 +70,8 @@ namespace mlir
 
         std::unique_ptr<Pass> createLinAlgMatmulTilingPass();
         std::unique_ptr<Pass> createLinAlgMatmulMicroKernelPass();
+        std::unique_ptr<Pass> createMatvecToParallelLoopsPass();
+
 
         // Optimize dense transpose (linalg.copy) based on the following paper:
         // HPTT: A High-Performance Tensor Transposition C++ Library