nonzero.mlir

./build_tools/ci/test_posix.sh                                                      
::group::Run ONNX e2e integration tests
TORCH_VERSION_FOR_COMPARISON = 2.6.0.dev20241107
Running tests sequentially with progress status
*** RUNNING TEST: AtenNonzero1DModule_one_nonzero ***
Compiling AtenNonzero1DModule_one_nonzero...

====================
ONNX RAW IR
module {
  func.func @main_graph(%arg0: !torch.vtensor<[?],i1>) -> !torch.vtensor<[?,1],si64> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 20 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.6.0"} {
    %none = torch.constant.none
    %0 = torch.operator "onnx.NonZero"(%arg0) : (!torch.vtensor<[?],i1>) -> !torch.vtensor<[1,?],si64> 
    %1 = torch.operator "onnx.Transpose"(%0) {torch.onnx.perm = [1 : si64, 0 : si64]} : (!torch.vtensor<[1,?],si64>) -> !torch.vtensor<[?,1],si64> 
    return %1 : !torch.vtensor<[?,1],si64>
  }
}


====================
Torch IR
module {
  func.func @main_graph(%arg0: !torch.vtensor<[?],i1>) -> !torch.vtensor<[1,1],si64> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 20 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.6.0"} {
    %0 = torch.vtensor.literal(dense<1> : tensor<1x1xsi64>) : !torch.vtensor<[1,1],si64>
    %str = torch.constant.str "floor"
    %str_0 = torch.constant.str "sum"
    %1 = torch.vtensor.literal(dense<0> : tensor<si64>) : !torch.vtensor<[],si64>
    %int-1 = torch.constant.int -1
    %none = torch.constant.none
    %false = torch.constant.bool false
    %int1 = torch.constant.int 1
    %int0 = torch.constant.int 0
    %int4 = torch.constant.int 4
    %2 = torch.aten._shape_as_tensor %arg0 : !torch.vtensor<[?],i1> -> !torch.vtensor<[1],si64>
    %3 = torch.aten.flatten.using_ints %arg0, %int0, %int0 : !torch.vtensor<[?],i1>, !torch.int, !torch.int -> !torch.vtensor<[?],i1>
    %4 = torch.aten.ne.Scalar %3, %int0 : !torch.vtensor<[?],i1>, !torch.int -> !torch.vtensor<[?],i1>
    %5 = torch.aten.to.dtype %4, %int4, %false, %false, %none : !torch.vtensor<[?],i1>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?],si64>
    %6 = torch.aten.cumsum %5, %int0, %none : !torch.vtensor<[?],si64>, !torch.int, !torch.none -> !torch.vtensor<[?],si64>
    %7 = torch.aten.sub.Scalar %6, %int1, %int1 : !torch.vtensor<[?],si64>, !torch.int, !torch.int -> !torch.vtensor<[?],si64>
    %8 = torch.aten.clamp %7, %int0, %none : !torch.vtensor<[?],si64>, !torch.int, !torch.none -> !torch.vtensor<[?],si64>
    %9 = torch.aten.arange.start_step %int0, %int-1, %int1, %none, %none, %none, %none : !torch.int, !torch.int, !torch.int, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?],si64>
    %10 = torch.aten.mul.Tensor %9, %5 : !torch.vtensor<[?],si64>, !torch.vtensor<[?],si64> -> !torch.vtensor<[?],si64>
    %11 = torch.aten.size.int %6, %int0 : !torch.vtensor<[?],si64>, !torch.int -> !torch.int
    %12 = torch.prim.ListConstruct %11 : (!torch.int) -> !torch.list<int>
    %13 = torch.aten.broadcast_to %1, %12 : !torch.vtensor<[],si64>, !torch.list<int> -> !torch.vtensor<[?],si64>
    %14 = torch.aten.scatter_reduce.two %13, %int0, %8, %10, %str_0, %false : !torch.vtensor<[?],si64>, !torch.int, !torch.vtensor<[?],si64>, !torch.vtensor<[?],si64>, !torch.str, !torch.bool -> !torch.vtensor<[?],si64>
    %15 = torch.aten.sum %5, %none : !torch.vtensor<[?],si64>, !torch.none -> !torch.vtensor<[],si64>
    %16 = torch.aten.Int.Tensor %15 : !torch.vtensor<[],si64> -> !torch.int
    %17 = torch.aten.slice.Tensor %14, %int0, %int0, %16, %int1 : !torch.vtensor<[?],si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[?],si64>
    %18 = torch.prim.ListConstruct %int0 : (!torch.int) -> !torch.list<int>
    %19 = torch.aten.flip %2, %18 : !torch.vtensor<[1],si64>, !torch.list<int> -> !torch.vtensor<[1],si64>
    %20 = torch.aten.cumprod %19, %int0, %none : !torch.vtensor<[1],si64>, !torch.int, !torch.none -> !torch.vtensor<[1],si64>
    %21 = torch.aten.flip %20, %18 : !torch.vtensor<[1],si64>, !torch.list<int> -> !torch.vtensor<[1],si64>
    %22 = torch.aten.slice.Tensor %21, %int0, %int1, %int1, %int1 : !torch.vtensor<[1],si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[0],si64>
    %23 = torch.aten.unsqueeze %17, %int1 : !torch.vtensor<[?],si64>, !torch.int -> !torch.vtensor<[?,1],si64>
    %24 = torch.aten.div.Tensor_mode %23, %0, %str : !torch.vtensor<[?,1],si64>, !torch.vtensor<[1,1],si64>, !torch.str -> !torch.vtensor<[?,1],si64>
    %25 = torch.aten.remainder.Tensor %24, %2 : !torch.vtensor<[?,1],si64>, !torch.vtensor<[1],si64> -> !torch.vtensor<[1,1],si64>
    %26 = torch.aten.transpose.int %25, %int0, %int1 : !torch.vtensor<[1,1],si64>, !torch.int, !torch.int -> !torch.vtensor<[1,1],si64>
    return %26 : !torch.vtensor<[1,1],si64>
  }
}


