[refactor] Move spmm and sddmm operators to operator module (#1873)

* upd

* upd

* upd

* upd

docs/source/api/python/index.rst

@@ -10,7 +10,7 @@ API Reference
    batch_heterograph
    nn
    init
-   backend
+   ops
    function
    traversal
    propagate

docs/source/api/python/backend.rst → docs/source/api/python/ops.rst

 .. _apibackend:
 
-.. currentmodule:: dgl.backend
+.. currentmodule:: dgl.ops
 
-dgl.backend
+dgl.ops
 ==================================
 
 Frame-agnostic operators for message passing on graphs.
@@ -36,7 +36,7 @@ here, you can enjoy the same convenience on other frameworks by similar usage):
 
    >>> import dgl
    >>> import torch as th
-   >>> import dgl.backend as F
+   >>> import dgl.ops as F
    >>> g = dgl.graph(([0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2]))  # 3 nodes, 6 edges
    >>> x = th.ones(3, 2, requires_grad=True)
    >>> x
@@ -130,7 +130,7 @@ The following is an example showing how GSDDMM works:
 
    >>> import dgl
    >>> import torch as th
-   >>> import dgl.backend as F
+   >>> import dgl.ops as F
    >>> g = dgl.graph(([0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2]))  # 3 nodes, 6 edges
    >>> x = th.ones(3, 2, requires_grad=True)
    >>> x

python/dgl/__init__.py

@@ -13,6 +13,7 @@ from . import container
 from . import distributed
 from . import random
 from . import sampling
+from . import ops
 
 from ._ffi.runtime_ctypes import TypeCode
 from ._ffi.function import register_func, get_global_func, list_global_func_names, extract_ext_funcs

python/dgl/backend/__init__.py

@@ -7,7 +7,6 @@ import importlib
 
 from . import backend
 from .set_default_backend import set_default_backend
-from itertools import product
 
 _enabled_apis = set()
 
@@ -19,191 +18,6 @@ def _gen_missing_api(api, mod_name):
                           ' the DGLBACKEND environment.' % (api, mod_name))
     return _missing_api
 
-_notes_docstring = r"""
-    Notes
-    -----
-    This function supports autograd (computing input gradients given the output gradient). If the
-    feature shape of two input operands do not match, we first broadcasts the features to a unified
-    shape (note that the memory usage will not increase accordingly) and then performs the operation.
-
-    Broadcasting follows NumPy semantics. Please see
-    https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
-    for more details about the NumPy broadcasting semantics."""
-
-def _gen_sddmm_func(lhs_target, rhs_target, binary_op):
-    name = "{}_{}_{}".format(lhs_target, binary_op, rhs_target)
-    target_dict = {
-        'u': "source node",
-        'e': "edge",
-        'v': "destination node"
-    }
-    lhs_str = target_dict[lhs_target]
-    rhs_str = target_dict[rhs_target]
-    docstring = r"""Generalized SDDMM function.
-    It computes edge features by {} {} features and {} features.
-
-    Parameters
-    ----------
-    g : DGLHeteroGraph
-        The input graph
-    x : tensor
-        The {} features.
-    y : tensor
-        The {} features.
-
-    Returns
-    -------
-    tensor
-        The result tensor.
-    {}""".format(binary_op, lhs_str, rhs_str,
-                 lhs_str, rhs_str,
-                 _notes_docstring)
-
-    def func(g, x, y):
-        return gsddmm(g, binary_op, x, y,
-                      lhs_target=lhs_target, rhs_target=rhs_target)
-    func.__name__ = name
-    func.__doc__ = docstring
-    return func
-
-def _gen_spmm_func(binary_op, reduce_op):
-    name = "u_{}_e_{}".format(binary_op, reduce_op)
-    docstring = """Generalized SpMM function.
-    It fuses two steps into one kernel.
-
-    1. Computes messages by {} source node and edge features.
-    2. Aggregate the messages by {} as the features on destination nodes.
-
-    Parameters
-    ----------
-    g : DGLHeteroGraph
-        The input graph
-    x : tensor
-        The source node features.
-    y : tensor
-        The edge features.
-
-    Returns
-    -------
-    tensor
-        The result tensor.
-    {}""".format(binary_op, reduce_op,
-                 _notes_docstring)
-
-    def func(g, x, y):
-        return gspmm(g, binary_op, reduce_op, x, y)
-    func.__name__ = name
-    func.__doc__ = docstring
-    return func
-
-def _gen_copy_reduce_func(binary_op, reduce_op):
-
-    name = "{}_{}".format(binary_op, reduce_op)
-    binary_str = {
-        "copy_u": "It copies node feature to edge as the message.",
-        'copy_e': "It regards edge feature as message."
-    }
-    x_str = {
-        "copy_u": "source node",
-        "copy_e": "edge"
-    }
-    docstring = lambda binary_op: """Generalized SpMM function. {}
-    Then aggregates the message by {} on destination nodes.
-
-    Parameters
-    ----------
-    g : DGLHeteroGraph
-        The input graph
-    x : tensor
-        The {} features.
-
-    Returns
-    -------
-    tensor
-        The result tensor.
-
-    Notes
-    -----
-    This function supports autograd (computing input gradients given the output gradient).
-    """.format(
-        binary_str[binary_op],
-        reduce_op,
-        x_str[binary_op],
-        _notes_docstring)
-
-    def func(g, x):
-        if binary_op == 'copy_u':
-            return gspmm(g, 'copy_lhs', reduce_op, x, None)
-        else:
-            return gspmm(g, 'copy_rhs', reduce_op, None, x)
-
-    func.__name__ = name
-    func.__doc__ = docstring(binary_op)
-    return func
-
-def _register_sddmm_func(mod, enabled_apis):
-    """Register sddmm functions"""
-    target = ["u", "v", "e"]
-    for lhs, rhs in product(target, target):
-        if lhs != rhs:
-            for binary_op in ["add", "sub", "mul", "div", "dot"]:
-                func = _gen_sddmm_func(lhs, rhs, binary_op)
-                setattr(mod, func.__name__, func)
-                enabled_apis.add(func.__name__)
-
-def _register_spmm_func(mod, enabled_apis):
-    """Register spmm functions"""
-    for binary_op in ["add", "sub", "mul", "div", "copy_u", "copy_e"]:
-        for reduce_op in ["sum", "max", "min"]:
-            if binary_op.startswith("copy"):
-                func = _gen_copy_reduce_func(binary_op, reduce_op)
-            else:
-                func = _gen_spmm_func(binary_op, reduce_op)
-            setattr(mod, func.__name__, func)
-            enabled_apis.add(func.__name__)
-
-def copy_u(g, x):
-    r"""Generalized SDDMM function that copies source node features to edges.
-
-    Parameters
-    ----------
-    g : DGLHeteroGraph
-        The input graph.
-    x : tensor
-        The source node features.
-
-    Returns
-    -------
-    tensor
-        The result tensor.
-
-    Notes
-    -----
-    This function supports autograd (computing input gradients given the output gradient).
-    """
-    return gsddmm(g, 'copy_lhs', x, None)
-
-def copy_v(g, x):
-    r"""Generalized SDDMM function that copies destination node features to edges.
-
-    Parameters
-    ----------
-    g : DGLHeteroGraph
-        The input graph.
-    x : tensor
-        The destination node features.
-
-    Returns
-    -------
-    tensor
-        The result tensor.
-
-    Notes
-    -----
-    This function supports autograd (computing input gradients given the output gradient).
-    """
-    return gsddmm(g, 'copy_rhs', None, x)
-
 def load_backend(mod_name):
     print('Using backend: %s' % mod_name, file=sys.stderr)
     mod = importlib.import_module('.%s' % mod_name, __name__)
