colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add layout to mhlo::InfeedOp td. 

PiperOrigin-RevId: 356286875
This commit is contained in:
A. Unique TensorFlower 2021-02-08 09:45:40 -08:00 committed by TensorFlow MLIR Team
parent 80d753c1fe
commit 4a29ca3b1d
2 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
View File

@ -401,12 +401,18 @@ def HLO_InfeedOp : HLO_Op<"infeed", []> {
of the data. Multiple Infeed operations are allowed in a computation, but
there must be a total order among the Infeed operations.
Attributes:
layout: Array attribute. Same shape as the output of the infeed, except
that every tensor is replaced by a minor_to_major array for the
tensor's layout.
See https://www.tensorflow.org/xla/operation_semantics#infeed.
}];
let arguments = (ins
HLO_Token:$token,
DefaultValuedAttr<StrAttr, "">:$infeed_config
DefaultValuedAttr<StrAttr, "">:$infeed_config,
OptionalAttr<ArrayAttr>:$layout
);
let results = (outs HLO_Tuple);
let hasCustomHLOConverter = 1;

4
tests/ops.mlir
View File

@ -442,7 +442,7 @@ func @dot_bad_precision_config(%arg0: tensor<2x2xi32>, %arg1: tensor<2x2xi32>) -
func @infeed_invalid_number_of_results(%token: !mhlo.token) -> tuple<tuple<tensor<i32>>, !mhlo.token, tensor<i32>> {
// expected-error@+1 {{result is expected to be a tuple of size 2, but got 3}}
%0 = "mhlo.infeed"(%token) {infeed_config = "foobar"} : (!mhlo.token) -> tuple<tuple<tensor<i32>>, !mhlo.token, tensor<i32>>
%0 = "mhlo.infeed"(%token) {infeed_config = "foobar", layout = [[[0]], unit, [0]]} : (!mhlo.token) -> tuple<tuple<tensor<i32>>, !mhlo.token, tensor<i32>>
return %0 : tuple<tuple<tensor<i32>>, !mhlo.token, tensor<i32>>
}
@ -450,7 +450,7 @@ func @infeed_invalid_number_of_results(%token: !mhlo.token) -> tuple<tuple<tenso
func @infeed_non_token_second_result(%token: !mhlo.token) -> tuple<tuple<tensor<i32>>, tensor<i32>> {
// expected-error@+1 {{second element of result tuple is expected to be of token type, but got 'tensor<i32>'}}
%0 = "mhlo.infeed"(%token) {infeed_config = "foobar"} : (!mhlo.token) -> tuple<tuple<tensor<i32>>, tensor<i32>>
%0 = "mhlo.infeed"(%token) {infeed_config = "foobar", layout = [[[0]], [0]]} : (!mhlo.token) -> tuple<tuple<tensor<i32>>, tensor<i32>>
return %0 : tuple<tuple<tensor<i32>>, tensor<i32>>
}