Add DMA loop subsumption transformation (#512)

Addresses #495.

DMA loop iteration subsumption tries to move scf.for loops inside the
DMA operations by updating the DMA access patterns and hoisting them out
of the loop. See the issue above for more details.

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.td

@@ -241,28 +241,6 @@ def AMDAIE_NpuDmaCpyNdOp: AMDAIE_Op<"npu.dma_cpy_nd",
   let extraClassDeclaration = [{
     // Check whether this dma operation has a wait user.
     bool hasDmaWaitOpUser();
-
-    // Helper method to return the memory space of source as an integer. If
-    // no memory space attribute, this indicates a global memory space and
-    // we return 0. Else cast the memory space attribute to an integer. 
-    uint64_t getSourceMemorySpaceAsUInt() {
-      MemRefType sourceMemrefType =
-          cast<LogicalObjectFifoType>(getDmaCpyNdOp().getSourceType())
-          .getElementType();
-      Attribute memSpace = sourceMemrefType.getMemorySpace();
-      return memSpace ? dyn_cast<IntegerAttr>(memSpace).getInt() : 0;
-    }
-
-    // Helper method to return the memory space of target as an integer. If
-    // no memory space attribute, this indicates a global memory space and
-    // we return 0. Else cast the memory space attribute to an integer. 
-    uint64_t getTargetMemorySpaceAsUInt() {
-      MemRefType targetMemrefType =
-          cast<LogicalObjectFifoType>(getDmaCpyNdOp().getTargetType())
-          .getElementType();
-      Attribute memSpace = targetMemrefType.getMemorySpace();
-      return memSpace ? dyn_cast<IntegerAttr>(memSpace).getInt() : 0;
-    }
 
     // Check whether this operation has addressing on the source side.
     bool hasSourceAddressing() {
@@ -278,6 +256,47 @@ def AMDAIE_NpuDmaCpyNdOp: AMDAIE_Op<"npu.dma_cpy_nd",
     CircularDmaCpyNdOp getDmaCpyNdOp() {
       return dyn_cast<CircularDmaCpyNdOp>(getDma().getDefiningOp());
     }
+
+    // Return the source memref type. This is retrieved using information from
+    // the input DMA operation.
+    MemRefType getSourceMemrefType() { 
+      return cast<LogicalObjectFifoType>(getDmaCpyNdOp().getSourceType())
+          .getElementType();
+    }
+
+    // Return the source memory space as an attribute.
+    Attribute getSourceMemorySpace() {
+      return getSourceMemrefType().getMemorySpace();
+    }
+
+    // Helper method to return the source memory space as an integer. If no 
+    // memory space attribute, this indicates a global memory space and we
+    // return 0. Else cast the memory space attribute to an integer. 
+    uint64_t getSourceMemorySpaceAsUInt() {
+      Attribute memSpace = getSourceMemorySpace();
+      return memSpace ? cast<IntegerAttr>(memSpace).getInt() : 0;
+    }
+
+    // Return the target memref type. This is retrieved using information from
+    // the input DMA operation.
+    MemRefType getTargetMemrefType() { 
+      return cast<LogicalObjectFifoType>(getDmaCpyNdOp().getTargetType())
+          .getElementType();
+    }
+
+    // Return the target memory space as an attribute.
+    Attribute getTargetMemorySpace() {
+      return getTargetMemrefType().getMemorySpace();
+    }
+
+    // Helper method to return the target memory space as an integer. If no
+    // memory space attribute, this indicates a global memory space and we
+    // return 0. Else cast the memory space attribute to an integer. 
+    uint64_t getTargetMemorySpaceAsUInt() {
+      Attribute memSpace = getTargetMemorySpace();
+      return memSpace ? cast<IntegerAttr>(memSpace).getInt() : 0;
+    }
+
     // A utility to create a new doubly strided operation from this one with a
     // new set of source and target offsets, sizes and strides.
     DoublyStridedOpInterface createDoublyStridedOp(
@@ -519,7 +538,7 @@ def AMDAIE_LogicalObjectFifoFromMemrefOp
     // Else cast the memory space attribute to an integer. 
     uint64_t getMemorySpaceAsUInt() {
       Attribute memSpace = getMemorySpace();
-      return memSpace ? dyn_cast<IntegerAttr>(memSpace).getInt() : 0;
+      return memSpace ? cast<IntegerAttr>(memSpace).getInt() : 0;
     }
 
     // Return the source memref type.

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIETargetModel.h

@@ -0,0 +1,63 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// In absence of a complete hardware model interface, this file contains some
+// constants to describe hardware-related parameters used in transformations.
+// This is meant to be temporary.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef IREE_COMPILER_AMDAIE_TARGET_MODEL_H_
+#define IREE_COMPILER_AMDAIE_TARGET_MODEL_H_
+
+namespace mlir::iree_compiler::AMDAIE {
+
+//===----------------------------------------------------------------------===//
+//
+// DMA iteration dimensions
+//
+// DMAs support multi-dimensional addressing through buffer descriptors in two
+// ways:
+// 1. Intra-iteration access pattern. Specified via 'strides' ('steps' in buffer
+// descriptor lingo), 'sizes' ('wraps' in buffer descriptro lingo) and
+// 'padding'. When a DMA executes a buffer descriptor, it will access the data
+// (read/write) as specified by the intra-iteration access pattern.
+// 2. Inter-iteration access pattern. Specified via an iteration 'stride',
+// 'size' and 'current_iteration' ('stride' is the same as 'stepsize' and 'size'
+// is the same as 'wrap' in buffer descriptor lingo). Here, 'current_iteration'
+// keeps track of the current execution iteration of the buffer descriptor and
+// is incremented after buffer descriptor execution. the 'stride' is the offset
+// to be used for each execution of the buffer descriptor, relative to the
+// previous one. When 'iteration_current' is equal to 'size', the
+// 'iteration_current' is reset to zero.
+//
+// Although all DMAs use the same buffer descriptor format to describe the
+// execution configuration, the intra-iteration and inter-dimensions are
+// typically used for different purposes on different DMAs. See for example the
+// usage of these constants inside the DMA loop subsumption pass.
+//
+//===----------------------------------------------------------------------===//
+
+/// Shim DMAs support 3 intra-iteration dimensions + 1 inter-iteration
+/// dimension.
+static const int64_t kAMDAIEShimDmaNbIntraDims = 3;
+static const int64_t kAMDAIEShimDmaNbInterDims = 1;
+
+/// MemTile DMAs support 4 intra-iteration dimensions + 1 inter-iteration
+/// dimension.
+static const int64_t kAMDAIEMemTileDmaNbIntraDims = 4;
+static const int64_t kAMDAIEMemTileDmaNbInterDims = 1;
+
+/// Core DMAs support 3 intra-iteration dimensions + 1 inter-iteration
+/// dimension.
+static const int64_t kAMDAIECoreDmaNbIntraDims = 3;
+static const int64_t kAMDAIECoreDmaNbInterDims = 1;
+
+}  // namespace mlir::iree_compiler::AMDAIE
+
+#endif  // IREE_COMPILER_AMDAIE_TARGET_MODEL_H_

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/CMakeLists.txt

@@ -15,6 +15,7 @@ iree_cc_library(
     "AMDAIEAttrs.h"
     "AMDAIEDialect.h"
     "AMDAIEOps.h"
+    "AMDAIETargetModel.h"
     "AMDAIETypes.h"
   TEXTUAL_HDRS
     "AMDAIEAttrs.cpp.inc"