spmm.h 22.7 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
/*!
 *  Copyright (c) 2020 by Contributors
 * \file array/cpu/spmm.h
 * \brief SPMM CPU kernel function header.
 */
#ifndef DGL_ARRAY_CPU_SPMM_H_
#define DGL_ARRAY_CPU_SPMM_H_

#include <dgl/array.h>
#include <dgl/bcast.h>
11
#include <dgl/runtime/config.h>
12
#include <dgl/runtime/parallel_for.h>
13
#include <math.h>
14

15
#include <algorithm>
16
17
#include <limits>
#include <memory>
18
#include <vector>
19

20
21
#include "spmm_binary_ops.h"
#if !defined(_WIN32)
22
#ifdef USE_AVX
23
#include "intel/cpu_support.h"
24
25
26
#ifdef USE_LIBXSMM
#include "spmm_blocking_libxsmm.h"
#endif  // USE_LIBXSMM
27
28
#endif  // USE_AVX
#endif  // _WIN32
29
30
31
32
namespace dgl {
namespace aten {
namespace cpu {

33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#if !defined(_WIN32)
#ifdef USE_AVX
/*!
 * \brief CPU kernel of SpMM on Csr format using Xbyak.
 * \param cpu_spec JIT'ed kernel
 * \param bcast Broadcast information.
 * \param csr The Csr matrix.
 * \param X The feature on source nodes.
 * \param W The feature on edges.
 * \param O The result feature on destination nodes.
 * \note it uses node parallel strategy, different threads are responsible
 *       for the computation of different nodes. For each edge, it uses the
 *       JIT'ed kernel.
 */
template <typename IdType, typename DType, typename Op>
48
49
50
void SpMMSumCsrXbyak(
    dgl::ElemWiseAddUpdate<Op>* cpu_spec, const BcastOff& bcast,
    const CSRMatrix& csr, const DType* X, const DType* W, DType* O) {
51
52
53
54
55
  const bool has_idx = !IsNullArray(csr.data);
  const IdType* indptr = csr.indptr.Ptr<IdType>();
  const IdType* indices = csr.indices.Ptr<IdType>();
  const IdType* edges = csr.data.Ptr<IdType>();
  int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len, rhs_dim = bcast.rhs_len;
56
57
58
59
60
61
62
63
64
65

  runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
    for (auto rid = b; rid < e; ++rid) {
      const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
      DType* out_off = O + rid * dim;
      for (IdType j = row_start; j < row_end; ++j) {
        const IdType cid = indices[j];
        const IdType eid = has_idx ? edges[j] : j;
        cpu_spec->run(out_off, X + cid * lhs_dim, W + eid * rhs_dim, dim);
      }
66
    }
67
  });
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
}
#endif  // USE_AVX
#endif  // _WIN32

/*!
 * \brief Naive CPU kernel of SpMM on Csr format.
 * \param cpu_spec JIT'ed kernel
 * \param bcast Broadcast information.
 * \param csr The Csr matrix.
 * \param X The feature on source nodes.
 * \param W The feature on edges.
 * \param O The result feature on destination nodes.
 * \note it uses node parallel strategy, different threads are responsible
 *       for the computation of different nodes.
 */
template <typename IdType, typename DType, typename Op>
84
85
86
void SpMMSumCsrNaive(
    const BcastOff& bcast, const CSRMatrix& csr, const DType* X, const DType* W,
    DType* O) {
87
88
89
90
91
  const bool has_idx = !IsNullArray(csr.data);
  const IdType* indptr = csr.indptr.Ptr<IdType>();
  const IdType* indices = csr.indices.Ptr<IdType>();
  const IdType* edges = csr.data.Ptr<IdType>();
  int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len, rhs_dim = bcast.rhs_len;
92
93
94
95
96
97
98
99
100
101
102
  runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
    for (auto rid = b; rid < e; ++rid) {
      const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
      DType* out_off = O + rid * dim;
      for (IdType j = row_start; j < row_end; ++j) {
        const IdType cid = indices[j];
        const IdType eid = has_idx ? edges[j] : j;
        for (int64_t k = 0; k < dim; ++k) {
          const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
          const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
          const DType* lhs_off =
103
              Op::use_lhs ? X + cid * lhs_dim + lhs_add : nullptr;
104
          const DType* rhs_off =
105
              Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
106
107
          out_off[k] += Op::Call(lhs_off, rhs_off);
        }
108
109
      }
    }
110
  });
