Add aten::nll_loss_backward op

Prashant Kumar · Prashant Kumar · commit 3b6a7457e7d7 · 2021-12-07T19:54:09.000+05:30
The lowering of aten::nll_loss_backward op has been added
from torch to linalg dialect. The changes has been made as
a part of -torch-convert-to-linalg pass.

Signed-off-by: Prashant Kumar &lt;prashant@nod-labs.com&gt;
diff --git a/e2e_testing/torchscript/nll_loss.py b/e2e_testing/torchscript/nll_loss.py
@@ -60,3 +60,62 @@ def forward(self, x, y):
 @register_test_case(module_factory=lambda: NllLossModule_ignore_index_out_of_bounds())
 def NllLossModule_ignore_index(module, tu: TestUtils):
   module.forward(tu.rand(2, 3), torch.tensor([0, 1]))
+
+
+class NllLossModule_backward(torch.nn.Module):
+
+  def __init__(self):
+    super().__init__()
+
+  @export
+  @annotate_args([
+      None,
+      ([-1], torch.float32, True),
+      ([-1, -1], torch.float32, True),
+      ([-1], torch.int64, True),
+      ([], torch.float32, True),
+  ])
+  def forward(self, grad_output, input, target, total_weight):
+    return torch.ops.aten.nll_loss_backward(grad_output=grad_output,
+                                           self=input,
+                                           target=target,
+                                           weight=None,
+                                           reduction=0,
+                                           ignore_index=10,
+                                           total_weight=total_weight)
+
+
+@register_test_case(module_factory=lambda: NllLossModule_backward())
+def NllLossModuleBackward_basic(module, tu: TestUtils):
+  module.forward(tu.rand(3), tu.rand(3, 4), torch.tensor([2, 3, 0]),
+                 torch.tensor(3.))
+
+
+class NllLossModule_backward_ignore_index_out_of_bounds(torch.nn.Module):
+
+  def __init__(self):
+    super().__init__()
+
+  @export
+  @annotate_args([
+      None,
+      ([-1], torch.float32, True),
+      ([-1, -1], torch.float32, True),
+      ([-1], torch.int64, True),
+      ([], torch.float32, True),
+  ])
+  def forward(self, grad_output, input, target, total_weight):
+    return torch.ops.aten.nll_loss_backward(grad_output=grad_output,
+                                            self=input,
+                                            target=target,
+                                            weight=None,
+                                            reduction=0,
+                                            ignore_index=1,
+                                            total_weight=total_weight)
+
+
+@register_test_case(
+    module_factory=lambda: NllLossModule_backward_ignore_index_out_of_bounds())
+def NllLossModuleBackward_ignore_index(module, tu: TestUtils):
+  module.forward(tu.rand(3), tu.rand(3, 4), torch.tensor([2, 3, 0]),
+                 torch.tensor(3.))
diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td
@@ -1592,6 +1592,26 @@ def Torch_AtenNllLossForwardOp : Torch_Op<"aten.nll_loss_forward", [
   let assemblyFormat = "$self `,` $target `,` $weight `,` $reduction `,` $ignore_index attr-dict `:` type($self) `,` type($target) `,` type($weight) `,` type($reduction) `,` type($ignore_index) `->` type($output) `,` type($total_weight)";
 }
 
+def Torch_AtenNllLossBackwardOp : Torch_Op<"aten.nll_loss_backward", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::nll_loss_backward : (Tensor, Tensor, Tensor, Tensor?, int, int, Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$grad_output,
+    AnyTorchTensorType:$self,
+    AnyTorchTensorType:$target,
+    AnyTorchOptionalTensorType:$weight,
+    Torch_IntType:$reduction,
+    Torch_IntType:$ignore_index,
+    AnyTorchTensorType:$total_weight
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$grad_output `,` $self `,` $target `,` $weight `,` $reduction `,` $ignore_index `,` $total_weight attr-dict `:` type($grad_output) `,` type($self) `,` type($target) `,` type($weight) `,` type($reduction) `,` type($ignore_index) `,` type($total_weight) `->` type($result)";
+}
+
 def Torch_AtenUnsqueezeOp : Torch_Op<"aten.unsqueeze", [
     AllowsTypeRefinement
   ]> {
diff --git a/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp b/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp
@@ -1263,6 +1263,107 @@ class ConvertAtenNllLossForwardOp
 };
 } // namespace
 
