Crash during canonicalization pass because of collapse followed by expand

[vivekkhandelwal1]

Just run, torch-mlir-opt --canonicalize IR.mlir, with the below IR:

module {
  func.func @main_graph(%arg0: tensor<?x64x1xf32>) -> tensor<?x384xf32> {
  %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2]] : tensor<?x64x1xf32> into tensor<?xf32>
  %c0 = arith.constant 0 : index
  %dim = tensor.dim %collapsed, %c0 : tensor<?xf32>
  %c384= arith.constant 384 : index
  %div = arith.divui %dim, %c384 : index
  %expanded = tensor.expand_shape %collapsed [[0, 1]] output_shape [%div, 384] : tensor<?xf32> into tensor<?x384xf32>
  return %expanded : tensor<?x384xf32>
  }
}

Note: To see the crash, you must have a torch-mlir build with assertions enabled.

[vinayakdsci]

@vivekkhandelwal1 This issue is present with mlir-opt too:

mlir-opt --canonicalize IR.mlir

produces the crash

mlir-opt: llvm-project/llvm/include/llvm/ADT/ArrayRef.h:257: const T &llvm::ArrayRef<long>::operator[](size_t) const [T = long]: Assertion `Index < Length && "Invalid index!"' failed.

Is it possible that the issue does not lie in Torch-MLIR?

[vinayakdsci]

@vivekkhandelwal1 The issue was indeed in MLIR upstream, Addressed by llvm/llvm-project#112774.

[vinayakdsci]

@pdhirajkumarprasad this issue should now be gone, after llvm/torch-mlir@f5d15ab