[AMDAIETemporaryAllocBufferization] Don't change operand of dealloc (n…

…od-ai#831)

Apologies -- I introduced this bug in
nod-ai#824

Basically, the operand of the memref.dealloc was becoming the
amdaie.buffer, it should remain the memref.alloc (canonicalization will
remove both the memref.dealloc and memref.alloc later, then).

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIETemporaryAllocBufferization.cpp

@@ -43,10 +43,12 @@ LogicalResult bufferizeTemporaryMemrefs(Operation *parentOp) {
       Operation *owner = operand.getOwner();
       bool inTile = owner->getParentOfType<CoreOp>() == coreOp;
       bool inGroup = owner->getParentOfType<WorkgroupOp>() == workgroupOp;
-      return inTile && inGroup;
+      bool isNotDealloc = !isa<memref::DeallocOp>(owner);
+      return inTile && inGroup && isNotDealloc;
     });
   }
 
+
   // Note: we don't erase allocs/deallocs, we leave this for canonicalization.
 
   return success();

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/temporary_alloc_bufferization.mlir

@@ -107,3 +107,25 @@ func.func @temp_buffer_2(){
   memref.dealloc %alloc : memref<4xi32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: @alloc_dealloc_left_untouched_for_later_canonicalization
+// CHECK: amdaie.buffer
+// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<4xi32>
+// CHECK: memref.dealloc %[[ALLOC]] : memref<4xi32>
+// CHECK: amdaie.end
+func.func @alloc_dealloc_left_untouched_for_later_canonicalization() {
+    %c0 = arith.constant 0 : index
+    %c2 = arith.constant 2 : index
+    %tile_0_2 = amdaie.tile(%c0, %c2)
+    %core_0_2 = amdaie.core(%tile_0_2, in : [], out : []) {
+      %alloc = memref.alloc() : memref<4xi32>
+      %reinterpret_cast = memref.reinterpret_cast %alloc
+          to offset: [0], sizes: [4], strides: [1] :
+          memref<4xi32> to memref<4xi32>
+      memref.dealloc %alloc : memref<4xi32>
+      amdaie.end
+    }
+    return
+}