111
112
}

113
114
115
116
117
118
119
120
121
122
123
/*!
 * \brief CPU kernel of SpMM on Csr format.
 * \param bcast Broadcast information.
 * \param csr The Csr matrix.
 * \param ufeat The feature on source nodes.
 * \param efeat The feature on edges.
 * \param out The result feature on destination nodes.
 * \note it uses node parallel strategy, different threads are responsible
 *       for the computation of different nodes.
 */
template <typename IdType, typename DType, typename Op>
124
125
126
void SpMMSumCsr(
    const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
    NDArray out) {
127
128
129
130
131
132
  const bool has_idx = !IsNullArray(csr.data);
  const IdType* indptr = csr.indptr.Ptr<IdType>();
  const IdType* indices = csr.indices.Ptr<IdType>();
  const IdType* edges = csr.data.Ptr<IdType>();
  const DType* X = ufeat.Ptr<DType>();
  const DType* W = efeat.Ptr<DType>();
133
  int64_t dim = bcast.out_len;
134
  DType* O = out.Ptr<DType>();
135
136
137
138
139
140
141
  CHECK_NOTNULL(indptr);
  CHECK_NOTNULL(O);
  if (Op::use_lhs) {
    CHECK_NOTNULL(indices);
    CHECK_NOTNULL(X);
  }
  if (Op::use_rhs) {
142
    if (has_idx) CHECK_NOTNULL(edges);
143
144
    CHECK_NOTNULL(W);
  }
145
#if !defined(_WIN32)
146
#ifdef USE_AVX
147
#ifdef USE_LIBXSMM
148
149
150
  const bool no_libxsmm = bcast.use_bcast ||
                          std::is_same<DType, double>::value ||
                          !dgl::runtime::Config::Global()->IsLibxsmmAvailable();
151
152
  if (!no_libxsmm) {
    SpMMSumCsrLibxsmm<IdType, DType, Op>(bcast, csr, ufeat, efeat, out);
153
  } else {
154
155
156
157
158
159
160
#endif  // USE_LIBXSMM
    typedef dgl::ElemWiseAddUpdate<Op> ElemWiseUpd;
    /* Prepare an assembler kernel */
    static std::unique_ptr<ElemWiseUpd> asm_kernel_ptr(
        (dgl::IntelKernel<>::IsEnabled()) ? new ElemWiseUpd() : nullptr);
    /* Distribute the kernel among OMP threads */
    ElemWiseUpd* cpu_spec = (asm_kernel_ptr && asm_kernel_ptr->applicable())
161
162
                                ? asm_kernel_ptr.get()
                                : nullptr;
163
164
165
    if (cpu_spec && dim > 16 && !bcast.use_bcast) {
      SpMMSumCsrXbyak<IdType, DType, Op>(cpu_spec, bcast, csr, X, W, O);
    } else {
166
167
#endif  // USE_AVX
#endif  // _WIN32
168
      SpMMSumCsrNaive<IdType, DType, Op>(bcast, csr, X, W, O);
169
#if !defined(_WIN32)
170
#ifdef USE_AVX
171
172
    }
#ifdef USE_LIBXSMM
173
  }
174
#endif  // USE_LIBXSMM
175
176
#endif  // USE_AVX
#endif  // _WIN32
177
178
179
180
181
182
183
184
185
186
187
188
189
190
}

/*!
 * \brief CPU kernel of SpMM on Coo format.
 * \param bcast Broadcast information.
 * \param coo The Coo matrix.
 * \param ufeat The feature on source nodes.
 * \param efeat The feature on edges.
 * \param out The result feature on destination nodes.
 * \note it uses node parallel strategy, different threads are responsible
 *       for the computation of different nodes. To avoid possible data hazard,
 *       we use atomic operators in the reduction phase.
 */