+// Given `grad_output`, `input`, `target`, `nll_loss_backward` is given by:
+//   for i in range(0, len(input[0])):
+//      for j in range(0, len(input[1])):
+//          nll_loss_backward[i][j] = (j == target[i]) ? grad_output[i] : 0
+// TODO: `weight` and `reduction` operands are still to be taken care of.
+namespace {
+class ConvertAtenNllLossBackwardOp
+    : public OpConversionPattern<AtenNllLossBackwardOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(AtenNllLossBackwardOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
+      return failure();
+    Location loc = op->getLoc();
+    Value input = adaptor.self();
+    Value target = adaptor.target();
+    Value weight = adaptor.weight();
+    Value grad_output = adaptor.grad_output();
+
+    int64_t reduce_dim;
+    if (!matchPattern(op.reduction(), m_TorchConstantInt(&reduce_dim)))
+      return rewriter.notifyMatchFailure(op, "dim must be constant");
+
+    // TODO: Handle reduction.
+    if (reduce_dim != 0)
+      return rewriter.notifyMatchFailure(
+          op, "reduction along dimensions is not supported.");
+
+    // TODO: Incorporate the weight argument.
+    if (!weight.getType().isa<mlir::torch::Torch::NoneType>())
+      return rewriter.notifyMatchFailure(
+          op, "Unimplemented, the weight operand is not incorporated.");
+
+    Value ignoreIndex = adaptor.ignore_index();
+    Value ignoreIndexVal = castIntToIndex(rewriter, loc, ignoreIndex);
+
+    unsigned inputRank = input.getType().cast<RankedTensorType>().getRank();
+    unsigned targetRank = target.getType().cast<RankedTensorType>().getRank();
+
+    // TODO: Cases with targetRank != 1 where `Mean` reduction is required.
+    if (inputRank != 2 || targetRank != 1) {
+      return rewriter.notifyMatchFailure(
+          op, "expected  input and target to be rank 2 and 1 respectively");
+    }
+    RankedTensorType resultType = getTypeConverter()
+                          ->convertType(op->getResult(0).getType())
+                          .cast<RankedTensorType>();
+
+    Type elementType = resultType.getElementType();
+
+    // Given there is no reduction `grad_input` size is equal to `input` size.
+    auto outputSize = getTensorSizes(rewriter, loc, input);
+    Value initTensor0 =
+        createZeroInitTensor(rewriter, loc, outputSize, elementType);
+    Value zeroVal = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getZeroAttr(elementType));
+
+    SmallVector<AffineExpr> targetExpr{rewriter.getAffineDimExpr(0)};
+    SmallVector<AffineExpr> resultExpr{rewriter.getAffineDimExpr(0),
+                                       rewriter.getAffineDimExpr(1)};
+    SmallVector<StringRef> iteratorTypes{getParallelIteratorTypeName(),
+                                         getParallelIteratorTypeName()};
+    auto indexingMaps =
+        AffineMap::inferFromExprList({targetExpr, targetExpr, resultExpr});
+    Value finalRes =
+        rewriter
+            .create<linalg::GenericOp>(
+                loc, resultType, ValueRange{target, grad_output}, initTensor0,
+                /*indexingMaps=*/indexingMaps,
+                /*iteratorTypes=*/iteratorTypes,
+                [&](OpBuilder &b, Location loc, ValueRange args) {
+                  Value indTarget = rewriter.create<arith::IndexCastOp>(
+                      loc, rewriter.getIndexType(), args[0]);
+                  Value indJ = rewriter.create<linalg::IndexOp>(loc, 1);
+
+                  // The final result is given by:
+                  // grad_input[i][j] = (j == target[i]) ? grad_output[i] : 0
+                  Value cmpEq = rewriter.create<arith::CmpIOp>(
+                      loc, arith::CmpIPredicate::eq, indJ, indTarget);
+
+                  // The target index shouldn't be equal to `ignoreIndex`.
+                  Value cmpNEq = rewriter.create<arith::CmpIOp>(
+                      loc, arith::CmpIPredicate::ne, ignoreIndexVal, indTarget);
+                  Value finalPredicate =
+                      rewriter.create<arith::AndIOp>(loc, cmpEq, cmpNEq);
+                  Value negate =
+                      rewriter.create<arith::NegFOp>(loc, elementType, args[1]);
+                  Value selectFinal = rewriter.create<mlir::SelectOp>(
+                      loc, finalPredicate, negate, zeroVal);
+                  b.create<linalg::YieldOp>(loc, selectFinal);
+                })
+            .getResult(0);
+
+    rewriter.replaceOp(op, finalRes);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 // See comments at in convertMmOp and the heading for this section for general
 // considerations. This function needs to be auto-generated.
@@ -3470,6 +3571,8 @@ class ConvertTorchToLinalg
     patterns.add<ConvertAtenSliceTensorOp>(typeConverter, context);
     target.addIllegalOp<AtenNllLossForwardOp>();
     patterns.add<ConvertAtenNllLossForwardOp>(typeConverter, context);
+    target.addIllegalOp<AtenNllLossBackwardOp>();
+    patterns.add<ConvertAtenNllLossBackwardOp>(typeConverter, context);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
diff --git a/lib/Dialect/Torch/Transforms/RefineTypes.cpp b/lib/Dialect/Torch/Transforms/RefineTypes.cpp
@@ -454,11 +454,12 @@ class TypeAnalyzer : public ForwardDataFlowAnalysis<ValueKnowledge> {
       return visitAtenAddCLikeOp(op, operands);
     } else if (auto scalarOp = dyn_cast<AtenAddIntOp>(op)) {
       return visitBinaryScalarOp(scalarOp);
-    }else if (auto nllForwardOp = dyn_cast<AtenNllLossForwardOp>(op)) {
+    } else if (auto nllForwardOp = dyn_cast<AtenNllLossForwardOp>(op)) {
       return visitAtenNllLossForwardOp(nllForwardOp, operands);
+    } else if (auto nllBackwardOp = dyn_cast<AtenNllLossBackwardOp>(op)) {
+      return visitAtenNllLossBackwardOp(nllBackwardOp, operands);
     }
 
-
     // Otherwise, this is an unknown operation. Just mark all results as
     // having reached a pessimistic fixpoint.
     return markAllPessimisticFixpoint(op->getResults());
@@ -584,9 +585,13 @@ class TypeAnalyzer : public ForwardDataFlowAnalysis<ValueKnowledge> {
   visitAten_SoftmaxOp(Aten_SoftmaxOp op,
                       ArrayRef<LatticeElement<ValueKnowledge> *> operands);
 
-  ChangeResult
-  visitAtenNllLossForwardOp(AtenNllLossForwardOp op,
-                      ArrayRef<LatticeElement<ValueKnowledge> *> operands);
+  ChangeResult visitAtenNllLossForwardOp(
+      AtenNllLossForwardOp op,
+      ArrayRef<LatticeElement<ValueKnowledge> *> operands);
+
+  ChangeResult visitAtenNllLossBackwardOp(
+      AtenNllLossBackwardOp op,
+      ArrayRef<LatticeElement<ValueKnowledge> *> operands);
 };
 } // namespace
 
