torch-mlir icon indicating copy to clipboard operation
torch-mlir copied to clipboard

[Torch Dialect] emit aten.nonzero op, recompose aten.index.Tensor with aten.nonzero_numpy to aten.masked_select

Open Vremold opened this issue 2 years ago • 0 comments

In detail, this PR includes:

  1. emit aten.nonzero and aten.nonzero_numpy op;
  2. add missing shape infer function for aten.masked_select op
  3. add recomposition rule for aten.index.Tensor op. When it's indices come from aten.nonzero_numpy op and the shapes of associated tensor are exactly same, it's equivalent to a aten.masked_select op. But for now, no backend supports aten.masked_select or aten.nonzero_numpy op, this PR only adds a MLIR test case.

PS: This PR relies on canonicalizer of prim.ListConstruct op which is introduced in PR #2306 .

Vremold avatar Jul 16 '23 15:07 Vremold