template <typename IdType, typename DType, typename Op>
191
192
193
void SpMMSumCoo(
    const BcastOff& bcast, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
    NDArray out) {
194
195
196
197
198
199
  const bool has_idx = !IsNullArray(coo.data);
  const IdType* row = coo.row.Ptr<IdType>();
  const IdType* col = coo.col.Ptr<IdType>();
  const IdType* edges = coo.data.Ptr<IdType>();
  const DType* X = ufeat.Ptr<DType>();
  const DType* W = efeat.Ptr<DType>();
200
  int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len, rhs_dim = bcast.rhs_len;
201
202
203
204
205
206
207
208
209
  DType* O = out.Ptr<DType>();
  const int64_t nnz = coo.row->shape[0];
  // fill zero elements
  memset(O, 0, out.GetSize());
  // spmm
#pragma omp parallel for
  for (IdType i = 0; i < nnz; ++i) {
    const IdType rid = row[i];
    const IdType cid = col[i];
210
    const IdType eid = has_idx ? edges[i] : i;
211
212
213
214
    DType* out_off = O + cid * dim;
    for (int64_t k = 0; k < dim; ++k) {
      const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
      const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
215
      const DType* lhs_off =
216
          Op::use_lhs ? X + rid * lhs_dim + lhs_add : nullptr;
217
      const DType* rhs_off =
218
          Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
219
      const DType val = Op::Call(lhs_off, rhs_off);
220
      if (val != 0) {
221
#pragma omp atomic
222
223
        out_off[k] += val;
      }
224
225
226
227
228
229
230
231
232
233
234
    }
  }
}

/*!
 * \brief CPU kernel of SpMM-Min/Max on Csr format.
 * \param bcast Broadcast information.
 * \param csr The Csr matrix.
 * \param ufeat The feature on source nodes.
 * \param efeat The feature on edges.
 * \param out The result feature on destination nodes.
235
 * \param argu Arg-Min/Max on source nodes, which refers the source node indices
236
 *        correspond to the minimum/maximum values of reduction result on
237
 *        destination nodes. It's useful in computing gradients of Min/Max
238
239
240
 *        reducer.
 * \param arge Arg-Min/Max on edges. which refers the source node indices
          correspond to the minimum/maximum values of reduction result on
241
 *        destination nodes. It's useful in computing gradients of Min/Max
242
243
244
245
 *        reducer.
 * \note It uses node parallel strategy, different threads are responsible for
 *       the computation of different nodes.
 * \note The result will contain infinity for zero-degree nodes.
246
247
 */
template <typename IdType, typename DType, typename Op, typename Cmp>
248
249
250
void SpMMCmpCsr(
    const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
    NDArray out, NDArray argu, NDArray arge) {
251
252
253
  const bool has_idx = !IsNullArray(csr.data);
  const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices = static_cast<IdType*>(csr.indices->data);
254
  const IdType* edges =
255
      has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
256
257
258
  const DType* X = Op::use_lhs ? static_cast<DType*>(ufeat->data) : nullptr;
  const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
  const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
259
260
                rhs_dim = bcast.rhs_len;
  DType* O = static_cast<DType*>(out->data);
261
262
  IdType* argX = Op::use_lhs ? static_cast<IdType*>(argu->data) : nullptr;
  IdType* argW = Op::use_rhs ? static_cast<IdType*>(arge->data) : nullptr;
263
264
265
266
267
268
269
270
  CHECK_NOTNULL(indptr);
  CHECK_NOTNULL(O);
  if (Op::use_lhs) {
    CHECK_NOTNULL(indices);
    CHECK_NOTNULL(X);
    CHECK_NOTNULL(argX);
  }
  if (Op::use_rhs) {
271
    if (has_idx) CHECK_NOTNULL(edges);
272
273
274
275
276
277
278
    CHECK_NOTNULL(W);
    CHECK_NOTNULL(argW);
  }
#if !defined(_WIN32)
#ifdef USE_AVX
#ifdef USE_LIBXSMM

279
280
281
  const bool no_libxsmm = bcast.use_bcast ||
                          std::is_same<DType, double>::value ||
                          !dgl::runtime::Config::Global()->IsLibxsmmAvailable();
282
  if (!no_libxsmm) {
283
284
    SpMMCmpCsrLibxsmm<IdType, DType, Op, Cmp>(
        bcast, csr, ufeat, efeat, out, argu, arge);
285
286
287
288
289
  } else {
#endif  // USE_LIBXSMM
#endif  // USE_AVX
#endif  // _WIN32

290
291
292
293
294
295
296
297
298
299
300
301
302
    runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
      for (auto rid = b; rid < e; ++rid) {
        const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
        DType* out_off = O + rid * dim;
        IdType* argx_off = argX + rid * dim;
        IdType* argw_off = argW + rid * dim;
        for (IdType j = row_start; j < row_end; ++j) {
          const IdType cid = indices[j];
          const IdType eid = has_idx ? edges[j] : j;
          for (int64_t k = 0; k < dim; ++k) {
            const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
            const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
            const DType* lhs_off =
303
                Op::use_lhs ? X + cid * lhs_dim + lhs_add : nullptr;
304
            const DType* rhs_off =
305
                Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
306
307
308
309
310
311
312
            const DType val = Op::Call(lhs_off, rhs_off);
            if (Cmp::Call(out_off[k], val)) {
              out_off[k] = val;
              if (Op::use_lhs) argx_off[k] = cid;
              if (Op::use_rhs) argw_off[k] = eid;
            }
          }
313
314
        }
      }
315
    });
