Add custom printer for amdaie.tile op (#977)

Copied from aie.tile. Makes it easier to read IR.
nod-ai · Dec 11, 2024 · 41ad785 · 41ad785
1 parent e623a6a
commit 41ad785
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 3 deletions.
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.cpp
@@ -1346,7 +1346,15 @@ SmallVector<AMDAIE::NpuDmaCpyNdOp> NpuDmaWaitOp::getDmaOps() {
 //===----------------------------------------------------------------------===//
 
 void TileOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
-  setNameFn(getResult(), "tile");
+  std::optional<int64_t> iCol = getConstantIntValue(getCol());
+  std::optional<int64_t> iRow = getConstantIntValue(getRow());
+  std::string name{"tile"};
+  if (iCol.has_value() && iRow.has_value()) {
+    std::string sCol = std::to_string(iCol.value());
+    std::string sRow = std::to_string(iRow.value());
+    name += "_" + sCol + "_" + sRow;
+  }
+  setNameFn(getResult(), name);
 }
 
 bool TileOp::hasStaticLocation() {

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/test/roundtrip.mlir b/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/test/roundtrip.mlir
@@ -49,7 +49,7 @@ func.func @core() {
 // -----
 
 // CHECK-LABEL: func.func @logicalobjectfifo_from_memref
-// CHECK: %[[I0:.*]] = amdaie.logicalobjectfifo.from_memref %[[ARG0:.*]], {} 
+// CHECK: %[[I0:.*]] = amdaie.logicalobjectfifo.from_memref %[[ARG0:.*]], {}
 // CHECK-SAME: memref<1x1x8x16xi32, 1> -> !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>
 func.func @logicalobjectfifo_from_memref(%arg0: memref<1x1x8x16xi32, 1>) {
   %0 = amdaie.logicalobjectfifo.from_memref %arg0, {} : memref<1x1x8x16xi32, 1> -> !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>
@@ -487,3 +487,26 @@ func.func @reference_to() {
   %1 = amdaie.reference_to %0 : memref<1x1x8x4x8x4xi32>
   return
 }
+
+// -----
+
+// Test that if the row and column are statically known, the tile operation is
+// printed with the row and column in the SSA value.
+func.func @tile_a_b(%i : index) {
+  %c2 = arith.constant 2: index
+  %c3 = arith.constant 3 : index
+  amdaie.workgroup {
+    // CHECK: %tile_2_3 = amdaie.tile
+    %t_23 = amdaie.tile(%c2, %c3)
+    // CHECK: %tile_2_3_0 = amdaie.tile
+    %t_231 = amdaie.tile(%c2, %c3)
+    // CHECK: %tile = amdaie.tile
+    %t_i3 = amdaie.tile(%i, %c3)
+    // CHECK: %tile_1 = amdaie.tile
+    %t_2i = amdaie.tile(%c2, %i)
+    amdaie.controlcode {
+      amdaie.end
+    }
+  }
+  return
+}
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignTiles.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignTiles.cpp
@@ -7,7 +7,6 @@
 #include "iree-amd-aie/IR/AMDAIEOps.h"
 #include "iree-amd-aie/Transforms/Passes.h"
 #include "iree-amd-aie/Transforms/Utils/AMDAIEUtils.h"
-#include "iree-amd-aie/aie_runtime/Utils/ChannelGenerator.h"
 #include "iree-amd-aie/aie_runtime/iree_aie_runtime.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"