| 283 | // MatrixMul |
| 284 | template <> |
| 285 | uint64_t opr_footprint_func<opr::MatrixMul>(cg::OperatorNodeBase* opr) { |
| 286 | auto&& mopr = opr->cast_final_safe<opr::MatrixMul>(); |
| 287 | auto &&i0 = opr->input(0)->shape(), &&i1 = opr->input(1)->shape(); |
| 288 | mgb_assert(i0.ndim == 2 && i1.ndim == 2); |
| 289 | auto m = i0[0], k0 = i0[1], k1 = i1[0], n = i1[1]; |
| 290 | if (mopr.param().transposeA) { |
| 291 | std::swap(m, k0); |
| 292 | } |
| 293 | if (mopr.param().transposeB) { |
| 294 | std::swap(k1, n); |
| 295 | } |
| 296 | mgb_assert(k0 == k1); |
| 297 | // mul and add are counted as 2 operations |
| 298 | return m * k0 * n * 2; |
| 299 | } |
| 300 | |
| 301 | template <> |
| 302 | uint64_t opr_footprint_func<opr::LocalShareForward>(cg::OperatorNodeBase* opr) { |