====================
Torch Backend IR
module {
  func.func @main_graph(%arg0: !torch.vtensor<[?],i1>) -> !torch.vtensor<[1,1],si64> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 20 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.6.0"} {
    %0 = torch.vtensor.literal(dense<1> : tensor<1x1xsi64>) : !torch.vtensor<[1,1],si64>
    %str = torch.constant.str "floor"
    %str_0 = torch.constant.str "sum"
    %1 = torch.vtensor.literal(dense<0> : tensor<si64>) : !torch.vtensor<[],si64>
    %int-1 = torch.constant.int -1
    %none = torch.constant.none
    %false = torch.constant.bool false
    %int1 = torch.constant.int 1
    %int0 = torch.constant.int 0
    %int4 = torch.constant.int 4
    %2 = torch.aten._shape_as_tensor %arg0 : !torch.vtensor<[?],i1> -> !torch.vtensor<[1],si64>
    %3 = torch.aten.flatten.using_ints %arg0, %int0, %int0 : !torch.vtensor<[?],i1>, !torch.int, !torch.int -> !torch.vtensor<[?],i1>
    %4 = torch.aten.ne.Scalar %3, %int0 : !torch.vtensor<[?],i1>, !torch.int -> !torch.vtensor<[?],i1>
    %5 = torch.aten.to.dtype %4, %int4, %false, %false, %none : !torch.vtensor<[?],i1>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?],si64>
    %6 = torch.aten.cumsum %5, %int0, %none : !torch.vtensor<[?],si64>, !torch.int, !torch.none -> !torch.vtensor<[?],si64>
    %7 = torch.aten.sub.Scalar %6, %int1, %int1 : !torch.vtensor<[?],si64>, !torch.int, !torch.int -> !torch.vtensor<[?],si64>
    %8 = torch.aten.clamp %7, %int0, %none : !torch.vtensor<[?],si64>, !torch.int, !torch.none -> !torch.vtensor<[?],si64>
    %9 = torch.aten.arange.start_step %int0, %int-1, %int1, %none, %none, %none, %none : !torch.int, !torch.int, !torch.int, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?],si64>
    %10 = torch.aten.mul.Tensor %9, %5 : !torch.vtensor<[?],si64>, !torch.vtensor<[?],si64> -> !torch.vtensor<[?],si64>
    %11 = torch.aten.size.int %6, %int0 : !torch.vtensor<[?],si64>, !torch.int -> !torch.int
    %12 = torch.prim.ListConstruct %11 : (!torch.int) -> !torch.list<int>
    %13 = torch.aten.broadcast_to %1, %12 : !torch.vtensor<[],si64>, !torch.list<int> -> !torch.vtensor<[?],si64>
    %14 = torch.aten.scatter_reduce.two %13, %int0, %8, %10, %str_0, %false : !torch.vtensor<[?],si64>, !torch.int, !torch.vtensor<[?],si64>, !torch.vtensor<[?],si64>, !torch.str, !torch.bool -> !torch.vtensor<[?],si64>
    %15 = torch.aten.sum %5, %none : !torch.vtensor<[?],si64>, !torch.none -> !torch.vtensor<[],si64>
    %16 = torch.aten.Int.Tensor %15 : !torch.vtensor<[],si64> -> !torch.int
    %17 = torch.aten.slice.Tensor %14, %int0, %int0, %16, %int1 : !torch.vtensor<[?],si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[?],si64>
    %18 = torch.prim.ListConstruct %int0 : (!torch.int) -> !torch.list<int>
    %19 = torch.aten.flip %2, %18 : !torch.vtensor<[1],si64>, !torch.list<int> -> !torch.vtensor<[1],si64>
    %20 = torch.aten.cumprod %19, %int0, %none : !torch.vtensor<[1],si64>, !torch.int, !torch.none -> !torch.vtensor<[1],si64>
    %21 = torch.aten.flip %20, %18 : !torch.vtensor<[1],si64>, !torch.list<int> -> !torch.vtensor<[1],si64>
    %22 = torch.aten.slice.Tensor %21, %int0, %int1, %int1, %int1 : !torch.vtensor<[1],si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[0],si64>
    %23 = torch.aten.unsqueeze %17, %int1 : !torch.vtensor<[?],si64>, !torch.int -> !torch.vtensor<[?,1],si64>
    %24 = torch.aten.div.Tensor_mode %23, %0, %str : !torch.vtensor<[?,1],si64>, !torch.vtensor<[1,1],si64>, !torch.str -> !torch.vtensor<[?,1],si64>
    %25 = torch.aten.remainder.Tensor %24, %2 : !torch.vtensor<[?,1],si64>, !torch.vtensor<[1],si64> -> !torch.vtensor<[1,1],si64>
    %26 = torch.aten.transpose.int %25, %int0, %int1 : !torch.vtensor<[1,1],si64>, !torch.int, !torch.int -> !torch.vtensor<[1,1],si64>
    return %26 : !torch.vtensor<[1,1],si64>
  }
}


