TensorFlow TopK integration tests disabled due to failures
antiagainst opened this issue · comments
iree/tools/iree-opt -split-input-file -verify-diagnostics -pass-pipeline="func.func(iree-flow-dispatch-linalg-on-tensors-pass)" crashes on the following input:
func @top_k__2x2__f32__uniform__k_1(%arg0: !hal.buffer_view) -> (!hal.buffer_view, !hal.buffer_view) attributes {iree.abi.stub, iree.reflection = {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,2,2,2]],\22r\22:[[\22stuple\22,[\22ndarray\22,\22f32\22,2,2,1],[\22ndarray\22,\22i32\22,2,2,1]]],\22v\22:1}"}} {
%c0 = arith.constant 0 : index
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%true = arith.constant true
%0 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<2x2xf32>
%1 = linalg.init_tensor [2, 2] : tensor<2x2xi32>
%2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} outs(%1 : tensor<2x2xi32>) {
^bb0(%arg1: i32):
%38 = linalg.index 1 : index
%39 = arith.index_cast %38 : index to i32
linalg.yield %39 : i32
} -> tensor<2x2xi32>
cf.br ^bb1(%c0, %0, %2 : index, tensor<2x2xf32>, tensor<2x2xi32>)
^bb1(%3: index, %4: tensor<2x2xf32>, %5: tensor<2x2xi32>): // 2 preds: ^bb0, ^bb6
%6 = arith.cmpi slt, %3, %c2 : index
cf.cond_br %6, ^bb2(%c0, %4, %5 : index, tensor<2x2xf32>, tensor<2x2xi32>), ^bb7
^bb2(%7: index, %8: tensor<2x2xf32>, %9: tensor<2x2xi32>): // 2 preds: ^bb1, ^bb5
%10 = arith.cmpi slt, %7, %c2 : index
cf.cond_br %10, ^bb3, ^bb6
^bb3: // pred: ^bb2
%11 = tensor.extract %8[%3, %c0] : tensor<2x2xf32>
%12 = tensor.from_elements %11 : tensor<f32>
%13 = tensor.extract %8[%3, %c1] : tensor<2x2xf32>
%14 = tensor.from_elements %13 : tensor<f32>
%15 = tensor.extract %9[%3, %c0] : tensor<2x2xi32>
%16 = tensor.extract %9[%3, %c1] : tensor<2x2xi32>
%17 = linalg.init_tensor [] : tensor<i1>
%18 = linalg.generic {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} ins(%14, %12 : tensor<f32>, tensor<f32>) outs(%17 : tensor<i1>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: i1):
%38 = arith.cmpf ogt, %arg1, %arg2 : f32
linalg.yield %38 : i1
} -> tensor<i1>
%19 = linalg.generic {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} ins(%12, %14 : tensor<f32>, tensor<f32>) outs(%17 : tensor<i1>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: i1):
%38 = arith.cmpf ogt, %arg1, %arg2 : f32
linalg.yield %38 : i1
} -> tensor<i1>
%20 = tensor.extract %19[] : tensor<i1>
%21 = tensor.extract %18[] : tensor<i1>
%22 = arith.xori %20, %true : i1
%23 = arith.andi %22, %21 : i1
cf.cond_br %23, ^bb4, ^bb5(%8, %9 : tensor<2x2xf32>, tensor<2x2xi32>)
^bb4: // pred: ^bb3
%24 = tensor.insert %13 into %8[%3, %c0] : tensor<2x2xf32>
%25 = tensor.insert %11 into %24[%3, %c1] : tensor<2x2xf32>
%26 = tensor.insert %16 into %9[%3, %c0] : tensor<2x2xi32>
%27 = tensor.insert %15 into %26[%3, %c1] : tensor<2x2xi32>
cf.br ^bb5(%25, %27 : tensor<2x2xf32>, tensor<2x2xi32>)
^bb5(%28: tensor<2x2xf32>, %29: tensor<2x2xi32>): // 2 preds: ^bb3, ^bb4
%30 = arith.addi %7, %c1 : index
cf.br ^bb2(%30, %28, %29 : index, tensor<2x2xf32>, tensor<2x2xi32>)
^bb6: // pred: ^bb2
%31 = arith.addi %3, %c1 : index
cf.br ^bb1(%31, %8, %9 : index, tensor<2x2xf32>, tensor<2x2xi32>)
^bb7: // pred: ^bb1
%32 = tensor.extract_slice %4[0, 0] [2, 1] [1, 1] : tensor<2x2xf32> to tensor<2xf32>
%33 = tensor.expand_shape %32 [[0, 1]] : tensor<2xf32> into tensor<2x1xf32>
%34 = tensor.extract_slice %5[0, 0] [2, 1] [1, 1] : tensor<2x2xi32> to tensor<2xi32>
%35 = tensor.expand_shape %34 [[0, 1]] : tensor<2xi32> into tensor<2x1xi32>
%36 = hal.tensor.export %33 : tensor<2x1xf32> -> !hal.buffer_view
%37 = hal.tensor.export %35 : tensor<2x1xi32> -> !hal.buffer_view
return %36, %37 : !hal.buffer_view, !hal.buffer_view
}
The above is generated in TensorFlow TopK integration tests, which are disabled in #8913.
Fun, the PR above fixes this particular repro, but the topk tests still fail. Though 10 other tests are now "unexpectedly passing" :)
@antiagainst can you provide some details about how you obtained the IR (or if you have the IR that is input to IREE). Maybe there is some other issue.
Here is the original mhlo input to IREE:
module {
func @top_k__2x2__f32__uniform__k_1(%arg0: !iree_input.buffer_view) -> (!iree_input.buffer_view, !iree_input.buffer_view) attributes {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,2,2,2]],\22r\22:[[\22stuple\22,[\22ndarray\22,\22f32\22,2,2,1],[\22ndarray\22,\22i32\22,2,2,1]]],\22v\22:1}"} {
%0 = iree_input.cast.buffer_view_to_tensor %arg0 : !iree_input.buffer_view -> tensor<2x2xf32>
%1 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<2xi32>
%2 = "mhlo.broadcast_in_dim"(%1) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xi32>) -> tensor<2x2xi32>
%3:2 = "mhlo.sort"(%0, %2) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<i32>):
%8 = "mhlo.compare"(%arg1, %arg2) {compare_type = #mhlo<"comparison_type TOTALORDER">, comparison_direction = #mhlo<"comparison_direction GT">} : (tensor<f32>, tensor<f32>) -> tensor<i1>
"mhlo.return"(%8) : (tensor<i1>) -> ()
}) {dimension = 1 : i64, is_stable = true} : (tensor<2x2xf32>, tensor<2x2xi32>) -> (tensor<2x2xf32>, tensor<2x2xi32>)
%4 = "mhlo.slice"(%3#0) {limit_indices = dense<[2, 1]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<2x1xf32>
%5 = "mhlo.slice"(%3#1) {limit_indices = dense<[2, 1]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xi32>) -> tensor<2x1xi32>
%6 = iree_input.cast.tensor_to_buffer_view %4 : tensor<2x1xf32> -> !iree_input.buffer_view
%7 = iree_input.cast.tensor_to_buffer_view %5 : tensor<2x1xi32> -> !iree_input.buffer_view
return %6, %7 : !iree_input.buffer_view, !iree_input.buffer_view
}
func @top_k__2x2__f32__uniform__k_2(%arg0: !iree_input.buffer_view) -> (!iree_input.buffer_view, !iree_input.buffer_view) attributes {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,2,2,2]],\22r\22:[[\22stuple\22,[\22ndarray\22,\22f32\22,2,2,2],[\22ndarray\22,\22i32\22,2,2,2]]],\22v\22:1}"} {
%0 = iree_input.cast.buffer_view_to_tensor %arg0 : !iree_input.buffer_view -> tensor<2x2xf32>
%1 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<2xi32>
%2 = "mhlo.broadcast_in_dim"(%1) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xi32>) -> tensor<2x2xi32>
%3:2 = "mhlo.sort"(%0, %2) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<i32>):
%6 = "mhlo.compare"(%arg1, %arg2) {compare_type = #mhlo<"comparison_type TOTALORDER">, comparison_direction = #mhlo<"comparison_direction GT">} : (tensor<f32>, tensor<f32>) -> tensor<i1>
"mhlo.return"(%6) : (tensor<i1>) -> ()
}) {dimension = 1 : i64, is_stable = true} : (tensor<2x2xf32>, tensor<2x2xi32>) -> (tensor<2x2xf32>, tensor<2x2xi32>)
%4 = iree_input.cast.tensor_to_buffer_view %3#0 : tensor<2x2xf32> -> !iree_input.buffer_view
%5 = iree_input.cast.tensor_to_buffer_view %3#1 : tensor<2x2xi32> -> !iree_input.buffer_view
return %4, %5 : !iree_input.buffer_view, !iree_input.buffer_view
}
}
The error is something else now. It is somewhere in TF land.
Diagnostics:
/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py:1096:0: error: 'tf.InTopKV2' op : illegal op still exists
/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py:141:0: note: called from
/tmpfs/src/github/iree/integrations/tensorflow/test/python/iree_tf_tests/math/math_test.py:674:0: note: called from
/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/core/function_wrappers.py:117:0: note: called from
/tmpfs/src/github/iree/integrations/tensorflow/test/python/iree_tf_tests/math/math_test.py:674:0: note: called from
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py:1125:0: note: called from
/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py:672:0: note: called from
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py:1143:0: note: called from
/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py:3402:0: note: called from
/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py:3557:0: note: called from
/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py:1096:0: note: see current operation: %3 = "tf.InTopKV2"(%2, %1, %0) {device = ""} : (tensor<8x3xf32>, tensor<8xi32>, tensor<i32>) -> tensor<8xi1> loc(fused["InTopKV2:", callsite("in_top_k/InTopKV2@__inference_<lambda>_89"("/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py":1096:0) at callsite("/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py":141:0 at callsite("/tmpfs/src/github/iree/integrations/tensorflow/test/python/iree_tf_tests/math/math_test.py":674:0 at callsite("/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/core/function_wrappers.py":117:0 at callsite("/tmpfs/src/github/iree/integrations/tensorflow/test/python/iree_tf_tests/math/math_test.py":674:0 at callsite("/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py":1125:0 at callsite("/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py":672:0 at callsite("/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py":1143:0 at callsite("/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py":3402:0 at "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py":3557:0)))))))))])
<unknown>:0: error: The following illegal operations still remain:
tf.InTopKV2 (count: 1)
Just assigning this to Rob cause it needs to be assigned to someone. Seems like just might be something that might get fixed by itself.
I just noticed that the iree-dialects in integrations/tensorflow is out of date for a while. So landed #8943 to keep them in sync. It does include changes relating to topk. Though not sure it will fix whatever left.
Working for the static cases now so closing. We can fix dynamic topk when it proves useful.