Merge branch 'main' into nm_xclbin_merge

Author: nirvedhmeshram

.github/workflows/ci.yml

@@ -103,7 +103,7 @@ jobs:
         run: |
           python3 -m venv .venv
           source .venv/bin/activate
-          pip install https://github.com/Xilinx/mlir-aie/releases/download/latest-wheels/mlir_aie-0.0.1.2024060222+6f70bfe-py3-none-manylinux_2_35_x86_64.whl
+          pip install https://github.com/Xilinx/mlir-aie/releases/download/latest-wheels/mlir_aie-0.0.1.2024061222+3ac9566-py3-none-manylinux_2_35_x86_64.whl
 
           pip install -r tests/matmul/requirements.txt
 

build_tools/ci/run_matmul_test.sh

@@ -638,7 +638,7 @@ run_matmul_test \
     --lhs_rhs_type "bf16" \
     --acc_type "f32" \
     --m "64"  --n "64" --k "128" \
-    --expect_compile_failure "1"
+    --num_repeat_runs "0"
 
 run_matmul_test \
     --name_prefix "packPeelLarge" \

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

@@ -41,6 +41,18 @@ def LogicalObjectFifoPort: I32EnumAttr<"LogicalObjectFifoPort",
   let cppNamespace = "mlir::iree_compiler::AMDAIE";
 }
 
+def MemoryAccess: I32EnumAttr<"MemoryAccess",
+  "The memory access type",
+  [
+    I32EnumAttrCase<"None", 0>,
+    I32EnumAttrCase<"Read", 1>,
+    I32EnumAttrCase<"Write", 2>,
+    I32EnumAttrCase<"Any", 3>,
+  ]
+  > {
+  let cppNamespace = "mlir::iree_compiler::AMDAIE";
+}
+
 def AMDAIE_MemSpace_Global : I32EnumAttrCase<"Global", 0>;
 def AMDAIE_MemSpace_Shared : I32EnumAttrCase<"Shared", 1>;
 def AMDAIE_MemSpace_Local : I32EnumAttrCase<"Local", 2>;

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

@@ -328,6 +328,22 @@ LogicalObjectFifoFromMemrefOp CircularDmaCpyNdOp::getTargetObjectFifo() {
   return dyn_cast<LogicalObjectFifoFromMemrefOp>(getTarget().getDefiningOp());
 };
 
+//===----------------------------------------------------------------------===//
+// AMDAIE_LogicalObjectFifoAccessOp
+//===----------------------------------------------------------------------===//
+
+void LogicalObjectFifoAccessOp::build(OpBuilder &b,
+                                      mlir::OperationState &result, Value input,
+                                      MemoryAccess accessType) {
+  auto type = llvm::cast<LogicalObjectFifoType>(input.getType());
+  build(b, result, type.getElementType(), input, accessType);
+}
+
+LogicalObjectFifoFromMemrefOp
+LogicalObjectFifoAccessOp::getLogicalObjectFifo() {
+  return dyn_cast<LogicalObjectFifoFromMemrefOp>(getInput().getDefiningOp());
+};
+
 //===----------------------------------------------------------------------===//
 // AMDAIE_LogicalObjectFifoAcquire
 //===----------------------------------------------------------------------===//
