[Feature] x_dot_x builtin kernel support

docs/source/api/python/function.rst

@@ -40,6 +40,12 @@ Message functions
     e_sub_v
     e_mul_v
     e_div_v
+    u_dot_v
+    u_dot_e
+    v_dot_e
+    v_dot_u
+    e_dot_u
+    e_dot_v
 
 Reduce functions
 ----------------

docs/source/features/builtin.rst

@@ -98,6 +98,10 @@ Here is a cheatsheet of all the DGL builtins.
 |                         +----------------------------------------------------+-----------------------+
 |                         | ``e_add_v``, ``e_sub_v``, ``e_mul_v``, ``e_div_v`` |                       |
 |                         +----------------------------------------------------+-----------------------+
+|                         | ``u_dot_v``, ``u_dot_e``, ``v_dot_e``              |                       |
+|                         +----------------------------------------------------+-----------------------+
+|                         | ``v_dot_u``, ``e_dot_u``, ``e_dot_v``              |                       |
+|                         +----------------------------------------------------+-----------------------+
 |                         | ``src_mul_edge``                                   |  alias of ``u_mul_e`` |
 +-------------------------+----------------------------------------------------+-----------------------+
 | Reduce function         | ``max``                                            |                       |

python/dgl/backend/mxnet/tensor.py

