Add DMA loop subsumption transformation

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

@@ -255,6 +255,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(
@@ -496,7 +537,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"