@@ -341,6 +357,26 @@ void LogicalObjectFifoAcquire::build(OpBuilder &b, mlir::OperationState &result,
 // AMDAIE_LogicalObjectFifoFromMemrefOp
 //===----------------------------------------------------------------------===//
 
+/// Build with an array of static tile locations.
+void LogicalObjectFifoFromMemrefOp::build(
+    OpBuilder &b, mlir::OperationState &result, Value memref,
+    ArrayRef<std::pair<int64_t, int64_t>> tileLocations) {
+  SmallVector<Value> tiles;
+  tiles.reserve(tileLocations.size());
+  for (auto [column, row] : tileLocations) {
+    auto colIndex = b.create<arith::ConstantIndexOp>(b.getUnknownLoc(), column);
+    auto rowIndex = b.create<arith::ConstantIndexOp>(b.getUnknownLoc(), row);
+    auto tileOp =
+        b.create<AMDAIE::TileOp>(b.getUnknownLoc(), colIndex, rowIndex);
+    tiles.push_back(tileOp.getResult());
+  }
+  // For deterministic order.
+  llvm::sort(tiles.begin(), tiles.end(),
+             TileOp::tileValueColumnAndRowComparator);
+  auto type = LogicalObjectFifoType::get(cast<MemRefType>(memref.getType()));
+  build(b, result, type, memref, tiles);
+}
+
 LogicalResult LogicalObjectFifoFromMemrefOp::canonicalize(
     LogicalObjectFifoFromMemrefOp logicalObjectFifo,
     PatternRewriter &rewriter) {
@@ -349,23 +385,19 @@ LogicalResult LogicalObjectFifoFromMemrefOp::canonicalize(
     return success();
   }
 
-  auto comparator = [](Value a, Value b) -> bool {
-    TileOp tileA = dyn_cast<TileOp>(a.getDefiningOp());
-    TileOp tileB = dyn_cast<TileOp>(b.getDefiningOp());
-    int64_t colA = getConstantIntValue(tileA.getCol()).value();
-    int64_t rowA = getConstantIntValue(tileA.getRow()).value();
-    int64_t colB = getConstantIntValue(tileB.getCol()).value();
-    int64_t rowB = getConstantIntValue(tileB.getRow()).value();
-    if (colA == colB) return rowA < rowB;
-    return colA < colB;
-  };
   SmallVector<Value> tiles = logicalObjectFifo.getTiles();
-  if (llvm::is_sorted(tiles, comparator)) {
+  if (llvm::is_sorted(tiles, TileOp::tileValueColumnAndRowComparator)) {
+    // Still erase duplicates.
+    tiles.erase(std::unique(tiles.begin(), tiles.end()), tiles.end());
     return success();
   }
 
-  // If tiles are not sorted, sort them and replace the logical objectfifo
-  llvm::sort(tiles.begin(), tiles.end(), comparator);
+  // If tiles are not sorted, sort them, erase duplicates and replace the
+  // logical objectfifo.
+  llvm::sort(tiles.begin(), tiles.end(),
+             TileOp::tileValueColumnAndRowComparator);
+  tiles.erase(std::unique(tiles.begin(), tiles.end()), tiles.end());
+
   rewriter.replaceOpWithNewOp<AMDAIE::LogicalObjectFifoFromMemrefOp>(
       logicalObjectFifo,
       llvm::cast<LogicalObjectFifoType>(
@@ -532,6 +564,23 @@ bool TileOp::hasStaticLocation() {
   return getConstantIntValue(getCol()) && getConstantIntValue(getRow());
 }
 
+bool TileOp::tileColumnComparator(AMDAIE::TileOp &a, AMDAIE::TileOp &b) {
+  int64_t colA = getConstantIntValue(a.getCol()).value();
+  int64_t colB = getConstantIntValue(b.getCol()).value();
+  return colA < colB;
+}
+
+bool TileOp::tileValueColumnAndRowComparator(Value a, Value b) {
+  TileOp tileA = dyn_cast<AMDAIE::TileOp>(a.getDefiningOp());
+  TileOp tileB = dyn_cast<AMDAIE::TileOp>(b.getDefiningOp());
+  int64_t colA = getConstantIntValue(tileA.getCol()).value();
+  int64_t rowA = getConstantIntValue(tileA.getRow()).value();
+  int64_t colB = getConstantIntValue(tileB.getCol()).value();
+  int64_t rowB = getConstantIntValue(tileB.getRow()).value();
+  if (colA == colB) return rowA < rowB;
+  return colA < colB;
+};
+
 //===----------------------------------------------------------------------===//
 // AMDAIE_WorkgroupOp
 //===----------------------------------------------------------------------===//

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

@@ -110,6 +110,11 @@ def AMDAIE_TileOp: AMDAIE_Op<"tile", [
 
   let extraClassDeclaration = [{
     bool hasStaticLocation();
+    // Comparator for `amdaie.tile` based on column index.
+    static bool tileColumnComparator(AMDAIE::TileOp &a, AMDAIE::TileOp &b);
+    // Comparator for `amdaie.tile` values based on column index first and then
+    // row index.
+    static bool tileValueColumnAndRowComparator(Value a, Value b);
   }];
 }
 
@@ -319,6 +324,53 @@ def AMDAIE_NpuDmaWaitOp: AMDAIE_Op<"npu.dma_wait", []> {
 // IREE AMDAIE LogicalObjectFifo Ops
 //===----------------------------------------------------------------------===//
 
+def AMDAIE_LogicalObjectFifoAccessOp : AMDAIE_Op<"logicalobjectfifo.access"> {
+  let summary = "Operation to access the encapsulated memref from a logical"
+                "objectFifo.";
+  let description = [{
+    Returns the encapsulated memref from a logical objectFifo. This is meant to
+    be used within `amdaie.core` operations to access and operate on the memref.
+    Has a memory `access_type` argument that indicates the type of access being
+    done. This can be used to generate a correct (semaphore) synchronization
+    scheme to access the logical objectFifo's content.
+
+    Example:
+    ```mlir
+      %tile = amdaie.tile(%c1, %c3)
+      %alloc = memref.alloc() : memref<8x16xi32, 2>
+      %0 = amdaie.logicalobjectfifo.from_memref %alloc, {%tile} : memref<8x16xi32, 2>
+        -> !amdaie.logicalobjectfifo<memref<8x16xi32, 2>>
+      %core = amdaie.core(%tile) {
+        %1 = amdaie.logicalobjectfifo.access(%0, Read) : 
+          !amdaie.logicalobjectfifo<memref<8x16xi32, 2>> ->  memref<8x16xi32, 2>
+    ```
+  }];
+
+  let arguments = (
+    ins AnyAMDAIELogicalObjectFifoType:$input,
+        MemoryAccess:$access_type
+  ); 
+
+  let results = (outs AnyMemRef:$output);
+
+  let assemblyFormat = [{
+    `(` $input `,` $access_type `)`  attr-dict `:` type($input) `->` type($output)
+  }];
+
+  let builders = [
+    // Build a LogicalObjectFifoAccessOp with a logicalObjectFifo value and access
+    // type.
+    OpBuilder<(ins "mlir::Value":$input, "MemoryAccess":$access_type)>
+  ];
+
+  let extraClassDeclaration = [{
+    LogicalObjectFifoFromMemrefOp getLogicalObjectFifo();
+  }];
+
+  // let hasVerifier = 1;
+  let cppNamespace = "mlir::iree_compiler::AMDAIE";
+}
+
 def AMDAIE_LogicalObjectFifoAcquire: 
     AMDAIE_Op<"logicalobjectfifo.acquire", []> {
   let summary = "Semaphore operation to acquire objects from a logical"
@@ -430,7 +482,12 @@ def AMDAIE_LogicalObjectFifoFromMemrefOp
 
   // Build a LogicalObjectFifoFromMemrefOp with just a memref value.
   let builders = [
-    OpBuilder<(ins "mlir::Value":$memref)>
+    OpBuilder<(ins "mlir::Value":$memref)>,
+    // Build `LogicalObjectFifoFromMemrefOp` with an array of static tile
+    // locations.
+    OpBuilder<
+      (ins "mlir::Value":$memref, 
+           "::llvm::ArrayRef<std::pair<int64_t, int64_t>>":$tileLocations)>
   ];
 
   let extraClassDeclaration = [{

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/test/roundtrip.mlir

@@ -116,6 +116,15 @@ func.func @dma_cpy_nd_mixed(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32
 
 // -----
 
+// CHECK-LABEL: func.func @logicalobjectfifo_access
+// CHECK:       amdaie.logicalobjectfifo.access
+func.func @logicalobjectfifo_access(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 2>>) {
+  %0 = amdaie.logicalobjectfifo.access(%arg0, Write) : !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 2>> -> memref<1x1x8x16xi32, 2 : i32>
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func.func @logicalobjectfifo_acquire
 // CHECK:       %[[DMA:.+]] = amdaie.dma_cpy_nd
 // CHECK:       amdaie.logicalobjectfifo.acquire