mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
310eb6d8c6
Summary: Fixes https://github.com/pytorch/pytorch/issues/130311. We need to guard CUDA-only code in test_aoti_inference with macros so that it won't fail for CPU-only platform. Pull Request resolved: https://github.com/pytorch/pytorch/pull/134675 Approved by: https://github.com/atalman, https://github.com/chunyuan-w
51 lines
1.6 KiB
C++
51 lines
1.6 KiB
C++
#include <stdexcept>
|
|
|
|
#include <torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h>
|
|
#if defined(USE_CUDA) || defined(USE_ROCM)
|
|
#include <torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h>
|
|
#endif
|
|
|
|
#include "aoti_custom_class.h"
|
|
|
|
namespace torch::aot_inductor {
|
|
|
|
static auto registerMyAOTIClass =
|
|
torch::class_<MyAOTIClass>("aoti", "MyAOTIClass")
|
|
.def(torch::init<std::string, std::string>())
|
|
.def("forward", &MyAOTIClass::forward)
|
|
.def_pickle(
|
|
[](const c10::intrusive_ptr<MyAOTIClass>& self)
|
|
-> std::vector<std::string> {
|
|
std::vector<std::string> v;
|
|
v.push_back(self->lib_path());
|
|
v.push_back(self->device());
|
|
return v;
|
|
},
|
|
[](std::vector<std::string> params) {
|
|
return c10::make_intrusive<MyAOTIClass>(params[0], params[1]);
|
|
});
|
|
|
|
MyAOTIClass::MyAOTIClass(
|
|
const std::string& model_path,
|
|
const std::string& device)
|
|
: lib_path_(model_path), device_(device) {
|
|
if (device_ == "cpu") {
|
|
runner_ = std::make_unique<torch::inductor::AOTIModelContainerRunnerCpu>(
|
|
model_path.c_str());
|
|
#if defined(USE_CUDA) || defined(USE_ROCM)
|
|
} else if (device_ == "cuda") {
|
|
runner_ = std::make_unique<torch::inductor::AOTIModelContainerRunnerCuda>(
|
|
model_path.c_str());
|
|
#endif
|
|
} else {
|
|
throw std::runtime_error("invalid device: " + device);
|
|
}
|
|
}
|
|
|
|
std::vector<torch::Tensor> MyAOTIClass::forward(
|
|
std::vector<torch::Tensor> inputs) {
|
|
return runner_->run(inputs);
|
|
}
|
|
|
|
} // namespace torch::aot_inductor
|