Trace BufferOp through memref::CastOp (Xilinx#742)

nod-ai · Oct 21, 2024 · f2a164a · f2a164a
1 parent f4322ee
commit f2a164a
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp b/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp
@@ -577,6 +577,10 @@ AIE::BufferOp TileDMAAllocator::getBuffer(uint64_t, int64_t col, int64_t row,
                      ? (memcpyOp.getDstMemref())
                      : (memcpyOp.getSrcMemref());
   AIE::BufferOp bufferOp = buffer.getDefiningOp<AIE::BufferOp>();
+  // Memref cast
+  memref::CastOp castOp = buffer.getDefiningOp<memref::CastOp>();
+  if (!bufferOp && castOp)
+    bufferOp = castOp.getOperand().getDefiningOp<AIE::BufferOp>();
   return bufferOp;
 }
 
@@ -856,6 +860,10 @@ AIE::BufferOp MemTileDMAAllocator::getBuffer(uint64_t, int64_t col, int64_t row,
                      ? (memcpyOp.getDstMemref())
                      : (memcpyOp.getSrcMemref());
   AIE::BufferOp bufferOp = buffer.getDefiningOp<AIE::BufferOp>();
+  // Memref cast
+  memref::CastOp castOp = buffer.getDefiningOp<memref::CastOp>();
+  if (!bufferOp && castOp)
+    bufferOp = castOp.getOperand().getDefiningOp<AIE::BufferOp>();
   return bufferOp;
 }