diff --git a/lib/Transforms/InsertMovPass.cpp b/lib/Transforms/InsertMovPass.cpp index 5778a381..c3476b60 100644 --- a/lib/Transforms/InsertMovPass.cpp +++ b/lib/Transforms/InsertMovPass.cpp @@ -48,7 +48,15 @@ struct InsertMovForNeuraOps : public RewritePattern { state.addTypes(op->getResultTypes()); state.addAttributes(op->getAttrs()); + // Copies successors for terminator operations. + if (op->hasTrait()) { + for (Block *successor : op->getSuccessors()) { + state.addSuccessors(successor); + } + } + Operation *newOp = rewriter.create(state); + rewriter.replaceOp(op, newOp->getResults()); return success(); } diff --git a/test/neura/for_loop/test.mlir b/test/neura/for_loop/test.mlir index 1c72782d..c3ed9b81 100644 --- a/test/neura/for_loop/test.mlir +++ b/test/neura/for_loop/test.mlir @@ -8,7 +8,7 @@ // RUN: --assign-accelerator \ // RUN: --lower-llvm-to-neura \ // RUN: --fuse-patterns \ -// RN: --insert-mov \ +// RUN: --insert-mov \ // RUN: %t-kernel.mlir | FileCheck %s // Verifies the neura ops are generated. And fusion happens.