Failure to compile broadcast_in_dim

[jpienaar]

The following fails to compile

module @jit_encode_batch {
  func.func public @main(%arg0: tensor<i32>, %arg293: tensor<?x1024xi32>) -> tensor<?x1024xi32> {
    %0 = mhlo.constant dense<0> : tensor<i32>
    %1 = mhlo.convert %arg0 : tensor<i32>
    %2 = mhlo.reshape %1 : (tensor<i32>) -> tensor<1xi32>
    %3 = mhlo.constant dense<1024> : tensor<1xi32>
    %4 = "mhlo.concatenate"(%2, %3) {dimension = 0 : i64} : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %5 = "mhlo.dynamic_broadcast_in_dim"(%arg293, %4) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<?x1024xi32>, tensor<2xi32>) -> tensor<?x1024xi32>
    return %5 : tensor<?x1024xi32>
 }
}

while the following does

module @jit_encode_batch {
  func.func public @main(%arg0: tensor<i32>, %arg293: tensor<?x1024xi32>) -> tensor<?x1024xi32> {
    %0 = mhlo.constant dense<0> : tensor<i32>
    %1 = mhlo.convert %arg0 : tensor<i32>
    %2 = mhlo.reshape %1 : (tensor<i32>) -> tensor<1xi32>
    %3 = mhlo.constant dense<1024> : tensor<1xi32>
    %4 = "mhlo.concatenate"(%2, %3) {dimension = 0 : i64} : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %5 = "mhlo.dynamic_broadcast_in_dim"(%arg293, %4) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>,
      known_expanding_dimensions = dense<[0]> : tensor<1xi64>,
      known_nonexpanding_dimensions = dense<[1]> : tensor<1xi64>
    } : (tensor<?x1024xi32>, tensor<2xi32>) -> tensor<?x1024xi32>
    return %5 : tensor<?x1024xi32>
 }
}

Not sure if this is a frontend issue (the first form shouldn't have been produced with those attributes), luck that it compiles with those attributes, or expected that it compiles without those.

[jpienaar]

@d0k @sherhut if you are able to direct

[d0k]

Fails to compile using what pipeline? The first one is a lot harder to codegen because you don't know if the first dimension expands or not, we also can't derive this from the IR alone.

As far as I know kernelgen can compile the unannotated version into slow code, while IREE doesn't support this fully dynamic case. Ideally the frontend would produce the attributes, but it might not be possible depending on the input.

[jpienaar]

SG, so if one wants any performance here, you should be using the attributed form?

[jpienaar]

Ah, and while we have this question: what is the tool/config for the different pipelines here? (e.g., is it xla-opt with pipeline foo that captures all the passes)