Fix addmv handler - ignore input when beta is 0

Summary:
According to the pytorch documentation, input should be ignored when
beta is 0.

!ci_branch_mk2

Reviewers: #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, #pytorch, dariuszs

Reviewed By: #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, #pytorch, dariuszs

Subscribers: dariuszs

JIRA Issues: AFS-357

Differential Revision: https://phabricator.sourcevertex.net/D84820

Author: jreniecki

poptorch/source/popart_canonicalization/BlasOps.cpp

@@ -112,16 +112,38 @@ torch::jit::Node *baddbmmHandler(torch::jit::Graph *graph,
 
 torch::jit::Node *addmvHandler(torch::jit::Graph *graph,
                                torch::jit::Node *node) {
+  auto *input = node->input(0);
   auto *mat = node->input(1);
   auto *vec = node->input(2);
-  auto *t0 = createMatmul(graph, {mat, vec})->output();
-  auto *alpha = node->input(4);
-  auto *t1 = createMul(graph, {t0, alpha})->output();
-  auto *input = node->input(0);
   auto *beta = node->input(3);
-  auto *t2 = createMul(graph, {input, beta})->output();
-  // add(mul(matmul(mat, vec), alpha), mul(input, beta))
-  return createAdd(graph, {t1, t2});
+  auto *alpha = node->input(4);
+
+  const auto alpha_val = constantToFloat(alpha->node());
+  const auto beta_val = constantToFloat(beta->node());
+
+  if (alpha_val == 0 && beta_val == 0) {
+    return createConstantFloatLike(graph, input, {0}, {shapeFromTensor(input)});
+  }
+
+  torch::jit::Node *t1 = nullptr;
+  if (alpha_val != 0) {
+    auto *t0 = createMatmul(graph, {mat, vec})->output();
+    t1 = createMul(graph, {t0, alpha});
+  }
+
+  torch::jit::Node *output;
+  if (beta_val != 0) {
+    auto *t2 = createMul(graph, {input, beta});
+    if (t1 != nullptr) {
+      output = createAdd(graph, {t1->output(), t2->output()});
+    } else {
+      output = t2;
+    }
+  } else {
+    output = t1;
+  }
+
+  return output;
 }
 } // namespace
 

tests/blas_test.py

@@ -221,3 +221,51 @@ def forward(self, x1, x2, x3):
     ipu_result = poptorch.inferenceModel(model)(t1, t2, t3)
 
     helpers.assert_allclose(expected=cpu_result, actual=ipu_result)
+
+
+@pytest.mark.parametrize("input_shape", [(20, 10)])
+@pytest.mark.parametrize("beta", [0, .5])
+@pytest.mark.parametrize("alpha", [0, 1.5])
+@pytest.mark.parametrize("use_out", [True, False])
+def test_addmv(input_shape, beta, alpha, use_out):
+    torch.manual_seed(42)
+
+    mat = torch.randn(input_shape)
+    vec = torch.randn(input_shape[1])
+    inp = torch.randn(input_shape[0])
+
+    if beta == 0:
+        # NaNs in input should be ignored
+        inp[0] = float('nan')
+    if alpha == 0:
+        # NaNs in vec or mat should be ignored
+        mat[0, 0] = float('nan')
+        vec[0] = float('nan')
+
+    output = torch.empty(input_shape[0]) if use_out else None
+
+    class AddmvModel(torch.nn.Module):
+        def __init__(self, beta, alpha):
+            super().__init__()
+            self.beta = beta
+            self.alpha = alpha
+
+        def forward(self, inp, mat, vec, out=None):
+            result = torch.addmv(inp,
+                                 mat,
+                                 vec,
+                                 beta=self.beta,
+                                 alpha=self.alpha,
+                                 out=out)
+            if self.beta == 0 and self.alpha == 0:
+                # Avoid empty compute graph
+                result += torch.zeros_like(inp)
+            return result
+
+    model = AddmvModel(beta, alpha)
+    cpu_result = model(inp, mat, vec, out=output)
+    ipu_result = poptorch.inferenceModel(model)(inp, mat, vec, output)
+
+    helpers.assert_allclose(expected=cpu_result, actual=ipu_result)
+    if use_out is True:
+        helpers.assert_allclose(expected=cpu_result, actual=output)