aten::masked_fill (#4708)

Summary:
Pull Request resolved: #4708

Add torch_glow support for aten::masked_fill operator

Reviewed By: mortzur

Differential Revision: D22610504

fbshipit-source-id: 60fb272a95f7fb4758b4258a3b88625e9df49970

Author: jackm321

torch_glow/src/PyTorchModelLoader.cpp

@@ -642,6 +642,15 @@ struct FlattenInputs {
   };
 };
 
+/// Indexes of aten::masked_fill inputs.
+struct MaskedFillInputs {
+  enum {
+    input = 0,
+    mask = 1,
+    value = 2,
+  };
+};
+
 /// Indexes of aten::topk inputs.
 struct TopKInputs {
   enum {
@@ -803,6 +812,8 @@ PyTorchModelLoader::buildSymbolsMapping() {
       {{"aten::bmm"}, &PyTorchModelLoader::loadBmm},
       {{"aten::addmm"}, &PyTorchModelLoader::loadAddMM},
       {{"aten::flatten"}, &PyTorchModelLoader::loadFlatten},
+      {{"aten::masked_fill", "aten::masked_fill_"},
+       &PyTorchModelLoader::loadMaskedFill},
       {{"aten::prelu"}, &PyTorchModelLoader::loadPRelu},
       {{"aten::slice"}, &PyTorchModelLoader::loadSlice},
       {{"aten::softmax"}, &PyTorchModelLoader::loadSoftMax},
@@ -3522,6 +3533,48 @@ Error PyTorchModelLoader::loadTo(const torch::jit::Node *ptNode) {
   return addValueMapping(outputs[0], in);
 }
 
+Error PyTorchModelLoader::loadMaskedFill(const torch::jit::Node *ptNode) {
+  auto inputs = ptNode->inputs();
+  auto outputs = ptNode->outputs();
+  RETURN_IF_ERR(checkInputAndOutputSizes(inputs, 3, outputs, 1));
+
+  glow::NodeValue in;
+  ASSIGN_VALUE_OR_RETURN_ERR(
+      in, getGlowNodeValueForValue(inputs[MaskedFillInputs::input]));
+
+  glow::NodeValue mask;
+  ASSIGN_VALUE_OR_RETURN_ERR(
+      mask, getGlowNodeValueForValue(inputs[MaskedFillInputs::mask]));
+
+  size_t inSize = in.dims().size();
+  size_t maskSize = mask.dims().size();
+
+  RETURN_ERR_IF_NOT(
+      inSize >= maskSize,
+      strFormat("masked_fill must have inputs at least as large as mask got "
+                "input of size %zu and mask of size %zu",
+                inSize, maskSize));
+
+  size_t maskBroadcastAxis = inSize - maskSize;
+  if (maskBroadcastAxis > 0) {
+    mask = F_.createBroadcast("broadcast", mask, in.dims(), maskBroadcastAxis)
+               ->getNthResult(0);
+  }
+
+  float value;
+  ASSIGN_VALUE_OR_RETURN_ERR(value, iValToDouble(getGlowIValueForValue(
+                                        inputs[MaskedFillInputs::value])));
+
+  auto valueSplat =
+      F_.createSplat("masked_fill_value",
+                     F_.getParent()->uniqueType(ElemKind::FloatTy, in.dims()),
+                     value)
+          ->getResult();
+
+  auto out = F_.createSelect("masked_fill", mask, valueSplat, in);
+  return addValueMapping(outputs[0], out);
+}
+
 Error PyTorchModelLoader::loadFlatten(const torch::jit::Node *ptNode) {
   auto inputs = ptNode->inputs();
   auto outputs = ptNode->outputs();

torch_glow/src/PyTorchModelLoader.h

@@ -584,6 +584,10 @@ class PyTorchModelLoader {
   /// \returns error on failure.
   Error loadFlatten(const torch::jit::Node *ptNode);
 
+  /// Load a PyTorch aten::masked_fill node.
+  /// \returns error on failure.
+  Error loadMaskedFill(const torch::jit::Node *ptNode);
+
   /// Load a PyTorch topK node.
   /// \returns error on failure.
   Error loadTopK(const torch::jit::Node *ptNode);

torch_glow/tests/nodes/masked_fill_test.py

@@ -0,0 +1,43 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import torch
+from tests.utils import jitVsGlow
+import unittest
+
+
+class TestMaskedFill(unittest.TestCase):
+    def test_masked_fill_basic(self):
+        """Test of the PyTorch aten::masked_fill op on Glow."""
+
+        def masked_fill(a, mask):
+            return torch.masked_fill(a + a, mask, 42.0)
+
+        x = torch.randn([3])
+        mask = torch.tensor([True, False, True], dtype=torch.bool)
+
+        jitVsGlow(masked_fill, x, mask, expected_fused_ops={"aten::masked_fill"})
+
+    def test_masked_fill_broadcasted(self):
+        """Test of the PyTorch aten::masked_fill op on Glow with a
+        broadcasted mask"""
+
+        def masked_fill(a, mask):
+            return torch.masked_fill(a + a, mask, 42.0)
+
+        x = torch.randn([4, 1, 3])
+        mask = torch.tensor([True, False, True], dtype=torch.bool)
+
+        jitVsGlow(masked_fill, x, mask, expected_fused_ops={"aten::masked_fill"})
+
+    def test_masked_fill_inplace(self):
+        """Test of the PyTorch aten::masked_fill_ op on Glow"""
+
+        def masked_fill(a, mask):
+            b = a + a
+            b.masked_fill_(mask, 42.0)
+            return b
+
+        x = torch.randn([3])
+        mask = torch.tensor([True, False, True], dtype=torch.bool)
+
+        jitVsGlow(masked_fill, x, mask, expected_fused_ops={"aten::masked_fill_"})