Pattern fusion template

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<Pass> createLowerLlvmToNeuraPass();
+} // namespace neura
+} // namespace mlir
+
+#endif // NEURA_CONVERSION_LLVMTONEURA_LLVMTONEURAPASS_H

include/NeuraDialect/NeuraOps.td

@@ -1,13 +1,42 @@
 // NeuraOps.td - Custom operation definitions.
-include "mlir/IR/OpBase.td"
+
 include "NeuraDialect/NeuraDialect.td"
 
 // Defines an addition operation.
 def Neura_AddOp : Op<NeuraDialect, "add"> {
-  let summary = "Addition operation";
+  let summary = "Integer addition operation";
   let opName = "add";
-  let arguments = (ins F32:$lhs, F32:$rhs);
-  let results = (outs F32:$result);
-  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
+  let arguments = (ins AnyInteger:$lhs, AnyInteger:$rhs);
+  let results = (outs AnyInteger:$result);
+  // let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
+  let traits = [SameOperandsAndResultElementType];
+}
+
+// Defines an addition operation.
+def Neura_FAddOp : Op<NeuraDialect, "fadd"> {
+  let summary = "Floating addition operation";
+  let opName = "fadd";
+  let arguments = (ins AnyFloat:$lhs, AnyFloat:$rhs);
+  let results = (outs AnyFloat:$result);
+  // let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
+  let traits = [SameOperandsAndResultElementType];
+}
+
+def Neura_FAddFAddOp : Op<NeuraDialect, "fadd_fadd"> {
+  let summary = "Fused fadd(fadd(a, b), c)";
+  let arguments = (ins AnyFloat:$a, AnyFloat:$b, AnyFloat:$c);
+  let results = (outs AnyFloat:$result);
+  // let assemblyFormat = "$a `,` $b `,` $c attr-dict `:` type($result)";
+  let traits = [SameOperandsAndResultElementType];
+}
+
+
+// Defines a move operation for data communication.
+def Neura_MovOp : Op<NeuraDialect, "mov"> {
+  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];
 }

include/Transforms/FusePatternsPass.h

@@ -0,0 +1,11 @@
+#ifndef NEURA_TRANSFORMS_FUSEPATTERNSPASS_H
+#define NEURA_TRANSFORMS_FUSEPATTERNSPASS_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir::neura {
+std::unique_ptr<mlir::Pass> createFusePatternsPass();
+}
+
+#endif // NEURA_TRANSFORMS_FUSEPATTERNSPASS_H
+

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<mlir::Pass> createInsertMovPass();
+} // namespace neura
+} // namespace mlir
+
+#endif // NEURA_TRANSFORMS_INSERTMOVPASS_H
+

lib/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(NeuraDialect)
-add_subdirectory(Conversion/ArithToNeura)
+add_subdirectory(Conversion)
+add_subdirectory(Transforms)

lib/Conversion/ArithToNeura/ArithToNeura.cpp → lib/Conversion/ArithToNeura/ArithToNeuraPass.cpp

