Remove `hal.interface.binding.subspan` ops that are only used `memref.assume_alignment_ops`. (#12269)
diff --git a/compiler/src/iree/compiler/Codegen/Common/CleanupBufferAllocViewPass.cpp b/compiler/src/iree/compiler/Codegen/Common/CleanupBufferAllocViewPass.cpp
index cd2a5fc..f34deac 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CleanupBufferAllocViewPass.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/CleanupBufferAllocViewPass.cpp
@@ -112,7 +112,18 @@
if (!memEffect || !memEffect.hasEffect<MemoryEffects::Allocate>()) {
return failure();
}
- if (!op->use_empty()) return failure();
+ SmallVector<Operation *> deadUsers;
+ for (OpOperand &use : op->getUses()) {
+ if (auto user = dyn_cast<memref::AssumeAlignmentOp>(use.getOwner())) {
+ deadUsers.push_back(user);
+ continue;
+ }
+ // For any other use, return failure;
+ return failure();
+ }
+ for (auto user : deadUsers) {
+ rewriter.eraseOp(user);
+ }
rewriter.eraseOp(op);
return success();
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/remove_dead_allocs.mlir b/compiler/src/iree/compiler/Codegen/Common/test/remove_dead_allocs.mlir
index 146d3d0..50040b1 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/remove_dead_allocs.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/remove_dead_allocs.mlir
@@ -16,3 +16,13 @@
// CHECK-LABEL: func.func @alloc_keep
// CHECK-NEXT: %[[ALLOC:.+]] = memref.alloc
// CHECK-NEXT: return %[[ALLOC]]
+
+// -----
+
+func.func @cleanup_only_assume_alignment_uses() {
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : memref<42xf32>
+ memref.assume_alignment %0, 64 : memref<42xf32>
+ return
+}
+// CHECK-LABEL: func.func @cleanup_only_assume_alignment_uses()
+// CHECK-NEXT: return
