diff --git a/lib/NeuraDialect/Transforms/GenerateCodePass.cpp b/lib/NeuraDialect/Transforms/GenerateCodePass.cpp index 6c223f83..f672db95 100644 --- a/lib/NeuraDialect/Transforms/GenerateCodePass.cpp +++ b/lib/NeuraDialect/Transforms/GenerateCodePass.cpp @@ -34,8 +34,9 @@ struct GenerateCodePass for (auto func : module.getOps()) { auto accel_attr = func->getAttrOfType("accelerator"); - if (!accel_attr || accel_attr.getValue() != "neura") + if (!accel_attr || accel_attr.getValue() != "neura") { continue; + } llvm::json::Object func_obj; func_obj["name"] = func.getName().str(); diff --git a/lib/NeuraDialect/Transforms/LeveragePredicatedValuePass.cpp b/lib/NeuraDialect/Transforms/LeveragePredicatedValuePass.cpp index 48038f77..38db7e3c 100644 --- a/lib/NeuraDialect/Transforms/LeveragePredicatedValuePass.cpp +++ b/lib/NeuraDialect/Transforms/LeveragePredicatedValuePass.cpp @@ -32,7 +32,11 @@ struct LeveragePredicatedValuePass ModuleOp module = getOperation(); // Processes each function. - module.walk([&](func::FuncOp func) { + module.walk([&](FunctionOpInterface func) { + auto accel_attr = func->getAttrOfType("accelerator"); + if (!accel_attr || accel_attr.getValue() != "neura") { + return; + } // Converts block argument types to predicated values. func.walk([&](Block *block) { // skips the entry (first) block of the function. @@ -70,8 +74,8 @@ struct LeveragePredicatedValuePass private: // Gets operations in topological order. - void getOperationsInTopologicalOrder(func::FuncOp func, - SmallVector &ordered) { + void getOperationsInTopologicalOrder(FunctionOpInterface func, + SmallVector &ordered) { DenseSet visited; func.walk([&](Operation *op) { // Uses standard DFS to build topological order. diff --git a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp index db569bdf..5518005f 100644 --- a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp +++ b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp @@ -4,6 +4,7 @@ #include "NeuraDialect/NeuraPasses.h" #include "NeuraDialect/NeuraTypes.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/IR/Block.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Dominance.h" @@ -132,12 +133,13 @@ struct ControlFlowInfo { // Checks if all the live-out values in a block are dominated by the block's // arguments. -void assertLiveOutValuesDominatedByBlockArgs(func::FuncOp &func) { +void assertLiveOutValuesDominatedByBlockArgs(Region ®ion) { llvm::errs() << "[ctrl2data] Asserting live-out values dominated by block arguments\n"; - for (Block &block : func.getBlocks()) { - if (&block == &func.getBody().front()) + for (Block &block : region) { + if (&block == ®ion.front()) { continue; + } llvm::errs() << "[ctrl2data] Checking block: " << block << "\n"; DenseSet live_out_values; @@ -199,9 +201,9 @@ void assertLiveOutValuesDominatedByBlockArgs(func::FuncOp &func) { } // Builds control flow info for the given function. -void buildControlFlowInfo(func::FuncOp &func, ControlFlowInfo &ctrl_info, +void buildControlFlowInfo(Region ®ion, ControlFlowInfo &ctrl_info, DominanceInfo &dom_info) { - for (Block &block : func.getBlocks()) { + for (Block &block : region) { Operation *terminator = block.getTerminator(); if (auto cond_br = dyn_cast(terminator)) { @@ -305,12 +307,11 @@ Value getPrecessedCondition(Value condition, bool is_not_condition, return not_condition; } -void createReserveAndPhiOps( - func::FuncOp &func, ControlFlowInfo &ctrl_info, - llvm::MapVector &arg_to_reserve, - llvm::MapVector &arg_to_phi_result, - OpBuilder &builder) { - DominanceInfo dom_info(func); +void createReserveAndPhiOps(Region ®ion, ControlFlowInfo &ctrl_info, + llvm::MapVector &arg_to_reserve, + llvm::MapVector &arg_to_phi_result, + OpBuilder &builder) { + DominanceInfo dom_info(region.getParentOp()); // ================================================ // Step 1: Categorizes edges into six types. @@ -574,18 +575,18 @@ void createReserveAndPhiOps( } // Transforms control flow into data flow. -void transformControlFlowToDataFlow(func::FuncOp &func, +void transformControlFlowToDataFlow(Region ®ion, ControlFlowInfo &ctrl_info, DominanceInfo &dom_info, OpBuilder &builder) { // Asserts that all live-out values are dominated by block arguments. - assertLiveOutValuesDominatedByBlockArgs(func); + assertLiveOutValuesDominatedByBlockArgs(region); // Creates reserve and phi operations for each block argument. llvm::MapVector arg_to_reserve; llvm::MapVector arg_to_phi_result; - createReserveAndPhiOps(func, ctrl_info, arg_to_reserve, arg_to_phi_result, + createReserveAndPhiOps(region, ctrl_info, arg_to_reserve, arg_to_phi_result, builder); // Replaces blockarguments with phi results @@ -596,9 +597,9 @@ void transformControlFlowToDataFlow(func::FuncOp &func, } // Flattens blocks into the entry block. - Block *entryBlock = &func.getBody().front(); + Block *entryBlock = ®ion.front(); SmallVector blocks_to_flatten; - for (Block &block : func.getBody()) { + for (Block &block : region) { if (&block != entryBlock) blocks_to_flatten.push_back(&block); } @@ -650,23 +651,42 @@ struct TransformCtrlToDataFlowPass void getDependentDialects(DialectRegistry ®istry) const override { registry.insert(); + registry.insert(); } void runOnOperation() override { ModuleOp module = getOperation(); - module.walk([&](func::FuncOp func) { - OpBuilder builder(func.getContext()); - GrantPredicateInEntryBlock(&func.getBody().front(), builder); - - DominanceInfo dom_info(func); - - // Step 1: Analyzes the control flow and creates control flow info - // struct. - ControlFlowInfo ctrl_info; - buildControlFlowInfo(func, ctrl_info, dom_info); - - // Step 2: Transforms control flow into data flow. - transformControlFlowToDataFlow(func, ctrl_info, dom_info, builder); + module.walk([&](Operation *op) { + Region *region = nullptr; + DominanceInfo domInfo; + OpBuilder builder(op->getContext()); + + if (auto func = dyn_cast(op)) { + auto accel_attr = func->getAttrOfType("accelerator"); + if (!accel_attr || accel_attr.getValue() != "neura") { + return; + } + region = &func.getBody(); + domInfo = DominanceInfo(func); + GrantPredicateInEntryBlock(®ion->front(), builder); + assertLiveOutValuesDominatedByBlockArgs(*region); + } else if (auto llvmFunc = dyn_cast(op)) { + if (llvmFunc.isDeclaration()) return; + auto accel_attr = llvmFunc->getAttrOfType("accelerator"); + if (!accel_attr || accel_attr.getValue() != "neura") { + return; + } + region = &llvmFunc.getBody(); + domInfo = DominanceInfo(llvmFunc); + GrantPredicateInEntryBlock(®ion->front(), builder); + assertLiveOutValuesDominatedByBlockArgs(*region); + // Skips SSA live-out dominance assert. + } else { + return; + } + ControlFlowInfo ctrlInfo; + buildControlFlowInfo(*region, ctrlInfo, domInfo); + transformControlFlowToDataFlow(*region, ctrlInfo, domInfo, builder); }); } }; diff --git a/test/neura/for_loop/test.mlir b/test/neura/for_loop/test.mlir index b3c00bd5..ae5dfd8b 100644 --- a/test/neura/for_loop/test.mlir +++ b/test/neura/for_loop/test.mlir @@ -4,14 +4,36 @@ // TODO: Enable --insert-mov once the backward ctrl flow mov is supported. // Lowers to neura. +// TODO: Make `--leverage-predicated-value` work. Segmentation fault for now. +// https://github.com/coredac/dataflow/issues/84. // RUN: mlir-neura-opt \ // RUN: --assign-accelerator \ // RUN: --lower-llvm-to-neura \ -// RN: --transform-ctrl-to-data-flow \ +// RN: --leverage-predicated-value \ +// RUN: --transform-ctrl-to-data-flow \ // RUN: --fuse-patterns \ // RN: --insert-mov \ -// RUN: %t-kernel.mlir | FileCheck %s +// RUN: %t-kernel.mlir > %t-lowered.mlir + +// RUN: FileCheck %s < %t-lowered.mlir // Verifies the neura ops are generated. And fusion happens. -// CHECK: accelerator = "neura" -// CHECK-NOT: = llvm. +// CHECK: accelerator = "neura" +// CHECK-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> i64 +// CHECK-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> i64 +// CHECK-NEXT: %2 = "neura.constant"() <{predicate = true, value = 32 : i64}> : () -> i64 +// CHECK-NEXT: %3 = neura.reserve : i64 +// CHECK-NEXT: %4 = "neura.phi"(%3, %0) : (i64, i64) -> i64 +// CHECK-NEXT: %5 = "neura.gep"(%arg0, %4) : (!llvm.ptr, i64) -> !llvm.ptr +// CHECK-NEXT: %6 = "neura.load"(%5) : (!llvm.ptr) -> f32 +// CHECK-NEXT: %7 = "neura.gep"(%arg2, %4) : (!llvm.ptr, i64) -> !llvm.ptr +// CHECK-NEXT: %8 = "neura.load"(%7) : (!llvm.ptr) -> f32 +// CHECK-NEXT: %9 = "neura.load"(%arg1) : (!llvm.ptr) -> f32 +// CHECK-NEXT: %10 = "neura.fmul_fadd"(%6, %8, %9) : (f32, f32, f32) -> f32 +// CHECK-NEXT: "neura.store"(%10, %arg1) : (f32, !llvm.ptr) -> () +// CHECK-NEXT: %11 = "neura.add"(%4, %1) : (i64, i64) -> i64 +// CHECK-NEXT: %12 = "neura.icmp"(%11, %2) <{cmpType = "eq"}> : (i64, i64) -> i1 +// CHECK-NEXT: %13 = "neura.not"(%12) : (i1) -> i1 +// CHECK-NEXT: %14 = neura.grant_predicate %11, %13 : i64, i1 -> i64 +// CHECK-NEXT: neura.ctrl_mov %14 -> %3 : i64 i64 +// CHECK-NEXT: "neura.return"() : () -> () \ No newline at end of file