316
317
318
319
320
321
322
#if !defined(_WIN32)
#ifdef USE_AVX
#ifdef USE_LIBXSMM
  }
#endif  // USE_LIBXSMM
#endif  // USE_AVX
#endif  // _WIN32
323
324
}

325
326
327
328
329
330
331
332
333
334
/*!
 * \brief CPU kernel of SpMM-Min/Max on Csr format.
 * \param bcast Broadcast information.
 * \param csr The Csr matrix.
 * \param ufeat The feature on source nodes.
 * \param efeat The feature on edges.
 * \param out The result feature on destination nodes.
 * \param argu Arg-Min/Max on source nodes, which refers the source node indices
 *        correspond to the minimum/maximum values of reduction result on
 *        destination nodes. It's useful in computing gradients of Min/Max
335
 *        reducer.
336
337
 * \param arge Arg-Min/Max on edges. which refers the source node indices
 *        correspond to the minimum/maximum values of reduction result on
338
 *        destination nodes. It's useful in computing gradients of Min/Max
339
 *        reducer.
340
341
342
343
344
345
346
347
 * \param argu_ntype Node type of the arg-Min/Max on source nodes, which refers
 *        the source node types correspond to the minimum/maximum values of
 *        reduction result on destination nodes. It's useful in computing
 *        gradients of Min/Max reducer.
 * \param arge_etype Edge-type of the arg-Min/Max on edges. which refers the
 *        source node indices correspond to the minimum/maximum values of
 *        reduction result on destination nodes. It's useful in computing
 *        gradients of Min/Max reducer.
348
349
 * \param src_type Node type of the source nodes of an etype
 * \param etype Edge type
350
351
 */
template <typename IdType, typename DType, typename Op, typename Cmp>
352
353
354
355
void SpMMCmpCsrHetero(
    const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
    NDArray out, NDArray argu, NDArray arge, NDArray argu_ntype,
    NDArray arge_etype, const int ntype, const int etype) {
356
357
358
359
  const bool has_idx = !IsNullArray(csr.data);
  const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices = static_cast<IdType*>(csr.indices->data);
  const IdType* edges =
360
      has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
361
362
363
364
365
366
367
  const DType* X = Op::use_lhs ? static_cast<DType*>(ufeat->data) : nullptr;
  const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
  const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
                rhs_dim = bcast.rhs_len;
  DType* O = static_cast<DType*>(out->data);
  IdType* argX = Op::use_lhs ? static_cast<IdType*>(argu->data) : nullptr;
  IdType* argW = Op::use_rhs ? static_cast<IdType*>(arge->data) : nullptr;
368
369
370
371
  IdType* argX_ntype =
      Op::use_lhs ? static_cast<IdType*>(argu_ntype->data) : nullptr;
  IdType* argW_etype =
      Op::use_rhs ? static_cast<IdType*>(arge_etype->data) : nullptr;
372
373
374
375
376
377
378
379
  CHECK_NOTNULL(indptr);
  CHECK_NOTNULL(O);
  if (Op::use_lhs) {
    CHECK_NOTNULL(indices);
    CHECK_NOTNULL(X);
    CHECK_NOTNULL(argX);
  }
  if (Op::use_rhs) {
380
    if (has_idx) CHECK_NOTNULL(edges);
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
    CHECK_NOTNULL(W);
    CHECK_NOTNULL(argW);
  }
  // TODO(Israt): Use LIBXSMM. Homogeneous graph uses LIBXMM when enabled.
  runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
    for (auto rid = b; rid < e; ++rid) {
      const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
      DType* out_off = O + rid * dim;
      IdType* argx_off = argX + rid * dim;
      IdType* argw_off = argW + rid * dim;
      IdType* argx_ntype = argX_ntype + rid * dim;
      IdType* argw_etype = argW_etype + rid * dim;
      for (IdType j = row_start; j < row_end; ++j) {
        const IdType cid = indices[j];
        const IdType eid = has_idx ? edges[j] : j;
        for (int64_t k = 0; k < dim; ++k) {
          const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
          const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
          const DType* lhs_off =
400
              Op::use_lhs ? X + cid * lhs_dim + lhs_add : nullptr;
401
          const DType* rhs_off =
402
              Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
          const DType val = Op::Call(lhs_off, rhs_off);
          if (Cmp::Call(out_off[k], val)) {
            out_off[k] = val;
            if (Op::use_lhs) {
              argx_off[k] = cid;
              argx_ntype[k] = ntype;
            }
            if (Op::use_rhs) {
              argw_off[k] = eid;
              argw_etype[k] = etype;
            }
          }
        }
      }
    }
  });
}