@@ -235,12 +49,6 @@ def load_backend(mod_name):
                 setattr(thismod, api, mod.__dict__[api])
             else:
                 setattr(thismod, api, _gen_missing_api(api, mod_name))
-    _register_sddmm_func(thismod, _enabled_apis)
-    _register_spmm_func(thismod, _enabled_apis)
-    setattr(thismod, copy_u.__name__, copy_u)
-    _enabled_apis.add(copy_u.__name__)
-    setattr(thismod, copy_v.__name__, copy_v)
-    _enabled_apis.add(copy_v.__name__)
 
 
 def get_preferred_backend():

python/dgl/nn/mxnet/softmax.py

@@ -2,7 +2,7 @@
 # pylint: disable= no-member, arguments-differ, access-member-before-definition, unpacking-non-sequence
 import mxnet as mx
 
-from ... import backend as F
+from ... import ops as F
 from ...base import ALL, is_all
 
 __all__ = ['edge_softmax']

python/dgl/nn/pytorch/softmax.py

@@ -3,7 +3,7 @@
 import torch as th
 
 from ...base import ALL, is_all
-from ... import backend as F
+from ... import ops as F
 
 __all__ = ['edge_softmax']
 

python/dgl/ops/__init__.py

+"""dgl operator module."""
+from .spmm import *
+from .sddmm import *

python/dgl/ops/sddmm.py

