Tag kernel with accelerator attr for optional lowering

.github/workflows/main.yml

@@ -64,8 +64,6 @@ jobs:
     - name: setup dataflow tool-chain
       working-directory: ${{github.workspace}}
       run: |
-        git clone https://github.com/coredac/dataflow.git
-        cd dataflow
         mkdir build && cd build
         cmake -G Ninja .. \
           -DLLVM_DIR=${{github.workspace}}/llvm-project/build/lib/cmake/llvm \
@@ -102,6 +100,6 @@ jobs:
     - name: run test
       working-directory: ${{github.workspace}}
       run: |
-        cd ${{github.workspace}}/dataflow/test
+        cd ${{github.workspace}}/test
         ${{github.workspace}}/llvm-project/build/bin/llvm-lit * -v
 

include/Common/AcceleratorAttrs.h

@@ -0,0 +1,20 @@
+#ifndef COMMON_ACCELERATOR_ATTRS_H
+#define COMMON_ACCELERATOR_ATTRS_H
+
+#include "llvm/ADT/StringRef.h"
+
+namespace mlir {
+namespace accel {
+
+// Common attribute key.
+constexpr llvm::StringRef kAcceleratorAttr = "accelerator";
+
+// Common accelerator targets.
+constexpr llvm::StringRef kNeuraTarget = "neura";
+constexpr llvm::StringRef kGpuTarget   = "gpu";
+constexpr llvm::StringRef kTpuTarget   = "tpu";
+
+} // namespace accel
+} // namespace mlir
+
+#endif // COMMON_ACCELERATOR_ATTRS_H

include/NeuraDialect/NeuraOps.td

@@ -76,3 +76,26 @@ def Neura_MovOp : Op<NeuraDialect, "mov"> {
   let assemblyFormat = "$lhs attr-dict `:` type($lhs) `->` type($result)";
   // let traits = [Pure];
 }
+
+def Neura_ICmpOp : Op<NeuraDialect, "icmp"> {
+  let summary = "Integer compare operation";
+  let opName = "icmp";
+  let arguments = (ins AnyInteger:$lhs, AnyInteger:$rhs,
+                   StrAttr:$predicate);
+  let results = (outs I1:$result);
+  // let assemblyFormat = "$lhs `,` $rhs `,` $predicate attr-dict `:` type($result)";
+  // let traits = [SameOperandsAndResultElementType];
+}
+
+def Neura_LoadOp : Op<NeuraDialect, "load"> {
+  let arguments = (ins AnyType:$addr);
+  let results = (outs AnyType:$value);
+  // let assemblyFormat = "$addr attr-dict `:` type($value)";
+}
+
+def Neura_StoreOp : Op<NeuraDialect, "store"> {
+  let arguments = (ins AnyType:$value, AnyType:$addr);
+  let results = (outs);
+  // let assemblyFormat = "$value `,` $addr attr-dict";
+}
+

include/Transforms/AssignAcceleratorPass.h

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

lib/Conversion/LlvmToNeura/LlvmToNeuraPass.cpp

@@ -1,4 +1,5 @@
 #include "Conversion/LlvmToNeura/LlvmToNeura.h"
+#include "Common/AcceleratorAttrs.h"
 #include "NeuraDialect/NeuraDialect.h"
 #include "NeuraDialect/NeuraOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
@@ -75,6 +76,22 @@ struct LlvmVFMulToNeuraVFMul: public OpRewritePattern<mlir::LLVM::FMulOp> {
   }
 };
 
