call Class — pytorch Architecture

Source Code

aten/src/ATen/native/CPUBlas.cpp lines 1120–1170

  template <typename scalar_t_a, typename scalar_t_b, typename scalar_t_c>
  static void call(
      int64_t M,
      int64_t N,
      int64_t K,
      int64_t ld_a,
      int64_t ld_b,
      int64_t ld_c,
      const bool add_C,
      const scalar_t_a* A,
      const scalar_t_b* B,
      scalar_t_c* C) {
    auto&& key = BrgemmKey(
        M,
        N,
        K,
        1,
        ld_a,
        ld_b,
        ld_c,
        c10::CppTypeToScalarType<scalar_t_a>::value,
        c10::CppTypeToScalarType<scalar_t_b>::value,
        c10::CppTypeToScalarType<scalar_t_c>::value,
        add_C);
    // Fetch/create GemmHelper object
    auto&& value = fetch_or_create(key, [&]() {
      auto&& v = std::make_shared<GemmHelper>(
          M,
          N,
          K,
          1,
          ld_a,
          ld_b,
          ld_c,
          c10::CppTypeToScalarType<scalar_t_a>::value,
          c10::CppTypeToScalarType<scalar_t_b>::value,
          c10::CppTypeToScalarType<scalar_t_c>::value,
          add_C);
      (*v).brg.generate();
      return std::move(v);
    });
    if (get_current() != value) {
#if defined(ONEDNN_UKERNEL_1)
      dnnl::ukernel::brgemm::release_hw_context();
#endif
      ((*value).brg).set_hw_context();
      get_current() = value;
    }
    ((*value).brg)
        .execute(A, B, (*value).A_B_offsets, C, (*value).scratchpad.data());
  }