Fix annotations

Signed-off-by: Julia Koval <c_yuliak@xsightlabs.com>

Author: vaalfreja

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td

@@ -540,7 +540,8 @@ def CmpOp : P4HIR_Op<"cmp",
 def ScopeOp : P4HIR_Op<"scope", [
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        RecursivelySpeculatable,
-       RecursiveMemoryEffects, AutomaticAllocationScope,
+       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       AutomaticAllocationScope,
        NoRegionArguments, Annotated]> {
   let summary = "Represents a P4 scope";
   let description = [{

lib/Dialect/P4HIR/P4HIR_Ops.cpp

@@ -926,6 +926,28 @@ LogicalResult P4HIR::ScopeOp::verify() {
     return success();
 }
 
+void P4HIR::ScopeOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) {
+    if (auto annotations = getAnnotations()) {
+        auto attrs = annotations->getValue();
+        bool isOnlyAtomic = (attrs.size() == 1 && attrs[0].getName() == "atomic");
+        if (!isOnlyAtomic) {
+            effects.emplace_back(MemoryEffects::Write::get(), SideEffects::DefaultResource::get());
+            return;
+        }
+    }
+
+    for (Region &region : getOperation()->getRegions()) {
+        for (Block &block : region) {
+            for (Operation &op : block) {
+                if (auto opEffects = mlir::getEffectsRecursively(&op)) {
+                    effects.append(opEffects->begin(), opEffects->end());
+                }
+            }
+        }
+    }
+}
+
 LogicalResult P4HIR::ScopeOp::canonicalize(P4HIR::ScopeOp op, PatternRewriter &rewriter) {
     // Canonicalize scope: one without variables could be inlined
     if (op.getOps<VariableOp>().empty() && op.getScopeRegion().hasOneBlock() &&

lib/Transforms/FlattenCFG.cpp

@@ -86,9 +86,7 @@ class ScopeOpFlattening : public mlir::OpRewritePattern<P4HIR::ScopeOp> {
         mlir::OpBuilder::InsertionGuard guard(rewriter);
         auto loc = scopeOp.getLoc();
 
-        // Empty scope: just remove it.
-        // TODO: Decide if we'd need to do something with annotated scopes
-        if (scopeOp.isEmpty()) {
+        if (scopeOp.isEmpty() && !scopeOp.getAnnotations()) {
             rewriter.eraseOp(scopeOp);
             return mlir::success();
         }

test/Transforms/Folds/scope_simplification.mlir

@@ -68,12 +68,36 @@ module {
 
   // CHECK-LABEL: @empty_scope_with_annotations_preserved
   p4hir.func @empty_scope_with_annotations_preserved(%arg0: !ref_i32i) {
-    // TODO: Doesn't work(but should)
+    // Empty scopes with annotations should be preserved
+    // CHECK: p4hir.scope annotations {name = "test"} {
+    // CHECK-NEXT: }
+    // CHECK-NEXT: p4hir.return
     p4hir.scope annotations {name = "test"} {
     }
     p4hir.return
   }
 
+  // CHECK-LABEL: @empty_scope_with_atomic_removed
+  p4hir.func @empty_scope_with_atomic_removed(%arg0: !ref_i32i) {
+    // Empty scopes with only @atomic annotation should be removed
+    // CHECK-NOT: p4hir.scope
+    // CHECK: p4hir.return
+    p4hir.scope annotations {atomic} {
+    }
+    p4hir.return
+  }
+
+  // CHECK-LABEL: @empty_scope_with_atomic_and_name_preserved
+  p4hir.func @empty_scope_with_atomic_and_name_preserved(%arg0: !ref_i32i) {
+    // Empty scopes with @atomic plus other annotations should be preserved
+    // CHECK: p4hir.scope annotations {atomic, name = "debug"} {
+    // CHECK-NEXT: }
+    // CHECK-NEXT: p4hir.return
+    p4hir.scope annotations {atomic, name = "debug"} {
+    }
+    p4hir.return
+  }
+
   // CHECK-LABEL: @scope_in_if_else_without_vars
   p4hir.func @scope_in_if_else_without_vars(%cond: !p4hir.bool, %arg0: !ref_i32i) {
     // Scope without variables/annotations inside if else branch should be inlined