====================
LINALG Backend IR
#map = affine_map<(d0) -> (d0)>
#map1 = affine_map<(d0) -> ()>
#map2 = affine_map<(d0) -> (d0, 0)>
#map3 = affine_map<(d0, d1) -> (d0, 0)>
#map4 = affine_map<(d0, d1) -> (0, 0)>
#map5 = affine_map<(d0, d1) -> (d0, d1)>
#map6 = affine_map<(d0, d1) -> (0)>
module {
  ml_program.global private mutable @global_seed(dense<0> : tensor<i64>) : tensor<i64>
  func.func @main_graph(%arg0: tensor<?xi1>) -> tensor<1x1xi64> {
    %c-1 = arith.constant -1 : index
    %cst = arith.constant dense<0> : tensor<i64>
    %cst_0 = arith.constant dense<1> : tensor<1x1xi64>
    %c0 = arith.constant 0 : index
    %c0_i64 = arith.constant 0 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %dim = tensor.dim %arg0, %c0 : tensor<?xi1>
    %0 = arith.index_cast %dim : index to i64
    %from_elements = tensor.from_elements %0 : tensor<1xi64>
    %1 = tensor.empty(%dim) : tensor<?xi64>
    %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<?xi1>) outs(%1 : tensor<?xi64>) {
    ^bb0(%in: i1, %out: i64):
      %47 = arith.extui %in : i1 to i64
      linalg.yield %47 : i64
    } -> tensor<?xi64>
    %dim_1 = tensor.dim %2, %c0 : tensor<?xi64>
    %3 = tensor.empty(%dim_1) : tensor<?xi64>
    %4 = linalg.fill ins(%c0_i64 : i64) outs(%3 : tensor<?xi64>) -> tensor<?xi64>
    %5 = tensor.empty() : tensor<i64>
    %6 = linalg.fill ins(%c0_i64 : i64) outs(%5 : tensor<i64>) -> tensor<i64>
    %7:2 = tm_tensor.scan dimension(0) inclusive(true) ins(%2 : tensor<?xi64>) outs(%4, %6 : tensor<?xi64>, tensor<i64>) {
    ^bb0(%arg1: i64, %arg2: i64):
      %47 = arith.addi %arg1, %arg2 : i64
      tm_tensor.yield %47 : i64
    } -> tensor<?xi64>, tensor<i64>
    %dim_2 = tensor.dim %7#0, %c0 : tensor<?xi64>
    %8 = tensor.empty(%dim_2) : tensor<?xi64>
    %9 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%7#0 : tensor<?xi64>) outs(%8 : tensor<?xi64>) {
    ^bb0(%in: i64, %out: i64):
      %47 = arith.subi %in, %c1_i64 : i64
      linalg.yield %47 : i64
    } -> tensor<?xi64>
    %10 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%9 : tensor<?xi64>) outs(%8 : tensor<?xi64>) {
    ^bb0(%in: i64, %out: i64):
      %47 = arith.cmpi slt, %in, %c0_i64 : i64
      %48 = arith.select %47, %c0_i64, %in : i64
      linalg.yield %48 : i64
    } -> tensor<?xi64>
    %11 = tensor.empty(%c-1) : tensor<?xi64>
    %12 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%11 : tensor<?xi64>) {
    ^bb0(%out: i64):
      %47 = linalg.index 0 : index
      %48 = arith.index_cast %47 : index to i64
      linalg.yield %48 : i64
    } -> tensor<?xi64>
    %13 = arith.cmpi eq, %dim, %c-1 : index
    cf.assert %13, "mismatched size for broadcast"
    %14 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%12, %2 : tensor<?xi64>, tensor<?xi64>) outs(%11 : tensor<?xi64>) {
    ^bb0(%in: i64, %in_5: i64, %out: i64):
      %47 = arith.muli %in, %in_5 : i64
      linalg.yield %47 : i64
    } -> tensor<?xi64>
    %15 = arith.index_cast %dim_2 : index to i64
    %16 = arith.cmpi sge, %15, %c0_i64 : i64
    cf.assert %16, "negative values not allowed in new dimensions"
    %17 = arith.index_cast %15 : i64 to index
    %18 = tensor.empty(%17) : tensor<?xi64>
    %19 = linalg.generic {indexing_maps = [#map1, #map], iterator_types = ["parallel"]} ins(%cst : tensor<i64>) outs(%18 : tensor<?xi64>) {
    ^bb0(%in: i64, %out: i64):
      linalg.yield %in : i64
    } -> tensor<?xi64>
    %dim_3 = tensor.dim %10, %c0 : tensor<?xi64>
    %20 = tensor.empty(%dim_3) : tensor<?x1xi32>
    %21 = linalg.fill ins(%c0_i32 : i32) outs(%20 : tensor<?x1xi32>) -> tensor<?x1xi32>
    %22 = tensor.empty(%dim_3) : tensor<?xi64>
    %23 = linalg.fill ins(%c0_i64 : i64) outs(%22 : tensor<?xi64>) -> tensor<?xi64>
    %24:2 = linalg.generic {indexing_maps = [#map2, #map], iterator_types = ["parallel"]} outs(%21, %23 : tensor<?x1xi32>, tensor<?xi64>) {
    ^bb0(%out: i32, %out_5: i64):
      %47 = linalg.index 0 : index
      %48 = arith.remsi %47, %dim_3 : index
      %extracted_6 = tensor.extract %10[%48] : tensor<?xi64>
      %extracted_7 = tensor.extract %14[%48] : tensor<?xi64>
      %49 = arith.trunci %extracted_6 : i64 to i32
      linalg.yield %49, %extracted_7 : i32, i64
    } -> (tensor<?x1xi32>, tensor<?xi64>)
    %inserted_slice = tensor.insert_slice %24#0 into %21[0, 0] [%dim_3, 1] [1, 1] : tensor<?x1xi32> into tensor<?x1xi32>
    %25 = tm_tensor.scatter {dimension_map = array<i64: 0>} unique_indices(false) ins(%23, %inserted_slice : tensor<?xi64>, tensor<?x1xi32>) outs(%19 : tensor<?xi64>) {
    ^bb0(%arg1: i64, %arg2: i64):
      tm_tensor.yield %arg1 : i64
    } -> tensor<?xi64>
    %26 = tm_tensor.scatter {dimension_map = array<i64: 0>} unique_indices(false) ins(%24#1, %inserted_slice : tensor<?xi64>, tensor<?x1xi32>) outs(%25 : tensor<?xi64>) {
    ^bb0(%arg1: i64, %arg2: i64):
      %47 = arith.addi %arg1, %arg2 : i64
      tm_tensor.yield %47 : i64
    } -> tensor<?xi64>
    %27 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["reduction"]} ins(%2 : tensor<?xi64>) outs(%6 : tensor<i64>) {
    ^bb0(%in: i64, %out: i64):
      %47 = arith.addi %in, %out : i64
      linalg.yield %47 : i64
    } -> tensor<i64>
    %extracted = tensor.extract %27[] : tensor<i64>
    %dim_4 = tensor.dim %26, %c0 : tensor<?xi64>
    %28 = arith.index_cast %dim_4 : index to i64
    %29 = arith.cmpi slt, %28, %c0_i64 : i64
    %30 = arith.select %29, %28, %c0_i64 : i64
    %31 = arith.index_cast %30 : i64 to index
    %32 = arith.index_cast %extracted : i64 to index
    %33 = arith.cmpi slt, %32, %c0 : index
    %34 = arith.addi %32, %dim_4 : index
    %35 = arith.select %33, %34, %32 : index
    %36 = arith.cmpi slt, %35, %c0 : index
    %37 = arith.select %36, %c-1, %35 : index
    %38 = arith.cmpi sgt, %37, %dim_4 : index
    %39 = arith.select %38, %dim_4, %37 : index
    %40 = arith.subi %39, %31 : index
    %41 = arith.cmpi slt, %40, %c0 : index
    %42 = arith.select %41, %c0, %40 : index
    %extracted_slice = tensor.extract_slice %26[%31] [%42] [1] : tensor<?xi64> to tensor<?xi64>
    %expanded = tensor.expand_shape %extracted_slice [[0, 1]] output_shape [%42, 1] : tensor<?xi64> into tensor<?x1xi64>
    %43 = tensor.empty(%42) : tensor<?x1xi64>
    %44 = linalg.generic {indexing_maps = [#map3, #map4, #map5], iterator_types = ["parallel", "parallel"]} ins(%expanded, %cst_0 : tensor<?x1xi64>, tensor<1x1xi64>) outs(%43 : tensor<?x1xi64>) {
    ^bb0(%in: i64, %in_5: i64, %out: i64):
      %47 = arith.sitofp %in : i64 to f64
      %48 = arith.sitofp %in_5 : i64 to f64
      %49 = arith.divf %47, %48 : f64
      %50 = math.floor %49 : f64
      %51 = arith.fptosi %50 : f64 to i64
      linalg.yield %51 : i64
    } -> tensor<?x1xi64>
    %45 = tensor.empty() : tensor<1x1xi64>
    %46 = linalg.generic {indexing_maps = [#map3, #map6, #map5], iterator_types = ["parallel", "parallel"]} ins(%44, %from_elements : tensor<?x1xi64>, tensor<1xi64>) outs(%45 : tensor<1x1xi64>) {
    ^bb0(%in: i64, %in_5: i64, %out: i64):
      %47 = arith.remsi %in, %in_5 : i64
      %48 = arith.cmpi ne, %47, %c0_i64 : i64
      %49 = arith.cmpi slt, %in_5, %c0_i64 : i64
      %50 = arith.cmpi slt, %47, %c0_i64 : i64
      %51 = arith.xori %49, %50 : i1
      %52 = arith.andi %48, %51 : i1
      %53 = arith.addi %47, %in_5 : i64
      %54 = arith.select %52, %53, %47 : i64
      linalg.yield %54 : i64
    } -> tensor<1x1xi64>
    %transposed = linalg.transpose ins(%46 : tensor<1x1xi64>) outs(%45 : tensor<1x1xi64>) permutation = [1, 0] 
    return %transposed : tensor<1x1xi64>
  }
}

Running AtenNonzero1DModule_one_nonzero...
mismatched size for broadcast
./build_tools/ci/test_posix.sh: line 12: 3770074 Aborted                 (core dumped) python -m e2e_testing.main --config=onnx -v --filter AtenNonzero1DModule_one_nonzero