// RUN: npcomp-opt %s -aten-recognize-kernels -split-input-file |& FileCheck %s
// Note that this test is not exhaustive with respect to ops (since the facility
// is generic). Instead, it uses examplar ops to test various types of
// conversions.

// CHECK-LABEL: func @graph
func @graph(%arg0: !numpy.ndarray<*:?>, %arg1 : !numpy.ndarray<*:?>, %arg2 : si64) -> !numpy.ndarray<*:?> {
  // CHECK: %[[LHS:.*]] = numpy.copy_to_tensor %arg0
  // CHECK: %[[RHS:.*]] = numpy.copy_to_tensor %arg1
  // CHECK: %[[RESULT_IMM:.*]] = "aten.add"(%[[LHS]], %[[RHS]], %arg2) : (tensor<*x!basicpy.UnknownType>, tensor<*x!basicpy.UnknownType>, si64) -> tensor<*x!basicpy.UnknownType>
  // CHECK: %[[RESULT_MUT:.*]] = numpy.create_array_from_tensor %[[RESULT_IMM]] : (tensor<*x!basicpy.UnknownType>) -> !numpy.ndarray<*:?>
  %0 = torch.kernel_call "aten::add" %arg0, %arg1, %arg2 : (!numpy.ndarray<*:?>, !numpy.ndarray<*:?>, si64) -> !numpy.ndarray<*:?> {sigArgTypes = ["Tensor", "Tensor", "Scalar"], sigIsMutable = false, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor"]}
  // CHECK: return %[[RESULT_MUT]]
  return %0 : !numpy.ndarray<*:?>
}


// -----
// CHECK-LABEL: func @nll_loss2d_forward
// Contains a Tensor? type mapped to None.
func @nll_loss2d_forward(
    %arg0: !numpy.ndarray<[3,4,8,8]:f32>,
    %arg1: !numpy.ndarray<[3,8,8]:i64>,
    %arg2: !basicpy.NoneType,
    %arg3: i64,
    %arg4: i64) -> (!numpy.ndarray<[]:f32>, !numpy.ndarray<[]:f32>) {
  // CHECK: %[[TARG0:.*]] = numpy.copy_to_tensor %arg0
  // CHECK: %[[TARG1:.*]] = numpy.copy_to_tensor %arg1
  // CHECK: %[[TOUTPUT:.*]], %[[TTOTAL_WEIGHT:.*]] = "aten.nll_loss2d_forward"(%[[TARG0]], %[[TARG1]], %arg2, %arg3, %arg4) : (tensor<3x4x8x8xf32>, tensor<3x8x8xi64>, !basicpy.NoneType, i64, i64) -> (tensor<f32>, tensor<f32>)
  // CHECK: %[[AOUTPUT:.*]] = numpy.create_array_from_tensor %[[TOUTPUT]]
  // CHECK: %[[ATOTAL_WEIGHT:.*]] =  numpy.create_array_from_tensor %[[TTOTAL_WEIGHT]]
  %0:2 = torch.kernel_call "aten::nll_loss2d_forward"
      %arg0, %arg1, %arg2, %arg3, %arg4 :
      (!numpy.ndarray<[3,4,8,8]:f32>, !numpy.ndarray<[3,8,8]:i64>, !basicpy.NoneType, i64, i64) ->
      (!numpy.ndarray<[]:f32>, !numpy.ndarray<[]:f32>)
      {sigArgTypes = ["Tensor", "Tensor", "Tensor?", "int", "int"], sigIsMutable = false, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor", "Tensor"]}
  // CHECK: return %[[AOUTPUT]], %[[ATOTAL_WEIGHT]]
  return %0#0, %0#1 : !numpy.ndarray<[]:f32>, !numpy.ndarray<[]:f32>
}

// -----
// CHECK-LABEL: func @convolution
// Contains a Tensor?, bool, int and list types.
func @convolution(
    %arg0: !numpy.ndarray<[3,16,10,10]:f32>, %arg1: !numpy.ndarray<[4,16,3,3]:f32>,
    %arg2: !numpy.ndarray<[4]:f32>, %arg3: !basicpy.ListType, %arg4: !basicpy.ListType,
    %arg5: !basicpy.ListType, %arg6: i1, %arg7: !basicpy.ListType, %arg8: i64) -> !numpy.ndarray<[3,4,8,8]:f32> {
  // CHECK: %[[TARG0:.*]] = numpy.copy_to_tensor %arg0
  // CHECK: %[[TARG1:.*]] = numpy.copy_to_tensor %arg1
  // CHECK: %[[TARG2:.*]] = numpy.copy_to_tensor %arg2
  // CHECK: %[[TRESULT:.*]] = "aten.convolution"(%[[TARG0]], %[[TARG1]], %[[TARG2]], %arg3, %arg4, %arg5, %arg6, %arg7, %arg8) : (tensor<3x16x10x10xf32>, tensor<4x16x3x3xf32>, tensor<4xf32>, !basicpy.ListType, !basicpy.ListType, !basicpy.ListType, i1, !basicpy.ListType, i64) -> tensor<3x4x8x8xf32>
  %0 = torch.kernel_call "aten::convolution"
      %arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8 :
      (!numpy.ndarray<[3,16,10,10]:f32>, !numpy.ndarray<[4,16,3,3]:f32>,
       !numpy.ndarray<[4]:f32>, !basicpy.ListType, !basicpy.ListType,
       !basicpy.ListType, i1, !basicpy.ListType, i64) -> !numpy.ndarray<[3,4,8,8]:f32>
       {sigArgTypes = ["Tensor", "Tensor", "Tensor?", "int[]", "int[]", "int[]", "bool", "int[]", "int"], sigIsMutable = false, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor"]}
  return %0 : !numpy.ndarray<[3,4,8,8]:f32>
}

// -----
// CHECK-LABEL: func @convolution_backward
// Interesting because it has optional tensor returns.
func @convolution_backward(
    %arg0: !numpy.ndarray<[3,4,8,8]:f32>,
    %arg1: !numpy.ndarray<[3,16,10,10]:f32>,
    %arg2: !numpy.ndarray<[4,16,3,3]:f32>,
    %arg3: !basicpy.ListType,
    %arg4: !basicpy.ListType,
    %arg5: !basicpy.ListType,
    %arg6: i1,
    %arg7: !basicpy.ListType,
    %arg8: i64,
    %arg9: !basicpy.ListType) -> (!basicpy.NoneType, !numpy.ndarray<[4,16,3,3]:f32>, !numpy.ndarray<[4]:f32>) {
  // CHECK: %[[GRAD_INPUT:.*]], %[[GRAD_WEIGHT:.*]], %[[GRAD_BIAS:.*]] = "aten.convolution_backward"
  // Note that this kernel call masks out the input gradients, which will return as NoneType.
  %0:3 = torch.kernel_call "aten::convolution_backward"
    %arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9 :
    (!numpy.ndarray<[3,4,8,8]:f32>, !numpy.ndarray<[3,16,10,10]:f32>, !numpy.ndarray<[4,16,3,3]:f32>, !basicpy.ListType, !basicpy.ListType, !basicpy.ListType, i1, !basicpy.ListType, i64, !basicpy.ListType) ->
    (!basicpy.NoneType, !numpy.ndarray<[4,16,3,3]:f32>, !numpy.ndarray<[4]:f32>) {sigArgTypes = ["Tensor", "Tensor", "Tensor", "int[]", "int[]", "int[]", "bool", "int[]", "int", "bool[]"], sigIsMutable = false, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor", "Tensor", "Tensor"]}
  // CHECK: %[[AGRAD_WEIGHT:.*]] = numpy.create_array_from_tensor %[[GRAD_WEIGHT]]
  // CHECK: %[[AGRAD_BIAS:.*]] = numpy.create_array_from_tensor %[[GRAD_BIAS]]
  // Key thing: The return returns the raw NoneType from the masked input gradient
  // and it does not get converted to an array.
  // CHECK: return %[[GRAD_INPUT]], %[[AGRAD_WEIGHT]], %[[AGRAD_BIAS]]
  return %0#0, %0#1, %0#2 : !basicpy.NoneType, !numpy.ndarray<[4,16,3,3]:f32>, !numpy.ndarray<[4]:f32>
}

// -----
// CHECK-LABEL: func @copy_inplace
// Mutable/in-place op conversion, dropping result.
func @copy_inplace(%arg0: !numpy.ndarray<[4]:f32>, %arg1: !numpy.ndarray<[4]:f32>) -> !numpy.ndarray<[4]:f32> {
  // CHECK: %[[TARG1:.*]] = numpy.copy_to_tensor %arg1
  // CHECK: "aten.copy.inplace"(%arg0, %[[TARG1]]) : (!numpy.ndarray<[4]:f32>, tensor<4xf32>) -> ()
  %0 = torch.kernel_call "aten::copy_" %arg0, %arg1 : (!numpy.ndarray<[4]:f32>, !numpy.ndarray<[4]:f32>) -> !numpy.ndarray<[4]:f32> {sigArgTypes = ["Tensor", "Tensor", "bool"], sigIsMutable = true, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor"]}
  // CHECK: return %arg0
  return %0 : !numpy.ndarray<[4]:f32>
}