Skip to content

Commit

Permalink
Add tilize/untilize ttir generic region ops (#2274)
Browse files Browse the repository at this point in the history
### Ticket
#1948 

### What's changed
Add `ttir.tile_tilize_block` and `ttir.tile_untilize_block` ops to the
TTIR generic region ops.
Add respective verifiers to match input/output element types.

### Checklist
- [x] New/Existing tests provide coverage for changes
  • Loading branch information
jdesousa-TT authored Feb 25, 2025
1 parent bc092ea commit f15e704
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 0 deletions.
34 changes: 34 additions & 0 deletions include/ttmlir/Dialect/TTIR/IR/TTIRGenericRegionOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -136,4 +136,38 @@ def TTIR_TileMatmulBlockOp : TTIR_GenericRegionOp<"tile_matmul_block",
let hasVerifier = 1;
}

def TTIR_TileTilizeBlockOp : TTIR_GenericRegionOp<"tile_tilize_block",
[DestinationStyleOpInterface, MemoryEffects<[MemRead, MemWrite]>]> {
let summary = "TTIR Tile Tilize Block Op";
let description = [{
The `tile_tilize_block` operation tilizes the input row major memref block and outputs the memref containing the tilized data.
}];

let arguments = (ins AnyNon0RankedMemRef:$input,
AnyNon0RankedMemRef:$output);

let extraClassDeclaration = [{
MutableOperandRange getDpsInitsMutable() { return getOutputMutable(); }
}];

let hasVerifier = 1;
}

def TTIR_TileUntilizeBlockOp : TTIR_GenericRegionOp<"tile_untilize_block",
[DestinationStyleOpInterface, MemoryEffects<[MemRead, MemWrite]>]> {
let summary = "TTIR Tile Untilize Block Op";
let description = [{
The `tile_untilize_block` operation untilizes the input tilized memref block and outputs the memref contianing the row major data.
}];

let arguments = (ins AnyNon0RankedMemRef:$input,
AnyNon0RankedMemRef:$output);

let extraClassDeclaration = [{
MutableOperandRange getDpsInitsMutable() { return getOutputMutable(); }
}];

let hasVerifier = 1;
}

#endif // TTMLIR_TTMLIR_DIALECT_TTIR_TTIRGENERICREGIONOPS_TD
34 changes: 34 additions & 0 deletions lib/Dialect/TTIR/IR/TTIRGenericRegionOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,37 @@ ::mlir::LogicalResult mlir::tt::ttir::TileMatmulBlockOp::verify() {

return success();
}

// TileTilizeBlockOp verification
::mlir::LogicalResult mlir::tt::ttir::TileTilizeBlockOp::verify() {

if (llvm::isa<mlir::tt::TileType>(getInput().getType().getElementType())) {
return emitOpError(
"MemRef operand to TileTilizeBlock must not have tt.tile "
"element type");
}

if (!llvm::isa<mlir::tt::TileType>(getOutput().getType().getElementType())) {
return emitOpError("MemRef result of TileTilizeBlock must have tt.tile "
"element type");
}

return success();
}

// TileUntilizeBlockOp verification
::mlir::LogicalResult mlir::tt::ttir::TileUntilizeBlockOp::verify() {

if (!llvm::isa<mlir::tt::TileType>(getInput().getType().getElementType())) {
return emitOpError("MemRef operand to TileUntilizeBlock must have tt.tile "
"element type");
}

if (llvm::isa<mlir::tt::TileType>(getOutput().getType().getElementType())) {
return emitOpError(
"MemRef result of TileUntilizeBlock must not have tt.tile "
"element type");
}

return success();
}

0 comments on commit f15e704

Please sign in to comment.