[ONNX] Add OnnxToTorch lowering for Onnx.NegativeLogLikelihoodLoss Op

[123epsilon]

This implements the Onnx.NegativeLogLikelihoodLoss op using the signature provided here by replacing it with a NLLLossForward op.

Additionally, I included a helper function get_loss_reduction_enum to convert from a string reduction parameter to the corresponding intended integer value since this is an operation that will be reused for any loss function module. This differs from get_reduction_enum in TorchUpstream.cpp which handles the reduce parameter from scatter_reduce type operations.

[vivekkhandelwal1]

Hi @123epsilon @ScottTodd @andfau-amd, do we need this PR for this op lowering or not? Otherwise, we should close this and get this op working through @andfau-amd's patch.

[andfau-amd]

It does look like my patch can expand this op, e.g.:

$ python -m torch_mlir.tools.import_onnx ./mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_nllloss_NC/model.onnx
module {
  func.func public @test_nllloss_NC(%arg0: !torch.vtensor<[3,5],f32>, %arg1: !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
    %none = torch.constant.none
    %0 = call @"('NegativeLogLikelihoodLoss', '', 13, [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 5\0A    }\0A  }\0A}\0A, tensor_type {\0A  elem_type: 7\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], input: \22input\22\0Ainput: \22target\22\0Aoutput: \22loss\22\0Aop_type: \22NegativeLogLikelihoodLoss\22\0Aattribute {\0A  name: \22reduction\22\0A  s: \22none\22\0A  type: STRING\0A}\0A)"(%arg0, %arg1) : (!torch.vtensor<[3,5],f32>, !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32>
    return %0 : !torch.vtensor<[3],f32>
  }
  func.func private @"('NegativeLogLikelihoodLoss', '', 13, [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 5\0A    }\0A  }\0A}\0A, tensor_type {\0A  elem_type: 7\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], input: \22input\22\0Ainput: \22target\22\0Aoutput: \22loss\22\0Aop_type: \22NegativeLogLikelihoodLoss\22\0Aattribute {\0A  name: \22reduction\22\0A  s: \22none\22\0A  type: STRING\0A}\0A)"(%arg0: !torch.vtensor<[3,5],f32>, %arg1: !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "", torch.onnx_meta.producer_version = ""} {
    %none = torch.constant.none
    %0 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %1 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %2 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %3 = torch.operator "onnx.Unsqueeze"(%arg1, %2) : (!torch.vtensor<[3],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[3,1],si64>
    %4 = torch.operator "onnx.GatherElements"(%arg0, %3) {torch.onnx.axis = 1 : si64} : (!torch.vtensor<[3,5],f32>, !torch.vtensor<[3,1],si64>) -> !torch.vtensor<[3,1],f32>
    %5 = torch.operator "onnx.Neg"(%4) : (!torch.vtensor<[3,1],f32>) -> !torch.vtensor<[3,1],f32>
    %6 = torch.operator "onnx.Slice"(%5, %0, %1, %1) : (!torch.vtensor<[3,1],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[3,1],f32>
    %7 = torch.operator "onnx.Squeeze"(%6, %2) : (!torch.vtensor<[3,1],f32>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[3],f32>
    return %7 : !torch.vtensor<[3],f32>
  }
}

However, I haven't so far been able to test whether this expansion is actually correct/useful. FWIW, it does seem like TorchOnnxToTorch doesn't choke on it:

$ python -m torch_mlir.tools.import_onnx ./mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_nllloss_NC/model.onnx | build/bin/torch-mlir-opt --inline --convert-torch-onnx-to-torch --canonicalize
module {
  func.func public @test_nllloss_NC(%arg0: !torch.vtensor<[3,5],f32>, %arg1: !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
    %false = torch.constant.bool false
    %int1 = torch.constant.int 1
    %0 = torch.aten.unsqueeze %arg1, %int1 : !torch.vtensor<[3],si64>, !torch.int -> !torch.vtensor<[3,1],si64>
    %1 = torch.aten.gather %arg0, %int1, %0, %false : !torch.vtensor<[3,5],f32>, !torch.int, !torch.vtensor<[3,1],si64>, !torch.bool -> !torch.vtensor<[3,1],f32>
    %2 = torch.aten.neg %1 : !torch.vtensor<[3,1],f32> -> !torch.vtensor<[3,1],f32>
    %3 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %4 = torch.prims.squeeze %2, %3 : !torch.vtensor<[3,1],f32>, !torch.list<int> -> !torch.vtensor<[3],f32>
    return %4 : !torch.vtensor<[3],f32>
  }
}

[vivekkhandelwal1]

It does look like my patch can expand this op, e.g.:

$ python -m torch_mlir.tools.import_onnx ./mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_nllloss_NC/model.onnx
module {
  func.func public @test_nllloss_NC(%arg0: !torch.vtensor<[3,5],f32>, %arg1: !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
    %none = torch.constant.none
    %0 = call @"('NegativeLogLikelihoodLoss', '', 13, [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 5\0A    }\0A  }\0A}\0A, tensor_type {\0A  elem_type: 7\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], input: \22input\22\0Ainput: \22target\22\0Aoutput: \22loss\22\0Aop_type: \22NegativeLogLikelihoodLoss\22\0Aattribute {\0A  name: \22reduction\22\0A  s: \22none\22\0A  type: STRING\0A}\0A)"(%arg0, %arg1) : (!torch.vtensor<[3,5],f32>, !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32>
    return %0 : !torch.vtensor<[3],f32>
  }
  func.func private @"('NegativeLogLikelihoodLoss', '', 13, [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 5\0A    }\0A  }\0A}\0A, tensor_type {\0A  elem_type: 7\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A  }\0A}\0A], input: \22input\22\0Ainput: \22target\22\0Aoutput: \22loss\22\0Aop_type: \22NegativeLogLikelihoodLoss\22\0Aattribute {\0A  name: \22reduction\22\0A  s: \22none\22\0A  type: STRING\0A}\0A)"(%arg0: !torch.vtensor<[3,5],f32>, %arg1: !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "", torch.onnx_meta.producer_version = ""} {
    %none = torch.constant.none
    %0 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %1 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %2 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %3 = torch.operator "onnx.Unsqueeze"(%arg1, %2) : (!torch.vtensor<[3],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[3,1],si64>
    %4 = torch.operator "onnx.GatherElements"(%arg0, %3) {torch.onnx.axis = 1 : si64} : (!torch.vtensor<[3,5],f32>, !torch.vtensor<[3,1],si64>) -> !torch.vtensor<[3,1],f32>
    %5 = torch.operator "onnx.Neg"(%4) : (!torch.vtensor<[3,1],f32>) -> !torch.vtensor<[3,1],f32>
    %6 = torch.operator "onnx.Slice"(%5, %0, %1, %1) : (!torch.vtensor<[3,1],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[3,1],f32>
    %7 = torch.operator "onnx.Squeeze"(%6, %2) : (!torch.vtensor<[3,1],f32>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[3],f32>
    return %7 : !torch.vtensor<[3],f32>
  }
}

However, I haven't so far been able to test whether this expansion is actually correct/useful. FWIW, it does seem like TorchOnnxToTorch doesn't choke on it:

$ python -m torch_mlir.tools.import_onnx ./mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_nllloss_NC/model.onnx | build/bin/torch-mlir-opt --inline --convert-torch-onnx-to-torch --canonicalize
module {
  func.func public @test_nllloss_NC(%arg0: !torch.vtensor<[3,5],f32>, %arg1: !torch.vtensor<[3],si64>) -> !torch.vtensor<[3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
    %false = torch.constant.bool false
    %int1 = torch.constant.int 1
    %0 = torch.aten.unsqueeze %arg1, %int1 : !torch.vtensor<[3],si64>, !torch.int -> !torch.vtensor<[3,1],si64>
    %1 = torch.aten.gather %arg0, %int1, %0, %false : !torch.vtensor<[3,5],f32>, !torch.int, !torch.vtensor<[3,1],si64>, !torch.bool -> !torch.vtensor<[3,1],f32>
    %2 = torch.aten.neg %1 : !torch.vtensor<[3,1],f32> -> !torch.vtensor<[3,1],f32>
    %3 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %4 = torch.prims.squeeze %2, %3 : !torch.vtensor<[3,1],f32>, !torch.list<int> -> !torch.vtensor<[3],f32>
    return %4 : !torch.vtensor<[3],f32>
  }
}

To me, the above expansion does not seem correct, since I don't see the Log operation involved anywhere.

[andfau-amd]

Well, I don't see any mention of Log in the examples here: https://onnx.ai/onnx/operators/onnx__NegativeLogLikelihoodLoss.html

It seems like the "log" part is not actually in this operation, but rather the inputs are meant to already be logarithmic?

By the way, it seems like the imported ONNX node tests pass for this op: #3409 (comment)

[vivekkhandelwal1]

Well, I don't see any mention of Log in the examples here: https://onnx.ai/onnx/operators/onnx__NegativeLogLikelihoodLoss.html

It seems like the "log" part is not actually in this operation, but rather the inputs are meant to already be logarithmic?

By the way, it seems like the imported ONNX node tests pass for this op: #3409 (comment)

Sorry for the confusion, it's a fault at my end. Based on the op definition lowering looks fine.

[andfau-amd]

Maybe the expansion from this PR is better than the one in mine, since it uses the dedicated Torch NLL op?

[vivekkhandelwal1]

Maybe the expansion from this PR is better than the one in mine, since it uses the dedicated Torch NLL op?

I think taking a conservative approach, we can keep the op lowering now, and can delete it later once the function expansion thing is in and stable?

[andfau-amd]

Sounds good to me.