feat: add useMatmulForSingleBatch option to mlir::tosa::addTosaToLinalgPasses.

Author: ttjost

mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h

@@ -25,14 +25,16 @@ namespace mlir {
 namespace tosa {
 
 std::unique_ptr<Pass> createTosaToLinalg();
-std::unique_ptr<Pass> createTosaToLinalgNamed();
+std::unique_ptr<Pass>
+createTosaToLinalgNamed(bool useMatmulForSingleBatch = false);
 
 /// Populates passes to convert from TOSA to Linalg on buffers. At the end of
 /// the pass, the function will only contain linalg ops or standard ops if the
 /// pipeline succeeds.  The option to disable decompositions is available for
 /// benchmarking performance improvements from the canonicalizations.
 void addTosaToLinalgPasses(OpPassManager &pm,
-                           bool disableTosaDecompositions = false);
+                           bool disableTosaDecompositions = false,
+                           bool useMatmulForSingleBatch = false);
 
 /// Populates conversion passes from TOSA dialect to Linalg dialect.
 void populateTosaToLinalgConversionPatterns(RewritePatternSet *patterns);

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp

@@ -37,6 +37,11 @@ namespace {
 struct TosaToLinalgNamed
     : public impl::TosaToLinalgNamedBase<TosaToLinalgNamed> {
 public:
+  TosaToLinalgNamed() = default;
+  explicit TosaToLinalgNamed(bool useMatmulForSingleBatch) {
+    this->useMatmulForSingleBatch = useMatmulForSingleBatch;
+  }
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<arith::ArithDialect, linalg::LinalgDialect, math::MathDialect,
@@ -69,6 +74,7 @@ struct TosaToLinalgNamed
 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::tosa::createTosaToLinalgNamed() {
-  return std::make_unique<TosaToLinalgNamed>();
+std::unique_ptr<Pass>
+mlir::tosa::createTosaToLinalgNamed(bool useMatmulForSingleBatch) {
+  return std::make_unique<TosaToLinalgNamed>(useMatmulForSingleBatch);
 }

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp

@@ -75,15 +75,17 @@ std::unique_ptr<Pass> mlir::tosa::createTosaToLinalg() {
 }
 
 void mlir::tosa::addTosaToLinalgPasses(OpPassManager &pm,
-                                       bool disableTosaDecompositions) {
+                                       bool disableTosaDecompositions,
+                                       bool useMatmulForSingleBatch) {
   // Optional decompositions are designed to benefit linalg.
   if (!disableTosaDecompositions)
     pm.addNestedPass<func::FuncOp>(tosa::createTosaOptionalDecompositions());
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
 
   pm.addNestedPass<func::FuncOp>(tosa::createTosaInferShapesPass());
   pm.addNestedPass<func::FuncOp>(tosa::createTosaMakeBroadcastablePass());
-  pm.addNestedPass<func::FuncOp>(tosa::createTosaToLinalgNamed());
+  pm.addNestedPass<func::FuncOp>(
+      tosa::createTosaToLinalgNamed(useMatmulForSingleBatch));
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   // TODO: Remove pass that operates on const tensor and enable optionality
   pm.addNestedPass<func::FuncOp>(tosa::createTosaLayerwiseConstantFoldPass());