benvanik
commented
2 years ago Looks like some stuff around select handling; unrealized_conversion_cast sneaks in and doesn't go away:
%338 = arith.cmpi eq, %337, %c0_i32 : i32
%339 = select %338, %323, %331 : !stream.resource<*>
%340 = select %338, %321, %330 : index
%341 = builtin.unrealized_conversion_cast %339, %340 : !stream.resource<*>, index to !stream.resource<*>Should be pretty easy to fix up.
Note that whatever this model is doing is quite extreme in highlighting our current linalg fusion deficiencies. If we're expecting programs like this we really can't keep kicking the can on that.
%93 = "mhlo.compare"(%arg159, %3) {compare_type = "FLOAT", comparison_direction = "EQ"} : (tensor<1x1x1x3xf32>, tensor<1x1x1x3xf32>) -> tensor<1x1x1x3xi1>
%94 = mhlo.reduce %93, %83 ( {
^bb0(%arg322: tensor<i1>, %arg323: tensor<i1>): // no predecessors
%2359 = mhlo.and %arg322, %arg323 : tensor<i1>
"mhlo.return"(%2359) : (tensor<i1>) -> ()
}) {dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>} : (tensor<1x1x1x3xi1>, tensor<i1>) -> tensor<i1>
%95 = "mhlo.not"(%94) : (tensor<i1>) -> tensor<i1>
%96 = "mhlo.convert"(%95) : (tensor<i1>) -> tensor<i32>
%97 = tensor.extract %96[] : tensor<i32>
%98 = arith.cmpi eq, %97, %c0_i32 : i32
%99 = select %98, %arg321, %92 : tensor<1x224x224x3xf32>->
%5 = flow.dispatch.workgroups[%c3, %c1, %c1](%1) : (tensor<3xf32>) -> tensor<3xi1> =
(%arg1: !flow.dispatch.tensor<readonly:3xf32>, %arg2: !flow.dispatch.tensor<writeonly:3xi1>) {
%cst_1 = arith.constant -1.000000e+00 : f32
%c3_2 = arith.constant 3 : index
%781 = linalg.init_tensor [3] : tensor<3xi1>
%workgroup_size_0 = flow.dispatch.workgroup.size[0] : index
%workgroup_id_0 = flow.dispatch.workgroup.id[0] : index
%workgroup_count_0 = flow.dispatch.workgroup.count[0] : index
%782 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_0, %workgroup_size_0]
%783 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_0, %workgroup_size_0]
scf.for %arg3 = %782 to %c3_2 step %783 {
%784 = affine.min affine_map<(d0, d1) -> (d0, -d1 + 3)>(%workgroup_size_0, %arg3)
%785 = flow.dispatch.tensor.load %arg1, offsets = [%arg3], sizes = [%784], strides = [1] : !flow.dispatch.tensor<readonly:3xf32> -> tensor<?xf32>
%786 = affine.min affine_map<(d0, d1) -> (d0, -d1 + 3)>(%workgroup_size_0, %arg3)
%787 = tensor.extract_slice %781[%arg3] [%786] [1] : tensor<3xi1> to tensor<?xi1>
%788 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%785 : tensor<?xf32>) outs(%787 : tensor<?xi1>) {
^bb0(%arg4: f32, %arg5: i1): // no predecessors
%789 = arith.cmpf oeq, %arg4, %cst_1 : f32
linalg.yield %789 : i1
} -> tensor<?xi1>
flow.dispatch.tensor.store %788, %arg2, offsets = [%arg3], sizes = [%786], strides = [1] : tensor<?xi1> -> !flow.dispatch.tensor<writeonly:3xi1>
}
flow.return
}
%6 = flow.dispatch.workgroups[%c1, %c1, %c1](%5) : (tensor<3xi1>) -> tensor<i32> =
(%arg1: !flow.dispatch.tensor<readonly:3xi1>, %arg2: !flow.dispatch.tensor<writeonly:i32>) {
%true = arith.constant true
%781 = flow.dispatch.tensor.load %arg1, offsets = [], sizes = [], strides = [] : !flow.dispatch.tensor<readonly:3xi1> -> tensor<3xi1>
%782 = linalg.init_tensor [] : tensor<i32>
%783 = linalg.init_tensor [] : tensor<i1>
%784 = linalg.fill(%true, %783) : i1, tensor<i1> -> tensor<i1>
%785 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>], iterator_types = ["reduction"]} ins(%781 : tensor<3xi1>) outs(%784 : tensor<i1>) {
^bb0(%arg3: i1, %arg4: i1): // no predecessors
%787 = arith.andi %arg3, %arg4 : i1
linalg.yield %787 : i1
} -> tensor<i1>
%786 = linalg.generic {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} ins(%785 : tensor<i1>) outs(%782 : tensor<i32>) {
^bb0(%arg3: i1, %arg4: i32): // no predecessors
%787 = arith.xori %arg3, %true : i1
%788 = arith.extui %787 : i1 to i32
linalg.yield %788 : i32
} -> tensor<i32>
flow.dispatch.tensor.store %786, %arg2, offsets = [], sizes = [], strides = [] : tensor<i32> -> !flow.dispatch.tensor<writeonly:i32>
flow.return
}
%7 = tensor.extract %6[] : tensor<i32>
%8 = arith.cmpi eq, %7, %c0_i32 : i32
%9 = select %8, %0, %4 : tensor<1x224x224x3xf32>->
%5 = flow.dispatch @predict_dispatch_1::@predict_dispatch_1[%c3, %c1, %c1](%1) : (tensor<3xf32>) -> tensor<3xi1>
%6 = flow.dispatch @predict_dispatch_2::@predict_dispatch_2[%c1, %c1, %c1](%5) : (tensor<3xi1>) -> tensor<i32>
%7 = flow.tensor.load %6 : tensor<i32>
%8 = arith.cmpi eq, %7, %c0_i32 : i32
%9 = select %8, %0, %4 : tensor<1x224x224x3xf32>There are 54 of these in the model, meaning 54 command buffers and 54 device->host synchronizations. Detensoring won't help here as the result of the select (which tensor to use) is then consumed by another dispatch region - instead we need to be moving this kind of logic into the dispatch regions so that we keep only pure device->device dependencies. Is this kind of thing you're doing in the model yourself or something that comes from JAX?
stellaraccident
Input program: https://gist.github.com/stellaraccident/dc82380c55e68fbd726ce22b430b236d
Error:
Note that this model isn't great -- in a real inference situation, a number of the things inducing branchiness would be constants. However, it is important that this works at this level of generality.