Merge pull request #155 from ShangkunLi/enable-promote-funargs-in

Explictly use promote-func-arg-to-constant pass

Author: ShangkunLi

include/NeuraDialect/NeuraPasses.td

@@ -81,7 +81,6 @@ def CanonicalizeLiveIn : Pass<"canonicalize-live-in", "ModuleOp"> {
     This pass applies canonicalization transformations to Neura dialect operations.
     The canonicalization includes:
     1. Converting all live-in values of each basic block to block arguments.
-    2. Promoting function arguments to neura constant operations.
   }];
   let constructor = "neura::createCanonicalizeLiveInPass()";
 }

lib/NeuraDialect/NeuraPasses.cpp

@@ -26,6 +26,8 @@ void mlir::neura::registerNeuraConversionPassPipeline() {
         pm.addPass(mlir::createPrintOpGraphPass(os));
 
         pm.addPass(mlir::neura::createCanonicalizeCastPass());
+        pm.addPass(mlir::neura::createPromoteFuncArgToConstPass());
+        pm.addPass(mlir::neura::createFoldConstantPass());
         pm.addPass(mlir::neura::createCanonicalizeLiveInPass());
         pm.addPass(mlir::neura::createLeveragePredicatedValuePass());
         pm.addPass(mlir::createPrintOpGraphPass(os));

lib/NeuraDialect/Transforms/CanonicalizeLiveInPass.cpp

@@ -17,29 +17,6 @@ using namespace mlir;
 #include "NeuraDialect/NeuraPasses.h.inc"
 
 namespace {
-LogicalResult promoteFunctionArgsToConstants(Region &region) {
-  if (region.empty()) {
-    return success();
-  }
-
-  Block &entry_block = region.front();
-  OpBuilder builder(&entry_block, entry_block.begin());
-
-  // Collects all function arguments.
-  SmallVector<BlockArgument, 4> args(entry_block.getArguments().begin(),
-                                     entry_block.getArguments().end());
-
-  // Creates a constant operation for each function argument.
-  for (auto [idx, arg] : llvm::enumerate(args)) {
-    auto const_op = builder.create<neura::ConstantOp>(
-        arg.getLoc(), arg.getType(),
-        builder.getStringAttr("\%arg" + std::to_string(idx)));
-    arg.replaceAllUsesWith(const_op.getResult());
-  }
-
-  return success();
-}
-
 LogicalResult promoteLiveInValuesToBlockArgs(Region &region) {
   if (region.empty()) {
     return success();
@@ -286,8 +263,8 @@ LogicalResult promoteLiveInValuesToBlockArgs(Region &region) {
         if (needs_update) {
           OpBuilder builder(cond_br_op);
           builder.create<neura::CondBr>(
-              cond_br_op.getLoc(), cond_br_op.getCondition(),
-              true_operands, false_operands, cond_br_op.getTrueDest(),
+              cond_br_op.getLoc(), cond_br_op.getCondition(), true_operands,
+              false_operands, cond_br_op.getTrueDest(),
               cond_br_op.getFalseDest());
           cond_br_op.erase();
         }
@@ -336,11 +313,6 @@ struct CanonicalizeLiveInPass
         return;
       }
 
-      if (failed(promoteFunctionArgsToConstants(*region))) {
-        signalPassFailure();
-        return;
-      }
-
       if (failed(promoteLiveInValuesToBlockArgs(*region))) {
         signalPassFailure();
         return;

lib/NeuraDialect/Transforms/Optimizations/HwAgnosticOpt/FoldConstantPass.cpp

@@ -111,7 +111,7 @@ struct FuseRhsConstantPattern : public OpRewritePattern<OpType> {
     Value rhs = op.getRhs();
 
     if (isOriginConstantOp(lhs)) {
-      assert(false && "LHS constant folding not implemented yet.");
+      llvm::errs() << "LHS constant folding not supported yet.\n";
       return failure();
     }
 

test/affine2neura/bert/bert_node1/bert_node1.mlir

@@ -14,6 +14,7 @@
 // RUN: --lower-builtin-to-neura \
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
+// RUN: --promote-func-arg-to-const \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \

test/affine2neura/bert/bert_node28/bert_node28.mlir

@@ -14,6 +14,7 @@
 // RUN: --lower-builtin-to-neura \
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
+// RUN: --promote-func-arg-to-const \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \

test/c2llvm2mlir/nested_loop/test.mlir

@@ -3,6 +3,7 @@
 
 // RUN: mlir-neura-opt --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
+// RUN:   --promote-func-arg-to-const \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -11,6 +12,7 @@
 
 // RUN: mlir-neura-opt --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
+// RUN:   --promote-func-arg-to-const \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \

test/controflow_fuse/complex_nested/complex_nested.mlir

@@ -14,6 +14,7 @@
 // RUN: --lower-builtin-to-neura \
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
+// RUN: --promote-func-arg-to-const \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \

test/controflow_fuse/non_perfect_nested/non_perfect_nested.mlir

@@ -14,6 +14,7 @@
 // RUN: --lower-builtin-to-neura \
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
+// RUN: --promote-func-arg-to-const \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \

test/controflow_fuse/perfect_nested/perfect_nested.mlir

@@ -23,6 +23,7 @@
 // RUN: --lower-builtin-to-neura \
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
+// RUN: --promote-func-arg-to-const \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -35,6 +36,7 @@
 // RUN: --lower-builtin-to-neura \
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
+// RUN: --promote-func-arg-to-const \
 // RUN: --fold-constant \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \