Restrict GetDimensionSize HLO op result type to 32 bit integer

XLA implementation has this limitation and always uses 32 bit result for this instruction. This will cause mismatch between the result type in MLIR and XLA at the time of export.

This should be resolved once we have a special dialect mapping directly to HLOInstructionProto. Another option until then could be to introduce a pass to legalize mhlo itself to match XLA semantics.

PiperOrigin-RevId: 324286936
This commit is contained in:
Smit Hinsu 2020-07-31 14:37:20 -07:00 committed by TensorFlow MLIR Team
parent fcb91fb0b9
commit 1c535f1718
1 changed files with 4 additions and 1 deletions

View File

@ -1075,7 +1075,10 @@ def HLO_GetDimensionSizeOp: HLO_Op<"get_dimension_size", [NoSideEffect]>,
HLO_Tensor:$operand,
I32Attr:$dimension
);
let results = (outs HLO_IntTensor);
// TODO(hinsu): Allow 64-bit result types once XLA HLO dialect based on the
// XLA semantics is available. This limitation is because of the current XLA
// implementation.
let results = (outs I32Tensor);
}
def HLO_MapOp: HLO_Op<"map",