+"""dgl sddmm operator module."""
+from itertools import product
+import sys
+
+from ..backend import gsddmm
+
+__all__ = ['gsddmm', 'copy_u', 'copy_v']
+
+
+def _gen_sddmm_func(lhs_target, rhs_target, binary_op):
+    name = "{}_{}_{}".format(lhs_target, binary_op, rhs_target)
+    target_dict = {
+        'u': "source node",
+        'e': "edge",
+        'v': "destination node"
+    }
+    lhs_str = target_dict[lhs_target]
+    rhs_str = target_dict[rhs_target]
+    docstring = r"""Generalized SDDMM function.
+    It computes edge features by {op} {lhs} features and {rhs} features.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        The input graph
+    x : tensor
+        The {lhs} features.
+    y : tensor
+        The {rhs} features.
+
+    Returns
+    -------
+    tensor
+        The result tensor.
+
+    Notes
+    -----
+    This function supports autograd (computing input gradients given the output gradient). If the
+    feature shape of two input operands do not match, we first broadcasts the features to a unified
+    shape (note that the memory usage will not increase accordingly) and then performs the operation.
+
+    Broadcasting follows NumPy semantics. Please see
+    https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
+    for more details about the NumPy broadcasting semantics.
+    """.format(op=binary_op, lhs=lhs_str, rhs=rhs_str)
+
+    def func(g, x, y):
+        return gsddmm(g, binary_op, x, y,
+                      lhs_target=lhs_target, rhs_target=rhs_target)
+    func.__name__ = name
+    func.__doc__ = docstring
+    return func
+
+def _register_sddmm_func():
+    """Register sddmm functions"""
+    target = ["u", "v", "e"]
+    for lhs, rhs in product(target, target):
+        if lhs != rhs:
+            for binary_op in ["add", "sub", "mul", "div", "dot"]:
+                func = _gen_sddmm_func(lhs, rhs, binary_op)
+                setattr(sys.modules[__name__], func.__name__, func)
+                __all__.append(func.__name__)
+
+def copy_u(g, x):
+    r"""Generalized SDDMM function that copies source node features to edges.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        The input graph.
+    x : tensor
+        The source node features.
+
+    Returns
+    -------
+    tensor
+        The result tensor.
+
+    Notes
+    -----
+    This function supports autograd (computing input gradients given the output gradient).
+    """
+    return gsddmm(g, 'copy_lhs', x, None)
+
+def copy_v(g, x):
+    r"""Generalized SDDMM function that copies destination node features to edges.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        The input graph.
+    x : tensor
+        The destination node features.
+
+    Returns
+    -------
+    tensor
+        The result tensor.
+
+    Notes
+    -----
+    This function supports autograd (computing input gradients given the output gradient).
+    """
+    return gsddmm(g, 'copy_rhs', None, x)
+
+_register_sddmm_func()

python/dgl/ops/spmm.py

+"""dgl spmm operator module."""
+import sys
+from ..backend import gspmm
+
+__all__ = ['gspmm']
+
+
+def _gen_spmm_func(binary_op, reduce_op):
+    name = "u_{}_e_{}".format(binary_op, reduce_op)
+    docstring = """Generalized SpMM function.
+    It fuses two steps into one kernel.
+
+    1. Computes messages by {} source node and edge features.
+    2. Aggregate the messages by {} as the features on destination nodes.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        The input graph
+    x : tensor
+        The source node features.
+    y : tensor
+        The edge features.
+
+    Returns
+    -------
+    tensor
+        The result tensor.
+
+    Notes
+    -----
+    This function supports autograd (computing input gradients given the output gradient). If the
+    feature shape of two input operands do not match, we first broadcasts the features to a unified
+    shape (note that the memory usage will not increase accordingly) and then performs the operation.
+
+    Broadcasting follows NumPy semantics. Please see
+    https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
+    for more details about the NumPy broadcasting semantics.
+    """.format(binary_op, reduce_op)
+
+    def func(g, x, y):
+        return gspmm(g, binary_op, reduce_op, x, y)
+    func.__name__ = name
+    func.__doc__ = docstring
+    return func
+
+def _gen_copy_reduce_func(binary_op, reduce_op):
+
+    name = "{}_{}".format(binary_op, reduce_op)
+    binary_str = {
+        "copy_u": "It copies node feature to edge as the message.",
+        'copy_e': "It regards edge feature as message."
+    }
+    x_str = {
+        "copy_u": "source node",
+        "copy_e": "edge"
+    }
+    docstring = lambda binary_op: """Generalized SpMM function. {}
+    Then aggregates the message by {} on destination nodes.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        The input graph
+    x : tensor
+        The {} features.
+
+    Returns
+    -------
+    tensor
+        The result tensor.
+
+    Notes
+    -----
+    This function supports autograd (computing input gradients given the output gradient).
+    """.format(
+        binary_str[binary_op],
+        reduce_op,
+        x_str[binary_op])
+
+    def func(g, x):
+        if binary_op == 'copy_u':
+            return gspmm(g, 'copy_lhs', reduce_op, x, None)
+        else:
+            return gspmm(g, 'copy_rhs', reduce_op, None, x)
+
+    func.__name__ = name
+    func.__doc__ = docstring(binary_op)
+    return func
+
+def _register_spmm_func():
+    """Register spmm functions"""
+    for binary_op in ["add", "sub", "mul", "div", "copy_u", "copy_e"]:
+        for reduce_op in ["sum", "max", "min"]:
+            if binary_op.startswith("copy"):
+                func = _gen_copy_reduce_func(binary_op, reduce_op)
+            else:
+                func = _gen_spmm_func(binary_op, reduce_op)
+            setattr(sys.modules[__name__], func.__name__, func)
+            __all__.append(func.__name__)
+
+_register_spmm_func()