[aievec] Generalize vector passes

Right now, vectorization passes hook to FuncOp, which prevents conversion to AIEVec within other top level operations, like AIE.device ops. This patch makes all passes generic and allows for conversion within AIE.device.
Xilinx · Sep 15, 2023 · f90379d · f90379d
1 parent f21d029
commit f90379d
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 39 deletions.
diff --git a/include/aie/Dialect/AIEVec/Analysis/Passes.td b/include/aie/Dialect/AIEVec/Analysis/Passes.td
@@ -15,7 +15,7 @@
 
 include "mlir/Pass/PassBase.td"
 
-def AIEVecConvAnalysis : Pass<"aievec-convolution-analysis", "mlir::func::FuncOp"> {
+def AIEVecConvAnalysis : Pass<"aievec-convolution-analysis"> {
   let summary = "Find MAC chains that can be replaced by convolution ops in "
                 "AIE-ML";
   let constructor = "xilinx::aievec::createAIEVecConvolutionAnalysisPass()";

diff --git a/lib/Dialect/AIEVec/Transforms/AIEVecOptimizations.cpp b/lib/Dialect/AIEVec/Transforms/AIEVecOptimizations.cpp
@@ -234,8 +234,7 @@ configureAIEVecV2TransformationLegalizations(ConversionTarget &target) {
 // Lowering passes
 //===----------------------------------------------------------------------===//
 struct AIEVecTransformationPass
-    : public PassWrapper<AIEVecTransformationPass,
-                         OperationPass<func::FuncOp>> {
+    : public PassWrapper<AIEVecTransformationPass, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AIEVecTransformationPass)
 
   AIEVecTransformationPass() = default;
@@ -267,7 +266,7 @@ struct AIEVecTransformationPass
       llvm::cl::init("aie")};
 
   void runOnOperation() override {
-    auto func = getOperation();
+    auto op = getOperation();
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
@@ -277,7 +276,7 @@ struct AIEVecTransformationPass
       if (target == "aieml") {
         aieVersion = AIEArch::AIE_ML;
       } else if (target != "aie") {
-        func.emitError() << "unknown AIE target '" << aieTarget << "'";
+        op->emitError() << "unknown AIE target '" << aieTarget << "'";
         signalPassFailure();
         return;
       }
@@ -290,7 +289,7 @@ struct AIEVecTransformationPass
       configureAIEVecV2TransformationLegalizations(target);
     }
 
-    if (failed(applyPartialConversion(func, target, std::move(patterns)))) {
+    if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
       signalPassFailure();
     }
   }
@@ -302,8 +301,7 @@ createAIEVecTransformationPass(const OptimizeAIEVecOptions &options) {
 }
 
 struct AIEVecConvOpTransformationPass
-    : public PassWrapper<AIEVecConvOpTransformationPass,
-                         OperationPass<func::FuncOp>> {
+    : public PassWrapper<AIEVecConvOpTransformationPass, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AIEVecConvOpTransformationPass)
 
   AIEVecConvOpTransformationPass() = default;
@@ -343,7 +341,7 @@ struct AIEVecConvOpTransformationPass
       llvm::cl::init(0)};
 
   void runOnOperation() override {
-    auto func = getOperation();
+    auto op = getOperation();
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
@@ -353,7 +351,7 @@ struct AIEVecConvOpTransformationPass
       if (target == "aieml") {
         aieVersion = AIEArch::AIE_ML;
       } else if (target != "aie") {
-        func.emitError() << "unknown AIE target '" << aieTarget << "'";
+        op->emitError() << "unknown AIE target '" << aieTarget << "'";
         signalPassFailure();
         return;
       }
@@ -365,7 +363,7 @@ struct AIEVecConvOpTransformationPass
       configureAIEVecConvOpTransformationLegalizations(target, am);
     }
 
-    if (failed(applyPartialConversion(func, target, std::move(patterns)))) {
+    if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
       signalPassFailure();
     }
   }

diff --git a/lib/Dialect/AIEVec/Transforms/ConvertVectorToAIEVec.cpp b/lib/Dialect/AIEVec/Transforms/ConvertVectorToAIEVec.cpp
@@ -82,19 +82,19 @@ using SetInboundsToWriteOp = SetInboundsToReadStoreOpPattern<TransferWriteOp>;
 //===----------------------------------------------------------------------===//
 
 struct RedundantLoadStoreOptimizationPass
