update printing of tiles

nod-ai · Jan 13, 2025 · 072f706 · 072f706
1 parent a9955af
commit 072f706
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 9 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
@@ -1367,16 +1367,26 @@ SmallVector<AMDAIE::NpuDmaCpyNdOp> NpuDmaWaitOp::getDmaOps() {
 // AMDAIE_TileOp
 //===----------------------------------------------------------------------===//
 
+// Example: if the column is an integer value (3) and the row is not, the SSA
+// value might be `%tile_3_r`, where the `_r` denotes that the row is not known.
 void TileOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
   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;
+  std::ostringstream name;
+  name << "tile";
+
+  auto add = [&](std::optional<int64_t> maybeValue, char unknown) {
+    if (maybeValue.has_value()) {
+      name << '_' << maybeValue.value();
+    } else {
+      name << '_' << unknown;
+    }
+  };
+  if (iCol.has_value() || iRow.has_value()) {
+    add(iCol, 'c');
+    add(iRow, 'r');
   }
-  setNameFn(getResult(), name);
+  setNameFn(getResult(), name.str());
 }
 
 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
@@ -503,20 +503,28 @@ func.func @reference_to() {
 
 // -----
 
-// Test that if the row and column are statically known, the tile operation is
+// Test that if the row OR column is 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
+
+    // CHECK: %tile_c_3 = amdaie.tile
     %t_i3 = amdaie.tile(%i, %c3)
-    // CHECK: %tile_1 = amdaie.tile
+
+    // CHECK: %tile_2_r  = amdaie.tile
     %t_2i = amdaie.tile(%c2, %i)
+
+    // CHECK: %tile = amdaie.tile
+    %t_uu = amdaie.tile(%i, %i)
+
     amdaie.controlcode {
       amdaie.end
     }