From 380bb441e7465052762e083da5d5c230205684f0 Mon Sep 17 00:00:00 2001 From: tancheng Date: Tue, 13 May 2025 16:43:11 +0000 Subject: [PATCH 1/3] [feature] Declare and define mov --- include/NeuraDialect/NeuraOps.td | 10 +++ include/Transforms/InsertMovPass.h | 13 ++++ lib/CMakeLists.txt | 1 + lib/Conversion/ArithToNeura/ArithToNeura.cpp | 7 -- lib/Transforms/CMakeLists.txt | 11 +++ lib/Transforms/InsertMovPass.cpp | 82 ++++++++++++++++++++ test/neura/add.mlir | 9 +++ tools/mlir-neura-opt/CMakeLists.txt | 1 + tools/mlir-neura-opt/mlir-neura-opt.cpp | 4 + 9 files changed, 131 insertions(+), 7 deletions(-) create mode 100644 include/Transforms/InsertMovPass.h create mode 100644 lib/Transforms/CMakeLists.txt create mode 100644 lib/Transforms/InsertMovPass.cpp create mode 100644 test/neura/add.mlir diff --git a/include/NeuraDialect/NeuraOps.td b/include/NeuraDialect/NeuraOps.td index 0e90ae17..ac83aa1e 100644 --- a/include/NeuraDialect/NeuraOps.td +++ b/include/NeuraDialect/NeuraOps.td @@ -11,3 +11,13 @@ def Neura_AddOp : Op { let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)"; // let traits = [Pure]; } + +// Defines a move operation for data communication. +def Neura_MovOp : Op { + let summary = "Move operation"; + let opName = "mov"; + let arguments = (ins AnyType:$lhs); + let results = (outs AnyType:$result); + let assemblyFormat = "$lhs attr-dict `:` type($lhs) `->` type($result)"; + // let traits = [Pure]; +} diff --git a/include/Transforms/InsertMovPass.h b/include/Transforms/InsertMovPass.h new file mode 100644 index 00000000..06a0befa --- /dev/null +++ b/include/Transforms/InsertMovPass.h @@ -0,0 +1,13 @@ +#ifndef NEURA_TRANSFORMS_INSERTMOVPASS_H +#define NEURA_TRANSFORMS_INSERTMOVPASS_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace neura { + std::unique_ptr createInsertMovPass(); +} // namespace neura +} // namespace mlir + +#endif // NEURA_TRANSFORMS_INSERTMOVPASS_H + diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index 30d5c055..5487d736 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -1,2 +1,3 @@ add_subdirectory(NeuraDialect) add_subdirectory(Conversion/ArithToNeura) +add_subdirectory(Transforms) diff --git a/lib/Conversion/ArithToNeura/ArithToNeura.cpp b/lib/Conversion/ArithToNeura/ArithToNeura.cpp index cfb8c216..319ec5fe 100644 --- a/lib/Conversion/ArithToNeura/ArithToNeura.cpp +++ b/lib/Conversion/ArithToNeura/ArithToNeura.cpp @@ -15,12 +15,7 @@ struct ArithAddFOpLowering : public OpRewritePattern { LogicalResult matchAndRewrite(arith::AddFOp op, PatternRewriter &rewriter) const override { -llvm::errs() << "[cheng] step into matchAndRewriter()"; rewriter.replaceOpWithNewOp(op, op.getType(), op.getLhs(), op.getRhs()); - -llvm::errs() << "[cheng] Matched arith.addf: "; -// op.dump(); - return success(); } }; @@ -40,8 +35,6 @@ struct LowerArithToNeuraPass } void runOnOperation() override { - // getContext().loadDialect(); - RewritePatternSet patterns(&getContext()); llvm::errs() << "[cheng] check runOnOperation: "; getOperation().dump(); diff --git a/lib/Transforms/CMakeLists.txt b/lib/Transforms/CMakeLists.txt new file mode 100644 index 00000000..b6448528 --- /dev/null +++ b/lib/Transforms/CMakeLists.txt @@ -0,0 +1,11 @@ +add_mlir_library(NeuraTransforms + InsertMovPass.cpp + + LINK_LIBS PUBLIC + MLIRIR + MLIRFuncDialect + MLIRPass + MLIRSupport + MLIRTransformUtils + NeuraDialect +) diff --git a/lib/Transforms/InsertMovPass.cpp b/lib/Transforms/InsertMovPass.cpp new file mode 100644 index 00000000..467e3387 --- /dev/null +++ b/lib/Transforms/InsertMovPass.cpp @@ -0,0 +1,82 @@ +#include "NeuraDialect/NeuraDialect.h" +#include "NeuraDialect/NeuraOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" + +using namespace mlir; + +namespace { +struct InsertMovForNeuraOps : public RewritePattern { + InsertMovForNeuraOps(MLIRContext *context) + : RewritePattern(/*matchAnyOpTypeTag=*/MatchAnyOpTypeTag(), /*benefit=*/1, context) {} + + // using RewritePattern::RewritePattern; + + LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { + if (op->getDialect()->getNamespace() != "neura" || + isa(op)) { + return failure(); + } + + llvm::errs() << "[cheng] step into matching and rewrite"; + Location loc = op->getLoc(); + + // Wrap operands in mov + SmallVector newOperands; + for (Value operand : op->getOperands()) { + auto mov = rewriter.create(loc, operand.getType(), operand); + newOperands.push_back(mov); + } + + // Clone op with new operands + OperationState state(loc, op->getName()); + state.addOperands(newOperands); + state.addTypes(op->getResultTypes()); + state.addAttributes(op->getAttrs()); + Operation *newOp = rewriter.create(state); + + // Wrap each result in a mov + SmallVector newResults; + for (Value result : newOp->getResults()) { + auto mov = rewriter.create(loc, result.getType(), result); + newResults.push_back(mov); + } + + rewriter.replaceOp(op, newResults); + return success(); + } +}; + +struct InsertMovPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertMovPass) + + StringRef getArgument() const override { return "insert-mov"; } + StringRef getDescription() const override { + return "Insert neura.mov before and after all neura dialect operations."; + } + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } + + void runOnOperation() override { + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext()); + if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; +} // namespace + +namespace mlir { +namespace neura { + +std::unique_ptr createInsertMovPass() { + return std::make_unique(); +} + +} // namespace neura +} // namespace mlir diff --git a/test/neura/add.mlir b/test/neura/add.mlir new file mode 100644 index 00000000..a4564240 --- /dev/null +++ b/test/neura/add.mlir @@ -0,0 +1,9 @@ +// RN: mlir-neura-opt --lower-arith-to-neura --insert-mov %s | FileCheck %s +// RUN: mlir-neura-opt --insert-mov %s | FileCheck %s + +func.func @test(%a: f32) -> f32 { + %b = arith.constant 2.0 : f32 + %res = arith.addf %a, %b : f32 + // CHECK: neura.add + return %res : f32 +} diff --git a/tools/mlir-neura-opt/CMakeLists.txt b/tools/mlir-neura-opt/CMakeLists.txt index 8f617649..08fa6db8 100644 --- a/tools/mlir-neura-opt/CMakeLists.txt +++ b/tools/mlir-neura-opt/CMakeLists.txt @@ -15,6 +15,7 @@ target_link_libraries(mlir-neura-opt PRIVATE MLIRFuncDialect # Builtin dialect required by custom dialect MLIRArithDialect NeuraArithToNeura + NeuraTransforms ) # Includes directories. diff --git a/tools/mlir-neura-opt/mlir-neura-opt.cpp b/tools/mlir-neura-opt/mlir-neura-opt.cpp index 85df2f1e..b10cba16 100644 --- a/tools/mlir-neura-opt/mlir-neura-opt.cpp +++ b/tools/mlir-neura-opt/mlir-neura-opt.cpp @@ -7,6 +7,7 @@ #include "mlir/Tools/mlir-opt/MlirOptMain.h" #include "Conversion/ArithToNeura/ArithToNeura.h" #include "NeuraDialect/NeuraDialect.h" +#include "Transforms/InsertMovPass.h" int main(int argc, char **argv) { // Registers MLIR dialects. @@ -18,6 +19,9 @@ int main(int argc, char **argv) { mlir::registerPass([]() -> std::unique_ptr { return mlir::neura::createLowerArithToNeuraPass(); }); + mlir::registerPass([]() -> std::unique_ptr { + return mlir::neura::createInsertMovPass(); + }); // Runs the MLIR optimizer. return mlir::asMainReturnCode( From 35992b2f3c8c1f017e7f09a96a802796cb22e6db Mon Sep 17 00:00:00 2001 From: tancheng Date: Wed, 14 May 2025 05:42:49 +0000 Subject: [PATCH 2/3] [update] Avoid adding mov on op's result & fix hanging --- lib/Transforms/InsertMovPass.cpp | 32 ++++++++++++++++++-------------- test/neura/add.mlir | 5 +++-- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/lib/Transforms/InsertMovPass.cpp b/lib/Transforms/InsertMovPass.cpp index 467e3387..0ff05656 100644 --- a/lib/Transforms/InsertMovPass.cpp +++ b/lib/Transforms/InsertMovPass.cpp @@ -12,39 +12,43 @@ struct InsertMovForNeuraOps : public RewritePattern { InsertMovForNeuraOps(MLIRContext *context) : RewritePattern(/*matchAnyOpTypeTag=*/MatchAnyOpTypeTag(), /*benefit=*/1, context) {} - // using RewritePattern::RewritePattern; - LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { if (op->getDialect()->getNamespace() != "neura" || isa(op)) { return failure(); } - llvm::errs() << "[cheng] step into matching and rewrite"; + // Skips ops that already being inserted mov on the operands. + bool allInputsAreMov = llvm::all_of(op->getOperands(), [](Value v) { + return isa_and_nonnull(v.getDefiningOp()); + }); + if (allInputsAreMov) { + return failure(); + } + + // Makes sure none of the operand has being processed. + bool hasAnyMovInput = llvm::any_of(op->getOperands(), [](Value v) { + return isa_and_nonnull(v.getDefiningOp()); + }); + assert(!hasAnyMovInput && "Unexpected: operand already wrapped in neura.mov"); + Location loc = op->getLoc(); - // Wrap operands in mov + // Wraps operands in mov. SmallVector newOperands; for (Value operand : op->getOperands()) { auto mov = rewriter.create(loc, operand.getType(), operand); newOperands.push_back(mov); } - // Clone op with new operands + // Clones op with new operands. OperationState state(loc, op->getName()); state.addOperands(newOperands); state.addTypes(op->getResultTypes()); state.addAttributes(op->getAttrs()); - Operation *newOp = rewriter.create(state); - - // Wrap each result in a mov - SmallVector newResults; - for (Value result : newOp->getResults()) { - auto mov = rewriter.create(loc, result.getType(), result); - newResults.push_back(mov); - } - rewriter.replaceOp(op, newResults); + Operation *newOp = rewriter.create(state); + rewriter.replaceOp(op, newOp->getResults()); return success(); } }; diff --git a/test/neura/add.mlir b/test/neura/add.mlir index a4564240..df477461 100644 --- a/test/neura/add.mlir +++ b/test/neura/add.mlir @@ -1,9 +1,10 @@ -// RN: mlir-neura-opt --lower-arith-to-neura --insert-mov %s | FileCheck %s -// RUN: mlir-neura-opt --insert-mov %s | FileCheck %s +// RUN: mlir-neura-opt --lower-arith-to-neura --insert-mov %s | FileCheck %s func.func @test(%a: f32) -> f32 { %b = arith.constant 2.0 : f32 %res = arith.addf %a, %b : f32 + // CHECK: neura.mov %arg0 : f32 -> f32 + // CHECK: neura.mov %cst : f32 -> f32 // CHECK: neura.add return %res : f32 } From 2d4fc7d50b967c33f3ac98f4b128df383556b55f Mon Sep 17 00:00:00 2001 From: tancheng Date: Wed, 14 May 2025 17:45:20 +0000 Subject: [PATCH 3/3] [feature] Lower LLVM to Neura --- include/Conversion/LlvmToNeura/LlvmToNeura.h | 12 +++++ lib/CMakeLists.txt | 2 +- ...{ArithToNeura.cpp => ArithToNeuraPass.cpp} | 6 +-- lib/Conversion/ArithToNeura/CMakeLists.txt | 2 +- lib/Conversion/CMakeLists.txt | 2 + lib/Conversion/LlvmToNeura/CMakeLists.txt | 28 ++++++++++ .../LlvmToNeura/LlvmToNeuraPass.cpp | 51 +++++++++++++++++++ test/neura/{add.mlir => arith_add.mlir} | 0 test/neura/llvm_add.mlir | 10 ++++ tools/mlir-neura-opt/CMakeLists.txt | 6 ++- tools/mlir-neura-opt/mlir-neura-opt.cpp | 6 +++ 11 files changed, 116 insertions(+), 9 deletions(-) create mode 100644 include/Conversion/LlvmToNeura/LlvmToNeura.h rename lib/Conversion/ArithToNeura/{ArithToNeura.cpp => ArithToNeuraPass.cpp} (86%) create mode 100644 lib/Conversion/CMakeLists.txt create mode 100644 lib/Conversion/LlvmToNeura/CMakeLists.txt create mode 100644 lib/Conversion/LlvmToNeura/LlvmToNeuraPass.cpp rename test/neura/{add.mlir => arith_add.mlir} (100%) create mode 100644 test/neura/llvm_add.mlir diff --git a/include/Conversion/LlvmToNeura/LlvmToNeura.h b/include/Conversion/LlvmToNeura/LlvmToNeura.h new file mode 100644 index 00000000..cbe51792 --- /dev/null +++ b/include/Conversion/LlvmToNeura/LlvmToNeura.h @@ -0,0 +1,12 @@ +#ifndef NEURA_CONVERSION_LLVMTONEURA_LLVMTONEURAPASS_H +#define NEURA_CONVERSION_LLVMTONEURA_LLVMTONEURAPASS_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace neura { + std::unique_ptr createLowerLlvmToNeuraPass(); +} // namespace neura +} // namespace mlir + +#endif // NEURA_CONVERSION_LLVMTONEURA_LLVMTONEURAPASS_H diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index 5487d736..aec92864 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -1,3 +1,3 @@ add_subdirectory(NeuraDialect) -add_subdirectory(Conversion/ArithToNeura) +add_subdirectory(Conversion) add_subdirectory(Transforms) diff --git a/lib/Conversion/ArithToNeura/ArithToNeura.cpp b/lib/Conversion/ArithToNeura/ArithToNeuraPass.cpp similarity index 86% rename from lib/Conversion/ArithToNeura/ArithToNeura.cpp rename to lib/Conversion/ArithToNeura/ArithToNeuraPass.cpp index 319ec5fe..fed2c6ce 100644 --- a/lib/Conversion/ArithToNeura/ArithToNeura.cpp +++ b/lib/Conversion/ArithToNeura/ArithToNeuraPass.cpp @@ -27,7 +27,7 @@ struct LowerArithToNeuraPass StringRef getArgument() const override { return "lower-arith-to-neura"; } StringRef getDescription() const override { - return "Lower arithmetic operations to Neura dialect operations"; + return "Lower arith dialect operations to Neura dialect operations"; } void getDependentDialects(DialectRegistry ®istry) const override { @@ -36,11 +36,7 @@ struct LowerArithToNeuraPass void runOnOperation() override { RewritePatternSet patterns(&getContext()); - llvm::errs() << "[cheng] check runOnOperation: "; getOperation().dump(); - getOperation().walk([](Operation *op) { - llvm::errs() << "[cheng] Saw op: " << op->getName() << "\n"; - }); patterns.add(&getContext()); if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) { signalPassFailure(); diff --git a/lib/Conversion/ArithToNeura/CMakeLists.txt b/lib/Conversion/ArithToNeura/CMakeLists.txt index a5eee595..dbc4c462 100644 --- a/lib/Conversion/ArithToNeura/CMakeLists.txt +++ b/lib/Conversion/ArithToNeura/CMakeLists.txt @@ -1,5 +1,5 @@ add_mlir_library(NeuraArithToNeura - ArithToNeura.cpp + ArithToNeuraPass.cpp DEPENDS NeuraOpsIncGen diff --git a/lib/Conversion/CMakeLists.txt b/lib/Conversion/CMakeLists.txt new file mode 100644 index 00000000..b917c4f3 --- /dev/null +++ b/lib/Conversion/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(ArithToNeura) +add_subdirectory(LlvmToNeura) diff --git a/lib/Conversion/LlvmToNeura/CMakeLists.txt b/lib/Conversion/LlvmToNeura/CMakeLists.txt new file mode 100644 index 00000000..3a3c7c92 --- /dev/null +++ b/lib/Conversion/LlvmToNeura/CMakeLists.txt @@ -0,0 +1,28 @@ +add_mlir_library(NeuraLlvmToNeura + LlvmToNeuraPass.cpp + + DEPENDS + NeuraOpsIncGen + NeuraDialectIncGen + NeuraDialect + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRFuncDialect + MLIRLLVMDialect + MLIRIR + MLIRPass + MLIRTransforms +) + +target_include_directories(NeuraLlvmToNeura PUBLIC + ${CMAKE_BINARY_DIR}/lib/NeuraDialect + ${MLIR_INCLUDE_DIRS} + ${LLVM_INCLUDE_DIRS} + ${CMAKE_SOURCE_DIR}/include + ${CMAKE_BINARY_DIR}/include +) + +target_compile_definitions(NeuraLlvmToNeura + PRIVATE ${LLVM_DEFINITIONS} +) diff --git a/lib/Conversion/LlvmToNeura/LlvmToNeuraPass.cpp b/lib/Conversion/LlvmToNeura/LlvmToNeuraPass.cpp new file mode 100644 index 00000000..adff4bda --- /dev/null +++ b/lib/Conversion/LlvmToNeura/LlvmToNeuraPass.cpp @@ -0,0 +1,51 @@ +#include "Conversion/LlvmToNeura/LlvmToNeura.h" +#include "NeuraDialect/NeuraDialect.h" +#include "NeuraDialect/NeuraOps.h" +#include "mlir/Dialect/LLVMIR/LLVMAttrs.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMTypes.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +using namespace mlir; + +namespace { +struct LlvmAddFOpLowering : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(mlir::LLVM::FAddOp op, + PatternRewriter &rewriter) const override { + rewriter.replaceOpWithNewOp(op, op.getType(), op.getLhs(), op.getRhs()); + return success(); + } +}; + +struct LowerLlvmToNeuraPass + : public PassWrapper> { + + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerLlvmToNeuraPass) + + StringRef getArgument() const override { return "lower-llvm-to-neura"; } + StringRef getDescription() const override { + return "Lower LLVM operations to Neura dialect operations"; + } + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } + + void runOnOperation() override { + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext()); + if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) { + signalPassFailure(); + } + } +}; +} // namespace + +std::unique_ptr mlir::neura::createLowerLlvmToNeuraPass() { + return std::make_unique(); +} diff --git a/test/neura/add.mlir b/test/neura/arith_add.mlir similarity index 100% rename from test/neura/add.mlir rename to test/neura/arith_add.mlir diff --git a/test/neura/llvm_add.mlir b/test/neura/llvm_add.mlir new file mode 100644 index 00000000..56b111cd --- /dev/null +++ b/test/neura/llvm_add.mlir @@ -0,0 +1,10 @@ +// RUN: mlir-neura-opt --lower-llvm-to-neura --insert-mov %s | FileCheck %s + +func.func @test(%a: f32) -> f32 { + %b = llvm.mlir.constant(2.0 : f32) : f32 + %res = llvm.fadd %a, %b : f32 + // CHECK: [[LHS:%.*]] = neura.mov %{{.*}} : f32 -> f32 + // CHECK: [[RHS:%.*]] = neura.mov %{{.*}} : f32 -> f32 + // CHECK: [[RES:%.*]] = neura.add [[LHS]], [[RHS]] : f32 + return %res : f32 +} diff --git a/tools/mlir-neura-opt/CMakeLists.txt b/tools/mlir-neura-opt/CMakeLists.txt index 08fa6db8..5c774f2f 100644 --- a/tools/mlir-neura-opt/CMakeLists.txt +++ b/tools/mlir-neura-opt/CMakeLists.txt @@ -6,15 +6,17 @@ add_executable(mlir-neura-opt # Links MLIR libraries. target_link_libraries(mlir-neura-opt PRIVATE - MLIROptLib # MLIR optimizer library + MLIRDialect # MLIR Dialect MLIRIR # MLIR Core IR + MLIRLLVMDialect + MLIROptLib # MLIR optimizer library MLIRSupport # MLIR Support utilities MLIRTransforms # MLIR transformation passes - MLIRDialect # MLIR Dialect NeuraDialect # Custom dialect MLIRFuncDialect # Builtin dialect required by custom dialect MLIRArithDialect NeuraArithToNeura + NeuraLlvmToNeura NeuraTransforms ) diff --git a/tools/mlir-neura-opt/mlir-neura-opt.cpp b/tools/mlir-neura-opt/mlir-neura-opt.cpp index b10cba16..beaec22d 100644 --- a/tools/mlir-neura-opt/mlir-neura-opt.cpp +++ b/tools/mlir-neura-opt/mlir-neura-opt.cpp @@ -1,11 +1,13 @@ // tools/mlir-neura-opt/mlir-neura-opt.cpp +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/InitAllDialects.h" #include "mlir/InitAllPasses.h" #include "mlir/Support/FileUtilities.h" #include "mlir/Support/LogicalResult.h" #include "mlir/Tools/mlir-opt/MlirOptMain.h" #include "Conversion/ArithToNeura/ArithToNeura.h" +#include "Conversion/LlvmToNeura/LlvmToNeura.h" #include "NeuraDialect/NeuraDialect.h" #include "Transforms/InsertMovPass.h" @@ -15,10 +17,14 @@ int main(int argc, char **argv) { registry.insert(); registry.insert(); registry.insert(); + registry.insert(); mlir::registerPass([]() -> std::unique_ptr { return mlir::neura::createLowerArithToNeuraPass(); }); + mlir::registerPass([]() -> std::unique_ptr { + return mlir::neura::createLowerLlvmToNeuraPass(); + }); mlir::registerPass([]() -> std::unique_ptr { return mlir::neura::createInsertMovPass(); });