Folding no-op stream.async.update ops away. (iree-org#17458)

nod-ai · May 21, 2024 · 6c5198d · 6c5198d
1 parent 006af5d
commit 6c5198d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
@@ -1740,6 +1740,10 @@ OpFoldResult AsyncUpdateOp::fold(FoldAdaptor operands) {
       getUpdate().getType() == getType() && hasValueSemantics(getTarget())) {
     return getUpdate();
   }
+  // If copying to self then elide the copy.
+  if (getUpdate() == getTarget() && getUpdateSize() == getTargetSize()) {
+    return getTarget();
+  }
   return {};
 }
 

diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir
@@ -288,6 +288,19 @@ util.func private @FoldLocalAsyncUpdateOp(%arg0: !stream.resource<*>, %arg1: ind
 
 // -----
 
+// Tests that updates of a value into itself are no-oped.
+
+// CHECK-LABEL: @FoldNoOpAsyncUpdateOp
+util.func private @FoldNoOpAsyncUpdateOp(%arg0: !stream.resource<*>, %arg1: index) -> !stream.resource<*> {
+  %c0 = arith.constant 0 : index
+  // CHECK-NOT: stream.async.update
+  %0 = stream.async.update %arg0, %arg0[%c0 to %arg1] : !stream.resource<*>{%arg1} -> %arg0 as !stream.resource<*>{%arg1}
+  // CHECK: util.return %arg0
+  util.return %0 : !stream.resource<*>
+}
+
+// -----
+
 // CHECK-LABEL: @ElideInPlaceUpdateUpdate
 util.func private @ElideInPlaceUpdateUpdate(%arg0: !stream.resource<*>, %arg1: index, %arg2: !stream.resource<*>, %arg3: index) -> !stream.resource<*> {
   %c0 = arith.constant 0 : index