Skip to content

Commit

Permalink
[IR] Pretty-printer for LogicalObjectFifoFromMemrefOp (#1017)
Browse files Browse the repository at this point in the history
Makes IR easier to mentally parse -- it should be easier to figure out
which columns/rows a ` LogicalObjectFifoFromMemrefOp` is on.


```
%lof_3_2 = ... // column 3, row 2
%lof_3_r = ... // column 3, multiple rows or unkown row 
```

The lit tests have more examples
  • Loading branch information
newling authored Jan 13, 2025
1 parent a9955af commit 6d4d3ea
Show file tree
Hide file tree
Showing 3 changed files with 157 additions and 2 deletions.
52 changes: 52 additions & 0 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -625,6 +625,58 @@ LogicalObjectFifoFromBuffersOp::replaceWithNewTiles(
// AMDAIE_LogicalObjectFifoFromMemrefOp
//===----------------------------------------------------------------------===//

void LogicalObjectFifoFromMemrefOp::getAsmResultNames(
function_ref<void(Value, StringRef)> setNameFn) {
// 'lof' for 'logical object fifo'
constexpr const char *const name = "lof";

auto tiles = getTiles();

if (tiles.empty()) {
setNameFn(getResult(), name);
return;
}

// Denotes one or more tiles with multiple possible row/column values.
constexpr int64_t multiple{-2};
constexpr int64_t unset{-1};
int64_t col{unset};
int64_t row{unset};

for (Value index : tiles) {
TileOp tile = dyn_cast<TileOp>(index.getDefiningOp());
if (!tile) {
col = multiple;
row = multiple;
} else {
std::optional<int64_t> maybeCol = getConstantIntValue(tile.getCol());
if (!maybeCol) col = multiple;
if (col >= 0 && maybeCol.value() != col) col = multiple;
if (col == unset) col = maybeCol.value();

std::optional<int64_t> maybeRow = getConstantIntValue(tile.getRow());
if (!maybeRow) row = multiple;
if (row >= 0 && maybeRow.value() != row) row = multiple;
if (row == unset) row = maybeRow.value();
}
}

std::ostringstream namestream;
namestream << name << '_';

if (col >= 0) {
namestream << col;
} else {
namestream << 'c';
}
if (row >= 0) {
namestream << '_' << row;
} else {
namestream << '_' << 'r';
}
setNameFn(getResult(), namestream.str());
}

/// Build with an array of static tile locations.
void LogicalObjectFifoFromMemrefOp::build(
OpBuilder &b, mlir::OperationState &result, Value memref,
Expand Down
5 changes: 3 additions & 2 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -1279,8 +1279,9 @@ def AMDAIE_LogicalObjectFifoFromBuffersOp
}

def AMDAIE_LogicalObjectFifoFromMemrefOp
: AMDAIE_Op<"logicalobjectfifo.from_memref",
[DeclareOpInterfaceMethods<LogicalObjFifoOpInterface, ["replaceWithNewTiles"]>,
: AMDAIE_Op<"logicalobjectfifo.from_memref", [
DeclareOpInterfaceMethods<LogicalObjFifoOpInterface, ["replaceWithNewTiles"]>,
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
Pure]> {
let summary = "Create a logical objectFifo from a memref";
let description = [{
Expand Down
102 changes: 102 additions & 0 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/test/roundtrip.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -523,3 +523,105 @@ func.func @tile_a_b(%i : index) {
}
return
}

// -----

// CHECK-LABEL: func.func @from_memref_known_tiles
func.func @from_memref_known_tiles(%arg0 : memref<8xi32>, %t0 : index) {
%c2 = arith.constant 2: index
%c3 = arith.constant 3 : index
amdaie.workgroup {
%tile_2_3 = amdaie.tile(%c2, %c3)
%tile_3_3 = amdaie.tile(%c3, %c3)
%tile_3_2 = amdaie.tile(%c3, %c2)
%tile_2_2 = amdaie.tile(%c2, %c2)
// logicalobjectfifo without any tiles:
// CHECK: %lof =
%fifo0 = amdaie.logicalobjectfifo.from_memref %arg0, {} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
// logicalobjectfifo with one known tile:
// CHECK: %lof_2_3 =
%fifo3 = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_2_3} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
// logicalobjectfifo with two known tiles, in the same column.
// 'r' in the SSA value denotes multiple rows.
// CHECK: %lof_2_r =
%fifo4 = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_2_3, %tile_2_2} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
// logicalobjectfifo with two known tiles, in the same row.
// 'c' in the SSA value denotes multiple columns.
// CHECK: %lof_c_3 =
%fifo5 = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_2_3, %tile_3_3} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
// logicalobjectfifo with two known tiles, in different rows and columns:
// CHECK: %lof_c_r =
%fifo6 = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_2_3, %tile_3_2} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
// logicalobjectfifo with 4 tiles, spanning 2 rows and 2 columns:
// CHECK: %lof_c_r_0 =
%fifo7 = amdaie.logicalobjectfifo.from_memref %arg0,
{%tile_2_3, %tile_3_3, %tile_3_2, %tile_2_2} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
amdaie.controlcode {
amdaie.end
}
}
return
}

// -----

// CHECK-LABEL: func.func @from_memref_unknown_row
func.func @from_memref_unknown_row(%arg0 : memref<8xi32>, %t0 : index) {
%c2 = arith.constant 2: index
amdaie.workgroup {
%tile_2_u = amdaie.tile(%c2, %t0)
// CHECK: %lof_2_r =
%fifo = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_2_u} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
amdaie.controlcode {
amdaie.end
}
}
return
}

// -----

// CHECK-LABEL: func.func @from_memref_unknown_column
func.func @from_memref_unknown_column(%arg0 : memref<8xi32>, %t0 : index) {
%c3 = arith.constant 3 : index
amdaie.workgroup {
%tile_u_3 = amdaie.tile(%t0, %c3)
// CHECK: %lof_c_3 =
%fifo = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_u_3} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
amdaie.controlcode {
amdaie.end
}
}
return
}

// -----

// CHECK-LABEL: func.func @from_memref_unknown_row_column
func.func @from_memref_unknown_row_column(%arg0 : memref<8xi32>, %t0 : index) {
amdaie.workgroup {
%c2 = arith.constant 2: index
%tile_2_2 = amdaie.tile(%c2, %c2)
%tile_u_u = amdaie.tile(%t0, %t0)
// logicalobjectfifo with a single tile with unknown row and column:
// CHECK: %lof_c_r =
%fifo1 = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_u_u} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
// logicalobjectfifo with one unknown tile, and one known tile:
// CHECK: %lof_c_r_0 =
%fifo2 = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_2_2, %tile_u_u} :
memref<8xi32> -> !amdaie.logicalobjectfifo<memref<8xi32>>
amdaie.controlcode {
amdaie.end
}
}
return
}

0 comments on commit 6d4d3ea

Please sign in to comment.