-    : public PassWrapper<RedundantLoadStoreOptimizationPass,
-                         OperationPass<func::FuncOp>> {
+    : public PassWrapper<RedundantLoadStoreOptimizationPass, OperationPass<>> {
+
   void runOnOperation() override {
-    func::FuncOp funcOp = getOperation();
+    auto op = getOperation();
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
 
     patterns.add<SetInboundsToReadOp, SetInboundsToWriteOp>(
         patterns.getContext());
 
-    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+    (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
     IRRewriter rewriter(&getContext());
-    vector::transferOpflowOpt(rewriter, funcOp);
+    vector::transferOpflowOpt(rewriter, op);
   }
 };
 

diff --git a/lib/Dialect/AIEVec/Transforms/FoldMulAddChainToConvOp.cpp b/lib/Dialect/AIEVec/Transforms/FoldMulAddChainToConvOp.cpp
@@ -490,16 +490,16 @@ struct AIEVecConvAnalysis : public AIEVecConvAnalysisBase<AIEVecConvAnalysis> {
     markAllAnalysesPreserved();
     AnalysisManager am = getAnalysisManager();
     LongestConvMACChainAnalysis::am = &am;
-    func::FuncOp func = getOperation();
+    Operation *op = getOperation();
 
     // Compute all the chains
-    func.walk([&](arith::AddIOp addOp) {
+    op->walk([&](arith::AddIOp addOp) {
       if (isa<VectorType>(addOp.getResult().getType()))
         am.getChildAnalysis<LongestConvMACChainAnalysis>(addOp);
     });
 
     // Sort the chains, ready to split by group
-    func.walk([&](arith::AddIOp addOp) {
+    op->walk([&](arith::AddIOp addOp) {
       if (isa<VectorType>(addOp.getResult().getType())) {
         auto &analysis =
             am.getChildAnalysis<LongestConvMACChainAnalysis>(addOp);
@@ -509,7 +509,7 @@ struct AIEVecConvAnalysis : public AIEVecConvAnalysisBase<AIEVecConvAnalysis> {
     });
 
     if (printResult) {
-      func.walk([&](arith::AddIOp addOp) {
+      op->walk([&](arith::AddIOp addOp) {
         if (isa<VectorType>(addOp.getResult().getType())) {
           auto &macChainAnalysis =
               am.getChildAnalysis<LongestConvMACChainAnalysis>(addOp);

diff --git a/lib/Dialect/AIEVec/Transforms/VectorToAIEVecConversions.cpp b/lib/Dialect/AIEVec/Transforms/VectorToAIEVecConversions.cpp
@@ -2183,7 +2183,7 @@ static void configureAIEVecV2Legalizations(ConversionTarget &target,
 /// Lower incoming vector operations into their corresponding AIE vector
 /// intrinsics.
 struct LowerVectorToAIEVec
-    : public PassWrapper<LowerVectorToAIEVec, OperationPass<func::FuncOp>> {
+    : public PassWrapper<LowerVectorToAIEVec, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerVectorToAIEVec)
 
   LowerVectorToAIEVec() = default;
@@ -2213,7 +2213,7 @@ struct LowerVectorToAIEVec
       llvm::cl::init("aie")};
 
   void runOnOperation() override {
-    auto func = getOperation();
+    auto op = getOperation();
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
@@ -2223,7 +2223,7 @@ struct LowerVectorToAIEVec
       if (target == "aieml") {
         aieVersion = AIEArch::AIE_ML;
       } else if (target != "aie") {
-        func.emitError() << "unknown AIE target '" << aieTarget << "'";
+        op->emitError() << "unknown AIE target '" << aieTarget << "'";
         signalPassFailure();
         return;
       }
@@ -2239,7 +2239,7 @@ struct LowerVectorToAIEVec
       configureAIEVecV2Legalizations(target, am);
     }
 
-    if (failed(applyPartialConversion(func, target, std::move(patterns)))) {
+    if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
       signalPassFailure();
     }
   }
@@ -2259,7 +2259,8 @@ createLowerVectorToAIEVec(const LowerVectorToAIEVecOptions &options) {
 // additional common subexpressions with UPDs generated from unaligned
 // `transfer_read` ops.
 struct ExtendUPDOpsPass
-    : public PassWrapper<ExtendUPDOpsPass, OperationPass<func::FuncOp>> {
+    : public PassWrapper<ExtendUPDOpsPass, OperationPass<>> {
+
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
@@ -2272,8 +2273,8 @@ struct ExtendUPDOpsPass
              llvm::all_of(op->getUsers(),
                           [](Operation *op) { return isa<aievec::ExtOp>(op); });
     });
-    auto func = getOperation();
-    if (failed(applyPartialConversion(func, target, std::move(patterns)))) {
+    auto op = getOperation();
+    if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
       signalPassFailure();
     }
   }
@@ -2285,7 +2286,8 @@ struct ExtendUPDOpsPass
 // TODO: a single ext op of the top half, which might be a good opportunity to
 // TODO: further optimize wide UPDs.
 struct SimplifyUPDOpsPass
-    : public PassWrapper<SimplifyUPDOpsPass, OperationPass<func::FuncOp>> {
+    : public PassWrapper<SimplifyUPDOpsPass, OperationPass<>> {
+
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
@@ -2297,8 +2299,8 @@ struct SimplifyUPDOpsPass
       return !defOp || !isa<aievec::UPDOp>(defOp) || !defOp->hasOneUse() ||
              op.getIndex() != 0;
     });
-    auto func = getOperation();
-    if (failed(applyPartialConversion(func, target, std::move(patterns)))) {
+    auto op = getOperation();
+    if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
       signalPassFailure();
     }
   }

diff --git a/lib/Dialect/AIEVec/Transforms/VectorToVectorConversions.cpp b/lib/Dialect/AIEVec/Transforms/VectorToVectorConversions.cpp
@@ -328,8 +328,7 @@ populateAIEMLCanonicalizeConversionPatterns(RewritePatternSet &patterns) {
 //    2) Split unaligned transfer reads into a wider aligned transfer read
 //       followed by a `vector.extract_strided_slice` operation.
 struct CanonicalizeVectorForAIEVecPass
-    : public PassWrapper<CanonicalizeVectorForAIEVecPass,
-                         OperationPass<func::FuncOp>> {
+    : public PassWrapper<CanonicalizeVectorForAIEVecPass, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CanonicalizeVectorForAIEVecPass)
 
   CanonicalizeVectorForAIEVecPass() = default;
@@ -364,7 +363,7 @@ struct CanonicalizeVectorForAIEVecPass
       llvm::cl::init("aie")};
 
   void runOnOperation() override {
-    func::FuncOp funcOp = getOperation();
+    auto op = getOperation();
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
@@ -375,7 +374,7 @@ struct CanonicalizeVectorForAIEVecPass
       if (target == "aieml") {
         aieVersion = AIEArch::AIE_ML;
       } else if (target != "aie") {
-        funcOp.emitError() << "unknown AIE target '" << aieTarget << "'";
+        op->emitError() << "unknown AIE target '" << aieTarget << "'";
         signalPassFailure();
         return;
       }
@@ -391,7 +390,7 @@ struct CanonicalizeVectorForAIEVecPass
       configureAIEMLCanonicalizeLegalizations(target);
     }
 
-    if (failed(applyPartialConversion(funcOp, target, std::move(patterns)))) {
+    if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
       signalPassFailure();
     }
   }
@@ -403,16 +402,16 @@ static std::unique_ptr<::mlir::Pass> createCanonicalizeVectorForAIEVecPass(
 }
 
 struct HoistCastOpToDataSourcePass
-    : public PassWrapper<HoistCastOpToDataSourcePass,
-                         OperationPass<func::FuncOp>> {
+    : public PassWrapper<HoistCastOpToDataSourcePass, OperationPass<>> {
+
   void runOnOperation() override {
-    func::FuncOp funcOp = getOperation();
+    auto op = getOperation();
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
 
     patterns.add<HoistCastOpToDataSourcePattern>(patterns.getContext());
 
-    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+    (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
   }
 };