[mlir][async]: Make async.execute operation with RecursiveMemoryEffects trait

[AviadCo]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Aviad Cohen (AviadCo)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/116544.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Async/IR/AsyncOps.td (+2-1)
• (added) mlir/test/Dialect/Async/canonicalize.mlir (+10)

diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
index 33b67921752346..011bffd02a90c5 100644
--- a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
+++ b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
@@ -38,7 +38,8 @@ def Async_ExecuteOp :
                                                  ["getEntrySuccessorOperands",
                                                   "areTypesCompatible"]>,
                        AttrSizedOperandSegments,
-                       AutomaticAllocationScope]> {
+                       AutomaticAllocationScope,
+                       Pure]> {
   let summary = "Asynchronous execute operation";
   let description = [{
     The `body` region attached to the `async.execute` operation semantically
diff --git a/mlir/test/Dialect/Async/canonicalize.mlir b/mlir/test/Dialect/Async/canonicalize.mlir
new file mode 100644
index 00000000000000..1a74eaa344e169
--- /dev/null
+++ b/mlir/test/Dialect/Async/canonicalize.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s -split-input-file -canonicalize | FileCheck %s
+
+// CHECK-NOT: async.execute
+
+func.func @empty_execute() {
+  %token = async.execute {
+    async.yield
+  }
+  return
+}

[llvmbot]

@llvm/pr-subscribers-mlir-async

Author: Aviad Cohen (AviadCo)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/116544.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Async/IR/AsyncOps.td (+2-1)
• (added) mlir/test/Dialect/Async/canonicalize.mlir (+10)

diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
index 33b67921752346..011bffd02a90c5 100644
--- a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
+++ b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
@@ -38,7 +38,8 @@ def Async_ExecuteOp :
                                                  ["getEntrySuccessorOperands",
                                                   "areTypesCompatible"]>,
                        AttrSizedOperandSegments,
-                       AutomaticAllocationScope]> {
+                       AutomaticAllocationScope,
+                       Pure]> {
   let summary = "Asynchronous execute operation";
   let description = [{
     The `body` region attached to the `async.execute` operation semantically
diff --git a/mlir/test/Dialect/Async/canonicalize.mlir b/mlir/test/Dialect/Async/canonicalize.mlir
new file mode 100644
index 00000000000000..1a74eaa344e169
--- /dev/null
+++ b/mlir/test/Dialect/Async/canonicalize.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s -split-input-file -canonicalize | FileCheck %s
+
+// CHECK-NOT: async.execute
+
+func.func @empty_execute() {
+  %token = async.execute {
+    async.yield
+  }
+  return
+}

[joker-eph]

Can you elaborate on the motivation and the correctness in the description of this PR please?
This is surprising to me and I don't quite understand why this op can be Pure at the moment. Thanks.

[AviadCo]

Can you elaborate on the motivation and the correctness in the description of this PR please? This is surprising to me and I don't quite understand why this op can be Pure at the moment. Thanks.

Sure, the case I have is scf.if with result which means that I must have an else region. For example:

     %res = scf.if cond -> (!async.token) {
      ...
    } else {
      %token = async.execute {
        async.yield
      }
      scf.yield %token : !async.token
    }

In some phase, I reach the point where the async is being inside if that has no longer have result, for example:

    scf.if cond {
      ...
      %token = async.execute {
          .....
      }
      // Token is being used
    } else {
      async.execute {
        async.yield
      }
    }

In this case, I would expect that the whole else will be eliminated (folder) as the execute has no memory effect inside it and has no users. With Pure this will be optimized to:

    scf.if cond {
        %token = async.execute {
            .....
        }
        // Token is being used
    }

Keeping my specific flow apart, it seems legic to remove the execute region if it has no usages and have no memory effects like scf.for does.

[AviadCo]

Can you elaborate on the motivation and the correctness in the description of this PR please? This is surprising to me and I don't quite understand why this op can be Pure at the moment. Thanks.

Sure, the case I have is scf.if with result which means that I must have an else region. For example:

     %res = scf.if cond -> (!async.token) {
      ...
    } else {
      %token = async.execute {
        async.yield
      }
      scf.yield %token : !async.token
    }

In some phase, I reach the point where the async is being inside if that has no longer have result, for example:

    scf.if cond {
      ...
      %token = async.execute {
          .....
      }
      // Token is being used
    } else {
      async.execute {
        async.yield
      }
    }

In this case, I would expect that the whole else will be eliminated (folder) as the execute has no memory effect inside it and has no users. With Pure this will be optimized to:

    scf.if cond {
        %token = async.execute {
            .....
        }
        // Token is being used
    }

Keeping my specific flow apart, it seems legic to remove the execute region if it has no usages and have no memory effects like scf.for does.

@joker-eph I confused between Pure and RecursiveMemoryEffects - please see latest change

[joker-eph]

RecursiveMemoryEffects is more appropriate, thanks.

I have a hard time convincing me about correctness right now, but I can't imagine what can go wrong either: "submitting a task" isn't really side-effect free, but here it is more specifically "submitting a side-effect free task". So I would tend to agree with this patch.

It'd be nice if @sanjoy or @River707 would give a second opinion :)