mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 12:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[TPU][Mosaic][Easy] Add verification for AssumeMultipleOp.

Browse Source 
A user must use AssumeMultipleOp to annotate integer constants that are divisible by the given multiple.

PiperOrigin-RevId: 727699186
This commit is contained in:
jax authors 2025-02-16 21:15:24 -08:00
parent 936eb84422
commit a6fcb7415f
2 changed files with 22 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jaxlib/mosaic/dialect/tpu/tpu.td
View File

@ -608,6 +608,7 @@ def TPU_AssumeMultipleOp : TPU_Op<"assume_multiple", [Pure, SameOperandsAndResul
I32Attr:$multiple
);
let results = (outs AnyTypeOf<[Index, AnyInteger]>:$result);
let hasVerifier = 1;
}
def TPU_MemRefSliceOp : TPU_Op<"memref_slice", [Pure, AttrSizedOperandSegments]> {

21
jaxlib/mosaic/dialect/tpu/tpu_ops.cc
View File

@ -1225,6 +1225,27 @@ LogicalResult DynamicGatherOp::verify() {
return success();
}
LogicalResult AssumeMultipleOp::verify() {
auto operand_value = getValue();
auto divisor = getMultiple();
if (auto cst_op = operand_value.getDefiningOp<arith::ConstantOp>()) {
auto int_attr = dyn_cast<IntegerAttr>(cst_op.getValue());
// Illegal usage of AssumeMultipleOp.
if (!int_attr) {
return emitOpError(
"Illegal user annotation, expected an integer, but got ")
<< cst_op.getValue();
}
if (int_attr.getInt() % divisor != 0) {
return emitOpError(
"Illegal user annotation, expected an integer that is "
"divisible by the multiple, but got ")
<< int_attr.getInt() << " % " << divisor;
}
}
return success();
}
} // namespace tpu
} // namespace mlir