+struct LlvmICmpToNeuraICmp : public OpRewritePattern<LLVM::ICmpOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(LLVM::ICmpOp op,
+                                PatternRewriter &rewriter) const override {
+    auto pred = op.getPredicate();
+    auto lhs = op.getLhs();
+    auto rhs = op.getRhs();
+    auto resultType = op.getType();
+
+    rewriter.replaceOpWithNewOp<neura::ICmpOp>(
+        op, resultType, lhs, rhs, rewriter.getStringAttr(LLVM::stringifyICmpPredicate(pred)));
+    return success();
+  }
+};
+
 struct LowerLlvmToNeuraPass
     : public PassWrapper<LowerLlvmToNeuraPass, OperationPass<ModuleOp>> {
 
@@ -96,17 +113,21 @@ struct LowerLlvmToNeuraPass
     patterns.add<LlvmAddToNeuraAdd>(&getContext());
     patterns.add<LlvmFMulToNeuraFMul>(&getContext());
     patterns.add<LlvmVFMulToNeuraVFMul>(&getContext());
+    patterns.add<LlvmICmpToNeuraICmp>(&getContext());
     FrozenRewritePatternSet frozen(std::move(patterns));
 
     ModuleOp module_op = getOperation();
 
     // Applies to every region inside the module (regardless of func type,
     // e.g., mlir func or llvm func).
-    module_op.walk([&](Operation *op) {
-      if (!op->getRegions().empty()) {
-        for (Region &region : op->getRegions()) {
-          if (failed(applyPatternsAndFoldGreedily(region, frozen))) {
-            signalPassFailure();
+    module_op.walk([&](FunctionOpInterface func) {
+      if (func->hasAttr(mlir::accel::kAcceleratorAttr)) {
+        auto target = func->getAttrOfType<StringAttr>(mlir::accel::kAcceleratorAttr);
+        if (target && target.getValue() == mlir::accel::kNeuraTarget) {
+          for (Region &region : func->getRegions()) {
+            if (failed(applyPatternsAndFoldGreedily(region, frozen))) {
+              signalPassFailure();
+            }
           }
         }
       }

lib/Transforms/AssignAcceleratorPass.cpp

@@ -0,0 +1,42 @@
+#include "Common/AcceleratorAttrs.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct AssignAcceleratorPass : public PassWrapper<AssignAcceleratorPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AssignAcceleratorPass)
+
+  StringRef getArgument() const override { return "assign-accelerator"; }
+  StringRef getDescription() const override { return "Tags non-main functions as neura.kernel."; }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    Builder builder(&getContext());
+
+    module.walk([&](Operation *op) {
+      if (auto func = dyn_cast<FunctionOpInterface>(op)) {
+        if (func.getName() != "main" &&
+            !func.isExternal() &&
+            !func->hasAttr(mlir::accel::kAcceleratorAttr)) {
+          func->setAttr(mlir::accel::kAcceleratorAttr, builder.getStringAttr(mlir::accel::kNeuraTarget));
+        }
+      }
+    });
+  }
+};
+} // namespace
+
+/// Register the pass
+namespace mlir {
+namespace neura {
+std::unique_ptr<Pass> createAssignAcceleratorPass() {
+  return std::make_unique<AssignAcceleratorPass>();
+}
+} // namespace neura
+} // namespace mlir
+

lib/Transforms/CMakeLists.txt

@@ -1,15 +1,5 @@
-# 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
+  AssignAcceleratorPass.cpp
   InsertMovPass.cpp
   FusePatternsPass.cpp
 

test/neura/for_loop/test.mlir

@@ -4,15 +4,14 @@
 
 // Lowers to neura.
 // RUN: mlir-neura-opt \
+// RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --fuse-patterns \
 // RUN:   --insert-mov \
 // RUN:   %t-kernel.mlir | FileCheck %s
 
 // Verifies the neura ops are generated. And fusion happens.
-// CHECK:      "neura.vfmul"
-// CHECK:      "neura.add"
-// CHECK:      "neura.fmul_fadd"
-// CHECK:      [[LHS:%.*]] = neura.mov %{{.*}}
-// CHECK-NEXT: [[RHS:%.*]] = neura.mov %{{.*}}
-// CHECK-NEXT: [[RES:%.*]] = "neura.add"([[LHS]], [[RHS]])
+// CHECK: accelerator = "neura"
+// CHECK: neura.fmul_fadd
+// CHECK: neura.add
+// CHECK: neura.icmp

test/neura/llvm_add.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-neura-opt --lower-llvm-to-neura --insert-mov %s | FileCheck %s
+// RUN: mlir-neura-opt --assign-accelerator --lower-llvm-to-neura --insert-mov %s | FileCheck %s
 
 func.func @test(%a: f32) -> f32 {
   %b = llvm.mlir.constant(2.0 : f32) : f32

tools/mlir-neura-opt/mlir-neura-opt.cpp

@@ -10,6 +10,7 @@
 #include "Conversion/ArithToNeura/ArithToNeura.h"
 #include "Conversion/LlvmToNeura/LlvmToNeura.h"
 #include "NeuraDialect/NeuraDialect.h"
+#include "Transforms/AssignAcceleratorPass.h"
 #include "Transforms/InsertMovPass.h"
 #include "Transforms/FusePatternsPass.h"
 
@@ -25,6 +26,9 @@ int main(int argc, char **argv) {
   mlir::registerPass([]() -> std::unique_ptr<mlir::Pass> {
     return mlir::neura::createLowerArithToNeuraPass();
   });
+  mlir::registerPass([]() -> std::unique_ptr<mlir::Pass> {
+    return mlir::neura::createAssignAcceleratorPass();
+  });
   mlir::registerPass([]() -> std::unique_ptr<mlir::Pass> {
     return mlir::neura::createLowerLlvmToNeuraPass();
   });