Remove the dependency of the pre-registered op (#74)

Summary:
Pull Request resolved: #74

We can register the TorchScriptOp on the fly

Reviewed By: louisfeng

Differential Revision:
D41002899

Privacy Context Container: L1137347

fbshipit-source-id: dc12c86efc478054c7c75fc155a644fefb4516a1

train/compute/python/lib/config.py

@@ -6,6 +6,8 @@
 from .init_helper import get_logger
 from .iterator import config_iterator_map, ConfigIterator, DefaultConfigIterator
 from .operator import op_map, OperatorInterface
+from .pytorch.operator_impl import TorchScriptOp
+
 
 logger = get_logger()
 
@@ -64,8 +66,12 @@ def input_data_generator(self, value: Type[DataGenerator]):
 def make_op_config(op_name: str, op_info: Dict[str, Any], device: str):
     global op_map
     if (op_name not in op_map) or (not op_map[op_name]):
-        logger.warning(f"{op_name} has no valid callable defined, skipped.")
-        return None
+        if op_name.startswith("aten::"):
+            logger.debug(f"register op: {op_name}")
+            op_map[op_name] = TorchScriptOp(op_name)
+        else:
+            logger.warning(f"{op_name} has no valid callable defined, skipped.")
+            return None
 
     op = op_map[op_name]
     op.device = device

train/compute/python/workloads/pytorch/native_basic_ops.py

@@ -3,7 +3,7 @@
 import torch
 
 from ...lib.operator import OperatorInterface, register_operators
-from ...lib.pytorch.operator_impl import BuildableOp, CallableOp, TorchScriptOp, UnaryOp
+from ...lib.pytorch.operator_impl import BuildableOp, CallableOp, UnaryOp
 
 
 # Unary
@@ -37,15 +37,3 @@
     "torch.nn.Linear": BuildableOp(torch.nn.Linear),
 }
 register_operators(buildable_ops)
-
-
-# Operator schema: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml
-torchscript_ops: Dict[str, OperatorInterface] = {
-    "aten::add": TorchScriptOp("aten::add"),
-    "aten::add_": TorchScriptOp("aten::add_"),
-    "aten::matmul": TorchScriptOp("aten::matmul"),
-    "aten::mul": TorchScriptOp("aten::mul"),
-    "aten::sum": TorchScriptOp("aten::sum"),
-    "aten::linear": TorchScriptOp("aten::linear"),
-}
-register_operators(torchscript_ops)