Merge pull request coredac#223 from ShangkunLi/safe-return

Enable the safe return in neura dataflow ir

Author: ShangkunLi

include/NeuraDialect/NeuraOps.td

@@ -59,7 +59,7 @@ def Neura_KernelOp : Op<NeuraDialect, "kernel", [RecursiveMemoryEffects, SingleB
 }
 
 // Yield operation for fused_op and kernel regions.
-def Neura_YieldOp : Op<NeuraDialect, "yield", [Terminator, Pure, ReturnLike, HasParent<"FusedOp, KernelOp">]> {
+def Neura_YieldOp : Op<NeuraDialect, "yield", [Terminator, Pure, ReturnLike]> {
   let summary = "Yield values from a neura.kernel or neura.fused_op region.";
   let description = [{
     Returns values from a neura.kernel or neura.fused_op region to the parent operation.
@@ -79,6 +79,8 @@ def Neura_YieldOp : Op<NeuraDialect, "yield", [Terminator, Pure, ReturnLike, Has
   ];
 
   let assemblyFormat = [{($values^ `:` type($values))? attr-dict}];
+
+  let hasVerifier = 1;
 }
 
 
@@ -375,6 +377,46 @@ def Neura_ReturnOp : Op<NeuraDialect, "return", [Terminator]> {
   // let assemblyFormat = "($values^)? `,` attr-dict";
 }
 
+// Defines a return operation for void functions in dataflow mode.
+def Neura_ReturnVoidOp : Op<NeuraDialect, "return_void", []>{
+  let summary = "Return operation for void functions in dataflow mode.";
+  let description = [{
+    Terminates a void function. The trigger operand provides the execution
+    condition in dataflow mode - the return executes when the trigger becomes valid.
+    
+    Without a trigger (empty operands), this represents a static return that
+    always executes. After canonicalization, void returns should have triggers.
+
+    Example:
+      neura.return_void %cond : !neura.data<i1, i1>  // In dataflow mode
+      ...
+      neura.yield                       // Function ends here.
+
+    This is NOT a terminator - the block ends with neura.yield instead.
+  }];
+  let arguments = (ins Optional<AnyType>:$trigger);
+  let assemblyFormat = "($trigger^ `:` type($trigger))? attr-dict";
+}
+
+// Defines a return operation for non-void fnctions in dataflow mode.
+def Neura_ReturnValueOp : Op<NeuraDialect, "return_value", []>{
+  let summary = "Return operation for non-void functions in dataflow mode.";
+  let description = [{
+    Represents a return point for non-void functions in dataflow mode.
+    The value operand provides the return value and the trigger operand.
+
+    This is NOT a terminator - the block ends with neura.yield instead.
+
+    Example:
+      neura.return_value %ret_val : !neura.data<i32, i1> // In dataflow mode
+      ...
+      neura.yield                       // Function ends here.
+  }];
+  let arguments = (ins Variadic<AnyType>:$values);
+  let assemblyFormat = "($values^ `:` type($values))? attr-dict";
+}
+
+
 // Defines a cast operation for type conversion.
 def Neura_CastOp : Op<NeuraDialect, "cast">{
   let summary = "Generic type conversion operation";

include/NeuraDialect/NeuraPasses.h

@@ -25,6 +25,7 @@ std::unique_ptr<mlir::Pass> createTransformCtrlToDataFlowPass();
 std::unique_ptr<mlir::Pass> createLeveragePredicatedValuePass();
 std::unique_ptr<mlir::Pass> createMapToAcceleratorPass();
 std::unique_ptr<mlir::Pass> createGenerateCodePass();
+std::unique_ptr<mlir::Pass> createCanonicalizeReturnPass();
 std::unique_ptr<mlir::Pass> createCanonicalizeLiveInPass();
 std::unique_ptr<mlir::Pass> createPromoteFuncArgToConstPass();
 std::unique_ptr<mlir::Pass> createTransformToSteerControlPass();

include/NeuraDialect/NeuraPasses.td

@@ -75,6 +75,18 @@ def FuseLoopControl: Pass<"fuse-loop-control", "ModuleOp">{
   let constructor = "neura::createFuseLoopControlPass()";
 }
 
+def CanonicalizeReturn : Pass<"canonicalize-return", "ModuleOp">{
+  let summary = "Adds execution conditions to void return operations";
+  let description = [{
+    This pass processes void functions and adds trigger conditions to their
+    return operations. For void functions, the return operation needs an
+    execution condition to properly trigger in dataflow mode.
+    
+    This pass should run before CanonicalizeLiveIn and TransformCtrlToDataFlow.
+  }];
+  let constructor = "neura::createCanonicalizeReturnPass()";
+}
+
 def CanonicalizeLiveIn : Pass<"canonicalize-live-in", "ModuleOp"> {
   let summary = "Canonicalizes live-in values/operations in each basic block.";
   let description = [{

lib/NeuraDialect/Mapping/mapping_util.cpp

@@ -126,7 +126,8 @@ OperationKind getOperationKindFromMlirOp(Operation *op) {
 // Returns true if the operation does not need CGRA tile placement.
 bool is_non_materialized(Operation *op) {
   // Returns true if the operation does not need CGRA tile placement.
-  return mlir::isa<neura::ReserveOp, neura::CtrlMovOp, neura::DataMovOp>(op);
+  return mlir::isa<neura::ReserveOp, neura::CtrlMovOp, neura::DataMovOp,
+                   neura::YieldOp>(op);
 }
 
 } // namespace neura

lib/NeuraDialect/NeuraOps.cpp

@@ -1,12 +1,35 @@
 #include "NeuraDialect/NeuraOps.h"
 #include "NeuraDialect/NeuraDialect.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/Support/LogicalResult.h"
 
 using namespace mlir;
 using namespace mlir::neura;
 
+LogicalResult YieldOp::verify() {
+  Operation *parent_op = (*this)->getParentOp();
+
+  if (!parent_op) {
+    return emitOpError("must have a parent operation.");
+  }
+
+  // Allows yield in FusedOp and KernelOp
+  if (isa<FusedOp>(parent_op) || isa<KernelOp>(parent_op)) {
+    return success();
+  }
+
+  // Allows yield in func.func (for dataflow mode)
+  if (isa<func::FuncOp>(parent_op)) {
+    return success();
+  }
+
+  return emitOpError("expects parent op to be one of 'neura.fused_op', "
+                     "'neura.kernel', or 'func.func'");
+}
+
 LogicalResult PhiStartOp::verify() {
   // Checks if this phi_start is inside a fused_op.
   Operation *parent_op = getOperation()->getParentOp();

lib/NeuraDialect/NeuraPasses.cpp

@@ -25,6 +25,7 @@ void mlir::neura::registerNeuraConversionPassPipeline() {
         pm.addPass(mlir::createLowerLlvmToNeuraPass());
         pm.addPass(mlir::createPrintOpGraphPass(os));
 
+        pm.addPass(mlir::neura::createCanonicalizeReturnPass());
         pm.addPass(mlir::neura::createCanonicalizeCastPass());
         pm.addPass(mlir::neura::createPromoteFuncArgToConstPass());
         pm.addPass(mlir::neura::createFoldConstantPass());

lib/NeuraDialect/Transforms/CMakeLists.txt

@@ -10,6 +10,7 @@ add_mlir_library(
     LeveragePredicatedValuePass.cpp
     MapToAcceleratorPass.cpp
     GenerateCodePass.cpp
+    CanonicalizeReturnPass.cpp
     CanonicalizeLiveInPass.cpp
     CanonicalizeCastPass.cpp
     PromoteFuncArgToConstPass.cpp