Multi-Dimensional Buffer Descriptors in AIE2

include/aie/Dialect/AIE/IR/AIE.td

@@ -40,6 +40,7 @@ switch is referred to as `switchbox` to avoid confusion with the
 
   }];
   let useDefaultTypePrinterParser = 1;
+  let useDefaultAttributePrinterParser = 1;
   let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
@@ -797,6 +798,20 @@ def AIE_DMABDPACKETOp: AIE_Op<"dmaBdPacket", []> {
 
 }
 
+def AIE_DimTupleAttr : AttrDef<AIE_Dialect, "DimTuple", []> {
+  let mnemonic = "DimTuple";
+  let summary = 
+    "Tuple encoding the stride and wrap of one dimension in an "
+    "AIE2 n-dimensional buffer descriptor";
+  let parameters = (ins
+    "uint32_t" : $stepsize,
+    "uint16_t" : $wrap
+  );
+  let assemblyFormat = "`<` $stepsize `,` $wrap `>`";
+}
+
+def AIE_DimTupleArrayAttr : ArrayOfAttr<AIE_Dialect, "DimTupleArray", "DimTupleArray", "::xilinx::AIE::DimTupleAttr">;
+
 def AIE_DMABDOp: AIE_Op<"dmaBd", []> {
   let summary = "Declare a dma block descriptor op";
   let description = [{
@@ -823,17 +838,44 @@ def AIE_DMABDOp: AIE_Op<"dmaBd", []> {
     ```
     A DMA channel in a Memory Module can process one block descriptor after another by chaining them.
     There are 16 block descriptors per Memory Module. They are shared by four DMA channels.
+
+    On AIE-ML devices, an optional argument can be used to specify an array of 
+    step sizes and wraps to move data in more advanced patterns. Strides and
+    wraps are specified as tuples `<stride, wrap>`, and up to three dimensions
+    can be specified (or up to four dimensions on memtiles). 
+
+    The first element of the array gives the _highest-dimension_ stride and
+    wrap, the last element of the array gives the lowest-dimension.
+
+    Example:
+
+    ```
+    AIE.dmaBd(<%buf : memref<128xi32>, 0, 128>, 0, [<16, 8>, <1, 2>, <2, 8>])
+    ```
+
+    This corresponds to alternating between even and odd elements of the 
+    buffer/stream every 8 elements, like so, equivalent to nested loops like so:
+
+    ```
+    for(int i = 0; i < 8 /* wrap[0] */; i += 16 /* stepsize[0] */)
+      for(int j = 0; j < 2 /* wrap[1] */; j += 1 /* stepsize[1] */)
+        for(int k = 0; k < 8 /* wrap[2] */; k += 2 /* stepsize[2] */)
+          // access/store element at/to index (i + j + k)
+    ```
   }];
 
   let arguments = (
     ins AnyMemRef:$buffer,
         I32Attr:$offset,
         I32Attr:$len,
-        ConfinedAttr<I32Attr, [IntMinValue<0>, IntMaxValue<1>]>:$AB // 0: A, 1: B
+        ConfinedAttr<I32Attr, [IntMinValue<0>, IntMaxValue<1>]>:$AB, // 0: A, 1: B
+        OptionalAttr<AIE_DimTupleArrayAttr>:$dimensions
   );
 
+  let hasVerifier = 1;
+
   let assemblyFormat = [{
-    `(` `<` $buffer  `:` type($buffer) `,` $offset `,` $len `>` `,` $AB `)` attr-dict
+    `(` `<` $buffer  `:` type($buffer) `,` $offset `,` $len `>` `,` $AB (`,` $dimensions^ )? `)` attr-dict
   }];
 
   let extraClassDeclaration = [{
@@ -843,15 +885,14 @@ def AIE_DMABDOp: AIE_Op<"dmaBd", []> {
     bool isB() { return (getAB() == 1); }
   }];
 
-  // let builders = [
-  //   OpBuilder<(ins "Value":$buffer, "int":$offset, "int":$len, "int":$AB), [{
-  //     build($_builder, $_state,
-  //           buffer,
-  //           $_builder.getI32IntegerAttr(offset),
-  //           $_builder.getI32IntegerAttr(len),
-  //           $_builder.getI32IntegerAttr(AB));
-  //     }]>
-  // ];
+  let builders = [
+    OpBuilder<(ins "Value":$buffer, "int":$offset, "int":$len, "int":$AB), [{
+      odsState.addOperands(buffer);
+      odsState.addAttribute(getOffsetAttrName(odsState.name), $_builder.getI32IntegerAttr(offset));
+      odsState.addAttribute(getLenAttrName(odsState.name), $_builder.getI32IntegerAttr(len));
+      odsState.addAttribute(getABAttrName(odsState.name), $_builder.getI32IntegerAttr(AB));
+      }]>
+  ];
 }
 
 def AIE_DMAStartOp: AIE_Op<"dmaStart", [ParentOneOf<["MemOp", "MemTileDMAOp", "func::FuncOp", "ShimDMAOp"]>, Terminator]>,

include/aie/Dialect/AIE/IR/AIEDialect.h

@@ -65,67 +65,27 @@ class TileOp;
 /// Include the generated interface declarations.
 #include "aie/Dialect/AIE/IR/AIEInterfaces.h.inc"
 
+// Include dialect declarations such as parseAttributes, parseType
+#include "aie/Dialect/AIE/IR/AIEDialect.h.inc"
+
 namespace xilinx {
 namespace AIE {
 
 void registerAIETranslations();
 
-// FIXME: use this
-//#include "AIEDialect.h.inc"
-
-// The Dialect
-class AIEDialect : public mlir::Dialect {
-public:
-  explicit AIEDialect(mlir::MLIRContext *ctx);
-  static StringRef getDialectNamespace() { return "AIE"; }
-
-  /// Parse a type registered to this dialect. Overridding this method is
-  /// required for dialects that have custom types.
-  /// Technically this is only needed to be able to round-trip to textual IR.
-  mlir::Type parseType(DialectAsmParser &parser) const override;
-
-  /// Print a type registered to this dialect. Overridding this method is
-  /// only required for dialects that have custom types.
-  /// Technically this is only needed to be able to round-trip to textual IR.
-  void printType(mlir::Type type, DialectAsmPrinter &os) const override;
-};
+} // namespace AIE
+} // namespace xilinx
 
 ////////////////////////////////////////////////////////////////////////////////
 /////////////////////// Custom Types for the Dialect ///////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 
-// namespace detail {
-// struct AIEListTypeStorage;
-// }
-
-// /// LLVM-style RTTI: one entry per subclass to allow dyn_cast/isa.
-// enum AIETypeKind {
-//   // The enum starts at the range reserved for this dialect.
-//   AIE_TYPE = mlir::Type::FIRST_PRIVATE_EXPERIMENTAL_0_TYPE,
-//   AIE_LIST,
-// };
-
-// /// Type for Toy arrays.
-// /// In MLIR Types are reference to immutable and uniqued objects owned by the
-// /// MLIRContext. As such `AIEListType` only wraps a pointer to an uniqued
-// /// instance of `AIEListTypeStorage` (defined in our implementation file) and
-// /// provides the public facade API to interact with the type.
-// class AIEListType : public mlir::Type::TypeBase<AIEListType, mlir::Type,
-//                                                  detail::AIEListTypeStorage>
-//                                                  {
-// public:
-//   using Base::Base;
-
-//   /// Return the type of individual elements in the array.
-//   mlir::Type getElementType();
-
-//   /// Get the unique instance of this Type from the context.
-//   static AIEListType get(mlir::Type elementType);
-
-//   /// Support method to enable LLVM-style RTTI type casting.
-//   static bool kindof(unsigned kind) { return kind == AIETypeKind::AIE_LIST; }
-// };
+// Include generated TableGen-generated type definitions
+#define GET_TYPEDEF_CLASSES 1
+#include "aie/Dialect/AIE/IR/AIETypes.h.inc"
 
+namespace xilinx {
+namespace AIE {
 namespace detail {
 struct AIEObjectFifoTypeStorage;
 }
@@ -172,11 +132,24 @@ class AIEObjectFifoSubviewType
   mlir::Type getElementType();
 };
 
+} // namespace AIE
+} // namespace xilinx
+
+////////////////////////////////////////////////////////////////////////////////
+// Custom Attributes ///////////////////////////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+#define GET_ATTRDEF_CLASSES
+#include "aie/Dialect/AIE/IR/AIEAttrDefs.h.inc"
+
 ////////////////////////////////////////////////////////////////////////////////
 //////////////////// Custom Operations for the Dialect /////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 
-//#include "AIEOpInterfaces.h.inc"
+namespace xilinx {
+namespace AIE {
+
+// #include "AIEOpInterfaces.h.inc"
 
 typedef std::pair<WireBundle, int> Port;
 typedef std::pair<Port, Port> Connect;

include/aie/Dialect/AIE/IR/CMakeLists.txt

@@ -8,6 +8,11 @@
 add_mlir_dialect(AIE AIE)
 add_mlir_doc(AIE AIEDialect ./ -gen-dialect-doc)
 
+set(LLVM_TARGET_DEFINITIONS AIE.td)
+mlir_tablegen(AIEAttrDefs.h.inc -gen-attrdef-decls -attrdefs-dialect=AIE)
+mlir_tablegen(AIEAttrDefs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=AIE)
+add_public_tablegen_target(MLIRAIEAttrDefsIncGen)
+
 set(LLVM_TARGET_DEFINITIONS AIE.td)
 mlir_tablegen(AIEEnums.h.inc -gen-enum-decls)
 mlir_tablegen(AIEEnums.cpp.inc -gen-enum-defs)
@@ -17,9 +22,3 @@ set(LLVM_TARGET_DEFINITIONS AIEInterfaces.td)
 mlir_tablegen(AIEInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(AIEInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRAIEInterfacesIncGen)
-
-# set(LLVM_TARGET_DEFINITIONS AIE.td)
-# mlir_tablegen(LLVMConversions.inc -gen-llvmir-conversions)
-# mlir_tablegen(LLVMConversionEnumsToLLVM.inc -gen-enum-to-llvmir-conversions)
-# mlir_tablegen(LLVMConversionEnumsFromLLVM.inc -gen-enum-from-llvmir-conversions)
-# add_public_tablegen_target(MLIRLLVMConversionsIncGen)

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -16,9 +16,14 @@
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
+// Add TableGen'erated dialect definitions (including constructor)
+// We implement the initialize() function further below
+#include "aie/Dialect/AIE/IR/AIEDialect.cpp.inc"
+
 namespace {
 
 struct AIEInlinerInterface : public DialectInlinerInterface {
@@ -330,11 +335,16 @@ void AIEDialect::printType(mlir::Type type,
   }
 }
 
-// FIXME: use Tablegen'd dialect class
-AIEDialect::AIEDialect(mlir::MLIRContext *ctx)
-    : mlir::Dialect("AIE", ctx, ::mlir::TypeID::get<AIEDialect>()) {
-  // addTypes<AIEListType>();
+void AIEDialect::initialize() {
+  addTypes<
+#define GET_TYPE_LIST
+#include "aie/Dialect/AIE/IR/AIETypes.cpp.inc"
+      >();
   addTypes<AIEObjectFifoType, AIEObjectFifoSubviewType>();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "aie/Dialect/AIE/IR/AIEAttrDefs.cpp.inc"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "aie/Dialect/AIE/IR/AIE.cpp.inc"
@@ -624,9 +634,7 @@ const xilinx::AIE::AIETargetModel &xilinx::AIE::DeviceOp::getTargetModel() {
   return VC1902model;
 }
 
-LogicalResult xilinx::AIE::DeviceOp::verify() {
-  return success();
-}
+LogicalResult xilinx::AIE::DeviceOp::verify() { return success(); }
 
 LogicalResult xilinx::AIE::TileOp::verify() {
   const auto &target_model = getTargetModel(*this);
@@ -987,6 +995,64 @@ LogicalResult xilinx::AIE::MemTileDMAOp::verify() {
   return success();
 }
 
+// DMABDOp
+LogicalResult xilinx::AIE::DMABDOp::verify() {
+  // The following checks only apply if non-default strides/wraps are defined.
+  if (getDimensions()) {
+    ::mlir::MemRefType buffer = getBuffer().getType();
+    // We are restrictive about the type of the memref used as the input address
+    // to the DMABD when used with multi-dimensional strides/wraps. Since the
+    // BD will use the memref as a base address and copy from it in 32 bit
+    // chunks, while assuming the layout of the memref is contiguous, we
+    // disallow anything whose elemental size is not 32 bits, or where we
+    // cannot verify that the layout is contiguous.
+    if (!buffer.getElementType().isInteger(32) || buffer.getRank() > 1 ||
+        !buffer.getLayout().isIdentity()) {
+      return emitOpError() << "Specifying transfer step sizes and wraps is only"
+                              " supported for one-dimensional memrefs of 32 bit"
+                              " integer elements.";
+    }
+    uint64_t memref_size = 1; // in bytes
+    uint64_t max_idx = 0;
+    for (int64_t memref_dim : buffer.getShape()) {
+      memref_size *= 4 * memref_dim;
+    }
+    llvm::ArrayRef<xilinx::AIE::DimTupleAttr> dims = *getDimensions();
+    size_t max_n_dims = 3;
+    if (isa_and_nonnull<xilinx::AIE::MemTileDMAOp>((*this)->getParentOp())) {
+      max_n_dims = 4;
+    }
+    if (dims.size() > max_n_dims) {
+      return emitOpError() << "Cannot give more than "
+                           << std::to_string(max_n_dims)
+                           << " dimensions for step sizes and wraps in this "
+                              " tile (got "
+                           << std::to_string(dims.size()) << " dimensions).";
+    }
+    for (xilinx::AIE::DimTupleAttr dim : dims) {
+      max_idx += dim.getStepsize() * (dim.getWrap() - 1);
+      if (0 == dim.getStepsize()) {
+        return emitOpError()
+               << "Invalid step size; must be a positive integer.";
+      }
+      if (dim.getStepsize() > memref_size) {
+        return emitOpError()
+               << "Step size " << std::to_string(dim.getStepsize() * 4) << " "
+               << "bytes exceeds memref size " << std::to_string(memref_size);
+      }
+    }
+    if (memref_size <= 4 * max_idx) {
+      return emitOpError() << "Specified stepsize(s) and wrap(s) result in out "
+                              "of bounds access in buffer, for index "
+                           << std::to_string(max_idx) << ", accessing at "
+                           << std::to_string(4 * max_idx)
+                           << " byte offset in memref of length "
+                           << std::to_string(memref_size) << ".";
+    }
+  }
+  return success();
+}
+
 xilinx::AIE::TileOp xilinx::AIE::MemTileDMAOp::getTileOp() {
   return cast<xilinx::AIE::TileOp>(getTile().getDefiningOp());
 }
@@ -1206,3 +1272,7 @@ bool TileOp::isShimPLTile() {
 bool TileOp::isShimNOCorPLTile() { return isShimNOCTile() || isShimPLTile(); }
 } // namespace AIE
 } // namespace xilinx
+
+// Include implementations for custom attributes
+#define GET_ATTRDEF_CLASSES
+#include "aie/Dialect/AIE/IR/AIEAttrDefs.cpp.inc"

lib/Dialect/AIE/IR/CMakeLists.txt

@@ -13,6 +13,7 @@ add_mlir_dialect_library(AIE
 
   DEPENDS
   MLIRAIEIncGen
+  MLIRAIEAttrDefsIncGen
   MLIRAIEEnumsIncGen
   MLIRAIEInterfacesIncGen
 

lib/Dialect/AIEX/IR/CMakeLists.txt

@@ -15,6 +15,7 @@ add_mlir_dialect_library(AIEX
   MLIRAIEXEnumsIncGen
 
   LINK_LIBS PUBLIC
+  AIE
   MLIRIR
   MLIRSupport
   )