@@ -7,23 +7,20 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
-using namespace mlir;
+namespace mlir {
+namespace neura {
+// Uses arith2neura instead of llvm to avoid conflicts.
+namespace arith2neura {
 
-namespace {
-struct ArithAddFOpLowering : public OpRewritePattern<arith::AddFOp> {
-  using OpRewritePattern::OpRewritePattern;
+#include "ArithToNeuraPatterns.inc"
 
-  LogicalResult matchAndRewrite(arith::AddFOp op,
-                                PatternRewriter &rewriter) const override {
-llvm::errs() << "[cheng] step into matchAndRewriter()";
-    rewriter.replaceOpWithNewOp<neura::AddOp>(op, op.getType(), op.getLhs(), op.getRhs());
+} // namespace arith2neura
+} // namespace neura
+} // namespace mlir
 
-llvm::errs() << "[cheng] Matched arith.addf: ";
-// op.dump();
+using namespace mlir;
 
-    return success();
-  }
-};
+namespace {
 
 struct LowerArithToNeuraPass
     : public PassWrapper<LowerArithToNeuraPass, OperationPass<func::FuncOp>> {
@@ -32,23 +29,16 @@ 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 &registry) const override {
     registry.insert<mlir::neura::NeuraDialect>();
   }
 
   void runOnOperation() override {
-    // getContext().loadDialect<mlir::neura::NeuraDialect>();
-
     RewritePatternSet patterns(&getContext());
-    llvm::errs() << "[cheng] check runOnOperation: ";
-    getOperation().dump();
-    getOperation().walk([](Operation *op) {
-      llvm::errs() << "[cheng] Saw op: " << op->getName() << "\n";
-    });
-    patterns.add<ArithAddFOpLowering>(&getContext());
+    mlir::neura::arith2neura::populateWithGenerated(patterns);
     if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
       signalPassFailure();
     }

lib/Conversion/ArithToNeura/ArithToNeuraPatterns.td

@@ -0,0 +1,10 @@
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/Arith/IR/ArithOps.td"
+include "NeuraDialect/NeuraOps.td"
+
+def : Pat<
+  (Arith_AddFOp $lhs, $rhs, $_fastmath),
+  (Neura_FAddOp $lhs, $rhs)
+>;
+

lib/Conversion/ArithToNeura/CMakeLists.txt

@@ -1,7 +1,21 @@
+set(LLVM_TARGET_DEFINITIONS
+  ${CMAKE_CURRENT_SOURCE_DIR}/ArithToNeuraPatterns.td
+)
+
+mlir_tablegen(ArithToNeuraPatterns.inc
+  -gen-rewriters
+  -I ${MLIR_SOURCE_DIR}/include
+  -I ${MLIR_BINARY_DIR}/include
+  -I ${CMAKE_SOURCE_DIR}/include
+  -I ${CMAKE_CURRENT_SOURCE_DIR}
+)
+add_public_tablegen_target(ArithToNeuraPatternGen)
+
 add_mlir_library(NeuraArithToNeura
-  ArithToNeura.cpp
+  ArithToNeuraPass.cpp
 
   DEPENDS
+  ArithToNeuraPatternGen
   NeuraOpsIncGen
   NeuraDialectIncGen
   NeuraDialect
@@ -16,6 +30,7 @@ add_mlir_library(NeuraArithToNeura
 
 target_include_directories(NeuraArithToNeura PUBLIC
     ${CMAKE_BINARY_DIR}/lib/NeuraDialect
+    ${CMAKE_BINARY_DIR}/lib/Conversion/ArithToNeura
     ${MLIR_INCLUDE_DIRS}
     ${LLVM_INCLUDE_DIRS}
     ${CMAKE_SOURCE_DIR}/include
@@ -25,3 +40,4 @@ target_include_directories(NeuraArithToNeura PUBLIC
 target_compile_definitions(NeuraArithToNeura
   PRIVATE ${LLVM_DEFINITIONS}
 )
+

lib/Conversion/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_subdirectory(ArithToNeura)
+add_subdirectory(LlvmToNeura)

lib/Conversion/LlvmToNeura/CMakeLists.txt

@@ -0,0 +1,43 @@
+set(LLVM_TARGET_DEFINITIONS
+  ${CMAKE_CURRENT_SOURCE_DIR}/LlvmToNeuraPatterns.td
+)
+
+mlir_tablegen(LlvmToNeuraPatterns.inc
+  -gen-rewriters
+  -I ${MLIR_SOURCE_DIR}/include
+  -I ${MLIR_BINARY_DIR}/include
+  -I ${CMAKE_SOURCE_DIR}/include
+  -I ${CMAKE_CURRENT_SOURCE_DIR}
+)
+add_public_tablegen_target(LlvmToNeuraPatternGen)
+
+add_mlir_library(NeuraLlvmToNeura
+  LlvmToNeuraPass.cpp
+
+  DEPENDS
+  NeuraOpsIncGen
+  NeuraDialectIncGen
+  NeuraDialect
+  LlvmToNeuraPatternGen
+
+  LINK_LIBS PUBLIC
+  MLIRArithDialect
+  MLIRFuncDialect
+  MLIRLLVMDialect
+  MLIRIR
+  MLIRPass
+  MLIRTransforms
+)
+
+target_include_directories(NeuraLlvmToNeura PUBLIC
+    ${CMAKE_BINARY_DIR}/lib/NeuraDialect
+    ${CMAKE_BINARY_DIR}/lib/Conversion/LlvmToNeura
+    ${MLIR_INCLUDE_DIRS}
+    ${LLVM_INCLUDE_DIRS}
+    ${CMAKE_SOURCE_DIR}/include
+    ${CMAKE_BINARY_DIR}/include
+)
+
+target_compile_definitions(NeuraLlvmToNeura
+  PRIVATE ${LLVM_DEFINITIONS}
+)

lib/Conversion/LlvmToNeura/LlvmToNeuraPass.cpp

@@ -0,0 +1,53 @@
+#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"
+
+namespace mlir {
+namespace neura {
+// Uses llvm2neura instead of llvm to avoid conflicts.
+namespace llvm2neura {
+
+#include "LlvmToNeuraPatterns.inc"
+
+} // namespace llvm2neura
+} // namespace neura
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+struct LowerLlvmToNeuraPass
+    : public PassWrapper<LowerLlvmToNeuraPass, OperationPass<func::FuncOp>> {
+
+  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 &registry) const override {
+    registry.insert<mlir::neura::NeuraDialect>();
+  }
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    mlir::neura::llvm2neura::populateWithGenerated(patterns);
+    if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::neura::createLowerLlvmToNeuraPass() {
+  return std::make_unique<LowerLlvmToNeuraPass>();
+}

lib/Conversion/LlvmToNeura/LlvmToNeuraPatterns.td

@@ -0,0 +1,9 @@
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/LLVMIR/LLVMOps.td"
+include "NeuraDialect/NeuraOps.td"
+
+def : Pat<
+  (LLVM_FAddOp $lhs, $rhs, $_fastmath),
+  (Neura_FAddOp $lhs, $rhs)
+>;

lib/NeuraDialect/CMakeLists.txt

@@ -5,12 +5,22 @@ add_definitions(${MLIR_DEFINITIONS})
 
 set(LLVM_TARGET_DEFINITIONS ${CMAKE_CURRENT_SOURCE_DIR}/../../include/NeuraDialect/NeuraOps.td)
 
-mlir_tablegen(NeuraOps.h.inc -gen-op-decls -I${CMAKE_SOURCE_DIR}/include)
-mlir_tablegen(NeuraOps.cpp.inc -gen-op-defs -I${CMAKE_SOURCE_DIR}/include)
+mlir_tablegen(NeuraOps.h.inc -gen-op-decls
+  -I${MLIR_SOURCE_DIR}/include
+  -I${CMAKE_SOURCE_DIR}/include)
+mlir_tablegen(NeuraOps.cpp.inc -gen-op-defs
+  -I${MLIR_SOURCE_DIR}/include
+  -I${CMAKE_SOURCE_DIR}/include)
 add_public_tablegen_target(NeuraOpsIncGen)
 
-mlir_tablegen(NeuraDialect.h.inc -gen-dialect-decls -dialect=neura -I${CMAKE_SOURCE_DIR}/include)
-mlir_tablegen(NeuraDialect.cpp.inc -gen-dialect-defs -dialect=neura -I${CMAKE_SOURCE_DIR}/include)
+mlir_tablegen(NeuraDialect.h.inc -gen-dialect-decls
+  -dialect=neura
+  -I${MLIR_SOURCE_DIR}/include
+  -I${CMAKE_SOURCE_DIR}/include)
+mlir_tablegen(NeuraDialect.cpp.inc -gen-dialect-defs
+  -dialect=neura
+  -I${MLIR_SOURCE_DIR}/include
+  -I${CMAKE_SOURCE_DIR}/include)
 add_public_tablegen_target(NeuraDialectIncGen)
 
 add_public_tablegen_target(MLIRNeuraIncGen)

lib/Transforms/CMakeLists.txt

@@ -0,0 +1,23 @@
+# set(LLVM_TARGET_DEFINITIONS ${CMAKE_CURRENT_SOURCE_DIR}/FusePatterns.td)
+#
+# mlir_tablegen(FusePatterns.inc -gen-rewriters
+#   -I ${MLIR_SOURCE_DIR}/include
+#   -I ${MLIR_BINARY_DIR}/include
+#   -I ${CMAKE_SOURCE_DIR}/include
+#   -I ${CMAKE_CURRENT_SOURCE_DIR}
+# )
+#
+# add_public_tablegen_target(NeuraFusePatternsGen)
+
+add_mlir_library(NeuraTransforms
+  InsertMovPass.cpp
+  FusePatternsPass.cpp
+
+  LINK_LIBS PUBLIC
+    MLIRIR
+    MLIRFuncDialect
+    MLIRPass
+    MLIRSupport
+    MLIRTransformUtils
+    NeuraDialect
+)