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