shape-legalize-to-stablehlo emit segmentation fault when there is tensor.extract in IR

[qingyunqu]

What happened?

Testcase:

module attributes {torch.debug_module_name = "UnsafeViewCollapseDynamicWithAtenSizeIntModule"} {
  func.func @forward(%arg0: tensor<?x?x?x?x?x?xf32>, %arg1: tensor<i64>, %arg2: tensor<i64>) -> tensor<?x?x?x?x384xf32> {
    %c384_i64 = arith.constant 384 : i64
    %c0 = arith.constant 0 : index
    %c3 = arith.constant 3 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x?x?x?x?x?xf32>
    %0 = arith.index_cast %dim : index to i64
    %extracted = tensor.extract %arg1[] : tensor<i64>
    %extracted_0 = tensor.extract %arg2[] : tensor<i64>
    %dim_1 = tensor.dim %arg0, %c3 : tensor<?x?x?x?x?x?xf32>
    %1 = arith.index_cast %dim_1 : index to i64
    %from_elements = tensor.from_elements %0, %extracted, %extracted_0, %1, %c384_i64 : tensor<5xi64>
    %2 = stablehlo.dynamic_reshape %arg0, %from_elements : (tensor<?x?x?x?x?x?xf32>, tensor<5xi64>) -> tensor<?x?x?x?x384xf32>
    return %2 : tensor<?x?x?x?x384xf32>
  }
}

Run stablehlo-opt --shape-legalize-to-stablehlo would reproduce the error.
I prefer it to emit pass fail but segv.

Steps to reproduce your issue

1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

Version information

No response

[sdasgup3]

Thanks @qingyunqu for pointing poining this issue.

IMO, extracting value from scalar tensors, using empty indices, is not supported yet (from cs).

IMO this is a case of data dependent dynamism which is not supported (ref Dynamism-RFC::Out Of Scope[O6]). At least we should error out instead of just crashing. In any case, will get back to you soon.