Support AtenNonzeroOp via DecomposeComplexOps

[renxida]

Implements part of: #3293

Uses logic by GPetters from #615

Prototyped on google colab in this notebook

[renxida]

@gpetters94 does some shape-related things in his previous attempt at this op that I don't understand. I think I should maybe replicate those. Specifically:

1. The input is flattened coming into the decomposition: 33f9101#diff-bec402e56f7c88a9b22018c90df7e9ca721bdd097d4ffdd7c13d9d3f4c2e0688R681-R686

2. There is some stacking / transposing that I don't do: 33f9101#diff-bec402e56f7c88a9b22018c90df7e9ca721bdd097d4ffdd7c13d9d3f4c2e0688R764-R769

[lionsky]

Using below testcase, it seems to generate illegal aten.slice.Tensor

func.func @test_nonzero(%1199 : !torch.vtensor<[2,5],si64>) -> !torch.vtensor<[?,2],si64> {
    %1200 = torch.aten.nonzero %1199 : !torch.vtensor<[2,5],si64> -> !torch.vtensor<[?,2],si64>
    return %1200 : !torch.vtensor<[?,2],si64>

}

command line to reproduce

 torch-mlir-opt -pass-pipeline='builtin.module(func.func(torch-decompose-complex-ops),torch-backend-to-linalg-on-tensors-backend-pipeline)' nonzero.mlir

from the debug, one SliceTensor is illegal, since the slice dim is >= input rank.

  * Pattern : 'torch.aten.slice.Tensor -> ()' {
Trying to match "(anonymous namespace)::ConvertAtenSliceTensorOp"
    ** Insert  : 'arith.constant'(0x55651ccbcd90)
    ** Insert  : 'arith.constant'(0x55651ccbeea0)
    ** Insert  : 'arith.constant'(0x55651ccbef10)
    ** Failure : dim is statically invalid
"(anonymous namespace)::ConvertAtenSliceTensorOp" result 0
  } -> FAILURE : pattern failed to match

%20 rank is 1, but the dim is also %int1

%21 = torch.aten.slice.Tensor %20, %int1, %int1, %int2, %int1 : !torch.vtensor<[2],si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[1],si64>

Do you have any suggestion? @renxida

[lionsky]

It seems the second slice tensor should happen at the zero dimension, since the strides tensor is 1-D tensor

    Value slicedStrides = rewriter.create<AtenSliceTensorOp>(
        loc, slicedStrideType, flippedCumulativeProduct, /*dim*/ c(0),
        /*start=*/strideSliceStart, /*end=*/strideSliceEnd, /*step=*/c(1));

[renxida]

on this one today while Stella looks at my shortfin thing.

[renxida]

@lionsky nice catch. i'll fix this one!