@@ -376,8 +376,11 @@ def __init__(self, reducer, binary_op, graph, lhs, rhs, out_size, lhs_map,
     def forward(self, lhs_data, rhs_data):
         lhs_data_nd = zerocopy_to_dgl_ndarray(lhs_data)
         rhs_data_nd = zerocopy_to_dgl_ndarray(rhs_data)
-        feat_shape = K.infer_binary_feature_shape(lhs_data_nd, rhs_data_nd)
-        out_data = nd.empty((self.out_size,) + feat_shape,
+        feat_shape = K.infer_binary_feature_shape(self.binary_op, lhs_data_nd, rhs_data_nd)
+        out_shape = feat_shape
+        if self.binary_op == 'dot':
+            out_shape = feat_shape[:-1]
+        out_data = nd.empty((self.out_size,) + out_shape,
                             ctx=lhs_data.context, dtype=lhs_data.dtype)
         out_data_nd = zerocopy_to_dgl_ndarray_for_write(out_data)
         K.binary_op_reduce(
@@ -402,10 +405,10 @@ def forward(self, lhs_data, rhs_data):
             in_ones = nd.ones((n,), ctx=lhs_data.context, dtype=lhs_data.dtype)
             in_ones_nd = zerocopy_to_dgl_ndarray(in_ones)
             K.copy_reduce(
-                'sum', self.graph, target, in_ones_nd, degs_nd, 
+                'sum', self.graph, target, in_ones_nd, degs_nd,
                 in_map, self.out_map[0])
             # reshape
-            degs = degs.reshape((out_data.shape[0],) + (1,) * (out_data.ndim - 1)).clip(1, float('inf')) 
+            degs = degs.reshape((out_data.shape[0],) + (1,) * (out_data.ndim - 1)).clip(1, float('inf'))
             out_data = out_data / degs
         else:
             degs = None

python/dgl/backend/pytorch/tensor.py

@@ -288,11 +288,14 @@ def forward(ctx, reducer, binary_op, graph, lhs, rhs, lhs_data, rhs_data,
                 out_size, lhs_map, rhs_map, out_map):
         lhs_data_nd = zerocopy_to_dgl_ndarray(lhs_data)
         rhs_data_nd = zerocopy_to_dgl_ndarray(rhs_data)
-        feat_shape = K.infer_binary_feature_shape(lhs_data_nd, rhs_data_nd)
-        out_data = lhs_data.new_empty((out_size,) + feat_shape)
+        feat_shape = K.infer_binary_feature_shape(binary_op, lhs_data_nd, rhs_data_nd)
+        out_shape = feat_shape
+        if binary_op == 'dot':
+            out_shape = feat_shape[:-1]
+        out_data = lhs_data.new_empty((out_size,) + out_shape)
         out_data_nd = zerocopy_to_dgl_ndarray(out_data)
         K.binary_op_reduce(
-            reducer if reducer != 'mean' else 'sum', 
+            reducer if reducer != 'mean' else 'sum',
             binary_op, graph, lhs, rhs, lhs_data_nd, rhs_data_nd,
             out_data_nd, lhs_map[0], rhs_map[0], out_map[0])
         # normalize if mean reducer
@@ -311,7 +314,7 @@ def forward(ctx, reducer, binary_op, graph, lhs, rhs, lhs_data, rhs_data,
             in_ones = lhs_data.new_ones((n,))
             in_ones_nd = zerocopy_to_dgl_ndarray(in_ones)
             K.copy_reduce(
-                'sum', graph, target, in_ones_nd, degs_nd, in_map, out_map[0]) 
+                'sum', graph, target, in_ones_nd, degs_nd, in_map, out_map[0])
             # reshape
             degs = degs.reshape((out_data.shape[0],) + (1,) * (out_data.dim() - 1)).clamp(min=1)
             out_data = out_data / degs

python/dgl/function/message.py

@@ -147,12 +147,17 @@ def copy_e(e, out):
 
 ###############################################################################
 # Generate all following  builtin message functions:
+# element-wise message functions:
 # u_add_v, u_sub_v, u_mul_v, u_div_v
 # u_add_e, u_sub_e, u_mul_e, u_div_e
 # v_add_u, v_sub_u, v_mul_u, v_div_u
 # v_add_e, v_sub_e, v_mul_e, v_div_e
 # e_add_u, e_sub_u, e_mul_u, e_div_u
 # e_add_v, e_sub_v, e_mul_v, e_div_v
+#
+# dot message functions:
+# u_dot_v, u_dot_e, v_dot_e
+# v_dot_u, e_dot_u, e_dot_v
 
 _TARGET_MAP = {
     "u": TargetCode.SRC,
@@ -200,12 +205,11 @@ def _register_builtin_message_func():
     target = ["u", "v", "e"]
     for lhs, rhs in product(target, target):
         if lhs != rhs:
-            for binary_op in ["add", "sub", "mul", "div"]:
+            for binary_op in ["add", "sub", "mul", "div", "dot"]:
                 func = _gen_message_builtin(lhs, rhs, binary_op)
                 setattr(sys.modules[__name__], func.__name__, func)
                 __all__.append(func.__name__)
 
-
 _register_builtin_message_func()
 
 

python/dgl/kernel.py

@@ -4,11 +4,14 @@
 from ._ffi.function import _init_api
 from .ndarray import empty
 
-def infer_binary_feature_shape(lhs, rhs):
+# pylint: disable=invalid-name
+def infer_binary_feature_shape(op, lhs, rhs):
     """Infer the output feature shape after a binary operation between lhs and rhs.
 
     Parameter
     ---------
+    op : string
+        The binary_op name.
     lhs : dgl.ndarray.NDArray
         The lhs tensor.
     rhs : dgl.ndarray.NDArray
@@ -19,7 +22,7 @@ def infer_binary_feature_shape(lhs, rhs):
     tuple of int
         The output feature shape.
     """
-    ret = _CAPI_DGLKernelInferBinaryFeatureShape(lhs, rhs)
+    ret = _CAPI_DGLKernelInferBinaryFeatureShape(op, lhs, rhs)
     return tuple(ret.asnumpy())
 
 # pylint: disable=invalid-name

src/kernel/binary_reduce.cc

@@ -93,11 +93,25 @@ bool HasBcast(NDArray lhs, NDArray rhs) {
 //    e.g. (4, 1, 3, 3) and (4, 5, 3, 3) become (4, 1, 9) and (4, 5, 9)
 //
 // See also: BcastInfo (kernel/binary_reduce.h)
-BcastInfo CalcBcastInfo(NDArray lhs, NDArray rhs) {
+BcastInfo CalcBcastInfo(const std::string& op, NDArray lhs, NDArray rhs) {
   BcastInfo ret;
   const int max_ndim = std::max(lhs->ndim, rhs->ndim) - 1;
   int64_t accum = 0;
-  for (int j = 0; j < max_ndim; ++j) {
+  int j = 0;
+  // for dot operation: vector [dot] vector
+  // lhs_shape[ndim-1] == rhs_shape[ndim-1] = sizeof(vector)
+  // out_shape[ndim-1] = 1
+  if (op == binary_op::kDot) {
+    // get size of vector
+    ret.data_len = lhs->shape[lhs->ndim - 1];
+    // skip vector size dim
+    ++j;
+    ret.real_out_shape.push_back(ret.data_len);
+  } else {  // op != binary_op::kDot
+    ret.data_len = 1;
+  }
+
+  for (; j < max_ndim; ++j) {
     const int dl = (lhs->ndim - 1 - j < 1)? 1 : lhs->shape[lhs->ndim - 1 - j];
     const int dr = (rhs->ndim - 1 - j < 1)? 1 : rhs->shape[rhs->ndim - 1 - j];
     if (dl != dr) {
@@ -258,16 +272,18 @@ class BipartiteCSRWrapper : public CSRWrapper {
 
 
 std::vector<int64_t> InferBinaryFeatureShape(
+    const std::string& op,
     NDArray lhs,
     NDArray rhs) {
-  return CalcBcastInfo(lhs, rhs).real_out_shape;
+  return CalcBcastInfo(op, lhs, rhs).real_out_shape;
 }
 
 DGL_REGISTER_GLOBAL("kernel._CAPI_DGLKernelInferBinaryFeatureShape")
 .set_body([] (DGLArgs args, DGLRetValue* rv) {
-    NDArray lhs = args[0];
-    NDArray rhs = args[1];
-    const auto& shape = InferBinaryFeatureShape(lhs, rhs);
+    std::string op = args[0];
+    NDArray lhs = args[1];
+    NDArray rhs = args[2];
+    const auto& shape = InferBinaryFeatureShape(op, lhs, rhs);
     const int64_t len = shape.size();
     NDArray ret = NDArray::Empty(
         {len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
@@ -300,7 +316,7 @@ void BinaryOpReduce(
         rhs_mapping, lhs_mapping, out_mapping);
   } else {
     if (HasBcast(lhs_data, rhs_data)) {
-      BcastInfo info = CalcBcastInfo(lhs_data, rhs_data);
+      BcastInfo info = CalcBcastInfo(op, lhs_data, rhs_data);
       DGL_XPU_SWITCH(ctx.device_type, BinaryReduceBcastImpl,
           info, reducer, op, graph,
           lhs, rhs,
@@ -394,7 +410,7 @@ void BackwardLhsBinaryOpReduce(
         grad_out_data, grad_lhs_data);
   } else {
     if (HasBcast(lhs_data, rhs_data)) {
-      BcastInfo info = CalcBcastInfo(lhs_data, rhs_data);
+      BcastInfo info = CalcBcastInfo(op, lhs_data, rhs_data);
       DGL_XPU_SWITCH(ctx.device_type, BackwardBinaryReduceBcastImpl,
           info, reducer, op, graph,
           lhs, rhs,
@@ -468,7 +484,7 @@ void BackwardRhsBinaryOpReduce(
         grad_out_data, grad_rhs_data);
   } else {
     if (HasBcast(lhs_data, rhs_data)) {
-      BcastInfo info = CalcBcastInfo(lhs_data, rhs_data);
+      BcastInfo info = CalcBcastInfo(op, lhs_data, rhs_data);
       DGL_XPU_SWITCH(ctx.device_type, BackwardBinaryReduceBcastImpl,
           info, reducer, op, graph,
           lhs, rhs,

src/kernel/binary_reduce.h

@@ -30,6 +30,8 @@ struct BcastInfo {
   std::vector<int64_t> lhs_shape, lhs_stride;
   std::vector<int64_t> rhs_shape, rhs_stride;
   std::vector<int64_t> out_shape, out_stride;
+
+  int64_t data_len;
 };
 
 /*

src/kernel/binary_reduce_common.h

@@ -29,6 +29,7 @@ static const char kAdd[] = "add";
 static const char kSub[] = "sub";
 static const char kMul[] = "mul";
 static const char kDiv[] = "div";
+static const char kDot[] = "dot";
 static const char kUseLhs[] = "use_lhs";
 
 /*!
@@ -129,8 +130,8 @@ struct SwitchSrcDst<SelectDst> {
 // common binary functors
 template <typename DType>
 struct BinaryAdd {
-  static DGLDEVICE DGLINLINE DType Call(DType lhs, DType rhs) {
-    return lhs + rhs;
+  static DGLDEVICE DGLINLINE DType Call(DType *lhs, DType *rhs, int64_t len) {
+    return lhs[0] + rhs[0];
   }
   static DGLDEVICE DGLINLINE DType BackwardLhs(DType lhs, DType rhs, DType out) {
     return 1;
@@ -142,8 +143,8 @@ struct BinaryAdd {
 
 template <typename DType>
 struct BinaryMul {
-  static DGLDEVICE DGLINLINE DType Call(DType lhs, DType rhs) {
-    return lhs * rhs;
+  static DGLDEVICE DGLINLINE DType Call(DType *lhs, DType *rhs, int64_t len) {
+    return lhs[0] * rhs[0];
   }
   static DGLDEVICE DGLINLINE DType BackwardLhs(DType lhs, DType rhs, DType out) {
     return rhs;
@@ -155,8 +156,8 @@ struct BinaryMul {
 
 template <typename DType>
 struct BinarySub {
-  static DGLDEVICE DGLINLINE DType Call(DType lhs, DType rhs) {
-    return lhs - rhs;
+  static DGLDEVICE DGLINLINE DType Call(DType *lhs, DType *rhs, int64_t len) {
+    return lhs[0] - rhs[0];
   }
   static DGLDEVICE DGLINLINE DType BackwardLhs(DType lhs, DType rhs, DType out) {
     return 1;
@@ -168,8 +169,8 @@ struct BinarySub {
 
 template <typename DType>
 struct BinaryDiv {
-  static DGLDEVICE DGLINLINE DType Call(DType lhs, DType rhs) {
-    return lhs / rhs;
+  static DGLDEVICE DGLINLINE DType Call(DType *lhs, DType *rhs, int64_t len) {
+    return lhs[0] / rhs[0];
   }
   static DGLDEVICE DGLINLINE DType BackwardLhs(DType lhs, DType rhs, DType out) {
     return static_cast<DType>(1) / rhs;
@@ -181,8 +182,8 @@ struct BinaryDiv {
 
 template <typename DType>
 struct BinaryUseLhs {
-  static DGLDEVICE DGLINLINE DType Call(DType lhs, DType rhs) {
-    return lhs;
+  static DGLDEVICE DGLINLINE DType Call(DType *lhs, DType *rhs, int64_t len) {
+    return lhs[0];
   }
   static DGLDEVICE DGLINLINE DType BackwardLhs(DType lhs, DType rhs, DType out) {
     return 1;
@@ -192,6 +193,25 @@ struct BinaryUseLhs {
   }
 };
 
+template <typename DType>
+struct BinaryDot {
+  static DGLDEVICE DGLINLINE DType Call(DType *lhs, DType *rhs, int64_t len) {
+    DType out = 0;
+    // simple vector dot vector
+#pragma unroll
+    for (int i = 0; i < len; i ++)
+      out += lhs[i] * rhs[i];
+
+    return out;
+  }
+  static DGLDEVICE DGLINLINE DType BackwardLhs(DType lhs, DType rhs, DType out) {
+    return rhs;
+  }
+  static DGLDEVICE DGLINLINE DType BackwardRhs(DType lhs, DType rhs, DType out) {
+    return lhs;
+  }
+};
+
 // Macro for dispatching op enum code and target code into template arguments.
 // The macro dispatches following combinations:
 //  - Add(Src, Dst), Add(Src, Edge), Add(Dst, Edge)
@@ -201,6 +221,8 @@ struct BinaryUseLhs {
 //  - Div(Src, Dst), Div(Src, Edge), Div(Dst, Edge)
 //    Div(Dst, Src), Div(Edge, Src), Div(Edge, Dst)
 //  - UseLhs(Src, None), UseLhs(Edge, None)
+//  - Dot(Src, Dst), Dot(Src, Edge), Dot(Dst, Edge)
+//  - Dot(Dst, Src), Dot(Edge, Src), Dot(Edge, Dst)
 // Note that for commutative operators (e.g. Add and Mul), we only generate
 // kernels for lhs code smaller than rhs code.
 #define OP_TARGET_SWITCH(op, lhs, rhs, DType, OpType, LeftType, RightType, ...)   \
@@ -306,6 +328,36 @@ struct BinaryUseLhs {
     typedef SelectEdge LeftType;                               \
     typedef SelectNone RightType;                              \
     {__VA_ARGS__}                                              \
+  } else if (op == kDot && lhs == kSrc && rhs == kDst) {       \
+    typedef BinaryDot<DType> OpType;                           \
+    typedef SelectSrc LeftType;                                \
+    typedef SelectDst RightType;                               \
+    {__VA_ARGS__}                                              \
+  } else if (op == kDot && lhs == kSrc && rhs == kEdge) {      \
+    typedef BinaryDot<DType> OpType;                           \
+    typedef SelectSrc LeftType;                                \
+    typedef SelectEdge RightType;                              \
+    {__VA_ARGS__}                                              \
+  } else if (op == kDot && lhs == kDst && rhs == kEdge) {      \
+    typedef BinaryDot<DType> OpType;                           \
+    typedef SelectDst LeftType;                                \
+    typedef SelectEdge RightType;                              \
+    {__VA_ARGS__}                                              \
+  } else if (op == kDot && lhs == kDst && rhs == kSrc) {       \
+    typedef BinaryDot<DType> OpType;                           \
+    typedef SelectDst LeftType;                                \
+    typedef SelectSrc RightType;                               \
+    {__VA_ARGS__}                                              \
+  } else if (op == kDot && lhs == kEdge && rhs == kSrc) {      \
+    typedef BinaryDot<DType> OpType;                           \
+    typedef SelectEdge LeftType;                               \
+    typedef SelectSrc RightType;                               \
+    {__VA_ARGS__}                                              \
+  } else if (op == kDot && lhs == kEdge && rhs == kDst) {      \
+    typedef BinaryDot<DType> OpType;                           \
+    typedef SelectEdge LeftType;                               \
+    typedef SelectDst RightType;                               \
+    {__VA_ARGS__}                                              \
   } else {                                                     \
     LOG(FATAL) << "Unsupported operation: op=" << op           \
       << " lhs=" << lhs << " rhs=" << rhs;                     \
@@ -333,7 +385,13 @@ struct BinaryUseLhs {
   MSVC_EXPAND(GEN(__VA_ARGS__, SelectDst, SelectEdge, BinaryDiv))     \
   MSVC_EXPAND(GEN(__VA_ARGS__, SelectEdge, SelectDst, BinaryDiv))     \
   MSVC_EXPAND(GEN(__VA_ARGS__, SelectSrc, SelectNone, BinaryUseLhs))  \
-  MSVC_EXPAND(GEN(__VA_ARGS__, SelectEdge, SelectNone, BinaryUseLhs))
+  MSVC_EXPAND(GEN(__VA_ARGS__, SelectEdge, SelectNone, BinaryUseLhs)) \
+  MSVC_EXPAND(GEN(__VA_ARGS__, SelectSrc, SelectDst, BinaryDot))      \
+  MSVC_EXPAND(GEN(__VA_ARGS__, SelectSrc, SelectEdge, BinaryDot))     \
+  MSVC_EXPAND(GEN(__VA_ARGS__, SelectDst, SelectEdge, BinaryDot))     \
+  MSVC_EXPAND(GEN(__VA_ARGS__, SelectDst, SelectSrc, BinaryDot))      \
+  MSVC_EXPAND(GEN(__VA_ARGS__, SelectEdge, SelectSrc, BinaryDot))     \
+  MSVC_EXPAND(GEN(__VA_ARGS__, SelectEdge, SelectDst, BinaryDot))
 
 //////////////////////////////////////////////////////////////////////////
 // Defines reducer category. Each category is an empty structure.