421
422
423
424
425
426
427
/*!
 * \brief CPU kernel of SpMM-Min/Max on Coo format.
 * \param bcast Broadcast information.
 * \param coo The Coo matrix.
 * \param ufeat The feature on source nodes.
 * \param efeat The feature on edges.
 * \param out The result feature on destination nodes.
428
 * \param argu Arg-Min/Max on source nodes, which refers the source node indices
429
 *        correspond to the minimum/maximum values of reduction result on
430
 *        destination nodes. It's useful in computing gradients of Min/Max
431
432
433
 *        reducer.
 * \param arge Arg-Min/Max on edges. which refers the source node indices
 *        correspond to the minimum/maximum values of reduction result on
434
 *        destination nodes. It's useful in computing gradients of Min/Max
435
436
437
438
439
 *        reducer.
 * \note it uses node parallel strategy, different threads are responsible for
 *       the computation of different nodes. To avoid possible data hazard, we
 *       use atomic operators in the reduction phase.
 * \note The result will contain infinity for zero-degree nodes.
440
441
 */
template <typename IdType, typename DType, typename Op, typename Cmp>
442
443
444
void SpMMCmpCoo(
    const BcastOff& bcast, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
    NDArray out, NDArray argu, NDArray arge) {
445
446
447
  const bool has_idx = !IsNullArray(coo.data);
  const IdType* row = static_cast<IdType*>(coo.row->data);
  const IdType* col = static_cast<IdType*>(coo.col->data);
448
  const IdType* edges =
449
      has_idx ? static_cast<IdType*>(coo.data->data) : nullptr;
450
451
452
  const DType* X = Op::use_lhs ? static_cast<DType*>(ufeat->data) : nullptr;
  const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
  const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
453
454
                rhs_dim = bcast.rhs_len;
  DType* O = static_cast<DType*>(out->data);
455
456
  IdType* argX = Op::use_lhs ? static_cast<IdType*>(argu->data) : nullptr;
  IdType* argW = Op::use_rhs ? static_cast<IdType*>(arge->data) : nullptr;
457
458
459
460
461
462
463
464
  const int64_t nnz = coo.row->shape[0];
  // fill zero elements
  std::fill(O, O + out.NumElements(), Cmp::zero);
  // spmm
#pragma omp parallel for
  for (IdType i = 0; i < nnz; ++i) {
    const IdType rid = row[i];
    const IdType cid = col[i];
465
    const IdType eid = has_idx ? edges[i] : i;
466
    DType* out_off = O + cid * dim;
467
468
    IdType* argx_off = Op::use_lhs ? argX + cid * dim : nullptr;
    IdType* argw_off = Op::use_rhs ? argW + cid * dim : nullptr;
469
470
471
    for (int64_t k = 0; k < dim; ++k) {
      const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
      const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
472
      const DType* lhs_off =
473
          Op::use_lhs ? X + rid * lhs_dim + lhs_add : nullptr;
474
      const DType* rhs_off =
475
          Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
476
477
478
479
      const DType val = Op::Call(lhs_off, rhs_off);
#pragma omp critical
      if (Cmp::Call(out_off[k], val)) {
        out_off[k] = val;
480
481
        if (Op::use_lhs) argx_off[k] = rid;
        if (Op::use_rhs) argw_off[k] = eid;
482
483
484
485
486
      }
    }
  }
}