@@ -966,6 +971,27 @@ ChangeResult TypeAnalyzer::visitAtenNllLossForwardOp(
   return resultLattice;
 }
 
+ChangeResult TypeAnalyzer::visitAtenNllLossBackwardOp(
+    AtenNllLossBackwardOp op,
+    ArrayRef<LatticeElement<ValueKnowledge> *> operands) {
+  auto self = operands[1]->getValue();
+  auto knowledge =
+      ValueKnowledge::getNotNonePessimisticValueState(op.getContext());
+
+  knowledge.dtype = self.dtype;
+  int64_t reduction;
+  unsigned resultRank = self.sizes.size();
+
+  if (self.hasSizes &&
+      matchPattern(op.reduction(), m_TorchConstantInt(&reduction))) {
+    // reduction == 1 means reduce 1st dim.
+    resultRank = reduction == 1 ? resultRank - 1 : resultRank;
+  }
+  knowledge.sizes.resize(resultRank, kUnknownSize);
+  knowledge.hasSizes = true;
+  return getLatticeElement(op.getResult()).join(knowledge);
+}
+
 ChangeResult TypeAnalyzer::visitAtenUnsqueezeOp(
     AtenUnsqueezeOp op, ArrayRef<LatticeElement<ValueKnowledge> *> operands) {
   auto operand = operands[0]->getValue();
diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py
@@ -524,6 +524,7 @@ def emit_with_mutating_variants(key, **kwargs):
         emit("aten::_softmax : (Tensor, int, bool) -> (Tensor)")
         emit("aten::mean : (Tensor, int?) -> (Tensor)")
         emit("aten::nll_loss_forward : (Tensor, Tensor, Tensor?, int, int) -> (Tensor, Tensor)")
+        emit("aten::nll_loss_backward : (Tensor, Tensor, Tensor, Tensor?, int, int, Tensor) -> (Tensor)")
 
         # Misc tensor ops.
         emit("aten::unsqueeze : (Tensor, int) -> (Tensor)")
