Fix the MHLO to LMHLO lowering of 'gather'
The lowering assumes that the 'gather' op attributes are identical in both MHLO and LMHLO. But that's not true; some time ago the MHLO version was changed to pack 4 of its attributes into a struct. By doing the same for the LMHLO version we both fix the lowering for this op and resolve a longstanding TODO. PiperOrigin-RevId: 337943946
This commit is contained in:
parent
204cf7d544
commit
33c450e4cb
|
@ -602,11 +602,8 @@ def LHLO_GatherOp: LHLO_Op<"gather", []>, BASE_HLO_GatherOp {
|
|||
let arguments = (ins
|
||||
Arg<LHLO_Buffer, "", [MemRead]>:$operand,
|
||||
Arg<LHLO_IntBuffer, "", [MemRead]>:$start_indices,
|
||||
I64Attr:$index_vector_dim,
|
||||
I64ElementsAttr:$offset_dims,
|
||||
GatherDimensionNumbers:$dimension_numbers,
|
||||
I64ElementsAttr:$slice_sizes,
|
||||
I64ElementsAttr:$collapsed_slice_dims,
|
||||
I64ElementsAttr:$start_index_map,
|
||||
Arg<LHLO_Buffer, "", [MemWrite]>:$output
|
||||
);
|
||||
}
|
||||
|
|
|
@ -287,6 +287,28 @@ func @imag(%operand: memref<2x2xcomplex<f32>>, %result: memref<2x2xf32>) {
|
|||
|
||||
// -----
|
||||
|
||||
// BOTH-LABEL: func @gather
|
||||
func @gather(%operand: memref<13x7xf32>, %idxs: memref<5xi32>, %result: memref<5x7xf32>) {
|
||||
%tensor_operand = tensor_load %operand : memref<13x7xf32>
|
||||
%tensor_idxs = tensor_load %idxs : memref<5xi32>
|
||||
%tensor_result =
|
||||
"mhlo.gather"(%tensor_operand, %tensor_idxs)
|
||||
{ dimension_numbers =
|
||||
{ collapsed_slice_dims = dense<0> : tensor<1xi64>
|
||||
, index_vector_dim = 1 : i64
|
||||
, offset_dims = dense<1> : tensor<1xi64>
|
||||
, start_index_map = dense<0> : tensor<1xi64> }
|
||||
, indices_are_sorted = false
|
||||
, name = "gather.71"
|
||||
, slice_sizes = dense<[1, 7]> : tensor<2xi64> }
|
||||
: (tensor<13x7xf32>, tensor<5xi32>) -> tensor<5x7xf32>
|
||||
// BOTH: "lmhlo.gather"(%{{.*}}, %{{.*}}, %{{.*}})
|
||||
tensor_store %tensor_result, %result : memref<5x7xf32>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// BOTH-LABEL: func @imag_dyn
|
||||
func @imag_dyn(%operand: memref<?xcomplex<f32>>, %result: memref<?xf32>) {
|
||||
%tensor_operand = tensor_load %operand : memref<?xcomplex<f32>>
|
||||
|
|
Loading…
Reference in New Issue