487
488
489
490
491
492
493
494
495
/*!
 * \brief CPU kernel of Edge_softmax_csr_forward on Csr format.
 * \param bcast Broadcast information.
 * \param csr The Csr matrix.
 * \param ufeat The feature on source nodes.
 * \param efeat The feature on edges.
 * \param out The result of edge_softmax_forward.
 */
template <typename IdType, typename DType, typename Op>
496
497
498
void Edge_softmax_csr_forward(
    const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
    NDArray out) {
499
500
501
  const bool has_idx = !IsNullArray(csr.data);
  const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
  const IdType* edges =
502
      has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
503
504
505
506
507
  const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
  const int64_t dim = bcast.out_len, rhs_dim = bcast.rhs_len;
  runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
    for (auto rid = b; rid < e; ++rid) {
      const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
508
509
      std::vector<DType> data_e(row_end - row_start, 0);
      std::vector<IdType> num(row_end - row_start, 0);
510
511
512
513
514
515
      for (int64_t k = 0; k < dim; ++k) {
        DType max_v = -std::numeric_limits<DType>::infinity();
        for (IdType j = row_start; j < row_end; ++j) {
          const IdType eid = has_idx ? edges[j] : j;
          const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
          const DType* rhs_off =
516
517
518
              Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
          data_e[j - row_start] = *rhs_off;
          num[j - row_start] = eid * rhs_dim + rhs_add;
519
520
521
522
523
524
525
526
          max_v = std::max<DType>(max_v, (*rhs_off));
        }
        DType exp_sum = 0;
        for (auto& element : data_e) {
          element -= max_v;
          element = std::exp(element);
          exp_sum += element;
        }
527
528
        for (int i = 0; i < row_end - row_start; i++) {
          out.Ptr<DType>()[num[i]] = data_e[i] / exp_sum;
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
        }
      }
    }
  });
}

/*!
 * \brief CPU kernel of Edge_softmax_csr_backward on Csr format.
 * \param bcast Broadcast information.
 * \param csr The Csr matrix.
 * \param out The result of forward.
 * \param sds The result of gradiet * out.
 * \param back_out The result of edge_softmax_backward.
 */
template <typename IdType, typename DType, typename Op>
544
545
546
void Edge_softmax_csr_backward(
    const BcastOff& bcast, const CSRMatrix& csr, NDArray out, NDArray sds,
    NDArray back_out) {
547
548
549
  const bool has_idx = !IsNullArray(csr.data);
  const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
  const IdType* edges =
550
      has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
551
552
553
554
555
556
557
558
559
560
561
562
  const DType* W_out = Op::use_rhs ? static_cast<DType*>(out->data) : nullptr;
  const DType* W_sds = Op::use_rhs ? static_cast<DType*>(sds->data) : nullptr;
  const int64_t dim = bcast.out_len, rhs_dim = bcast.rhs_len;
  runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
    for (auto rid = b; rid < e; ++rid) {
      const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
      for (int64_t k = 0; k < dim; ++k) {
        DType sum_sds = 0;
        for (IdType j = row_start; j < row_end; ++j) {
          const IdType eid = has_idx ? edges[j] : j;
          const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
          const DType* rhs_off_sds =
563
              Op::use_rhs ? W_sds + eid * rhs_dim + rhs_add : nullptr;
564
565
          sum_sds += (*rhs_off_sds);
        }
566
        for (IdType j = row_start; j < row_end; ++j) {
567
568
569
          const IdType eid = has_idx ? edges[j] : j;
          const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
          const DType* rhs_off_out =
570
              Op::use_rhs ? W_out + eid * rhs_dim + rhs_add : nullptr;
571
          const DType* rhs_off_sds =
572
573
574
              Op::use_rhs ? W_sds + eid * rhs_dim + rhs_add : nullptr;
          back_out.Ptr<DType>()[eid * rhs_dim + rhs_add] =
              (*rhs_off_sds) - sum_sds * (*rhs_off_out);
575
576
577
578
579
580
        }
      }
    }
  });
}

581
582
583
584
585
}  // namespace cpu
}  // namespace aten
}  // namespace dgl

#endif  // DGL_ARRAY_CPU_SPMM_H_