ScaledGemmTunableOp Class — pytorch Architecture

Source Code

aten/src/ATen/cuda/tunable/TunableGemm.h lines 306–327

class ScaledGemmTunableOp : public TunableOp<ScaledGemmParams<CT>> {
 public:
  ScaledGemmTunableOp() {
    this->RegisterOp(std::string("Default"), std::make_unique<DefaultScaledGemmOp<CT>>());

#ifdef USE_ROCM
    for (auto&& [name, op] : GetHipBlasLtScaledGemmTypeStringAndOps<AT, BT, CT, ALayout, BLayout>()) {
      this->RegisterOp(std::move(name), std::move(op));
    }
#endif

    this->RegisterOp(std::string("Default"), std::make_unique<DefaultScaledGemmOp<CT>>());
  }

  std::string Signature() override {
    return fmt::sprintf("ScaledGemmTunableOp_%s_%s_%s_%c%c",
      TypeName<AT>(AT{}),
      TypeName<BT>(BT{}),
      TypeName<CT>(CT{}),
      BlasOpToString(ALayout), BlasOpToString(BLayout));
  }
};