mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
0a57a20c02
Summary: Currently, we define some C++ functions in one C++ Python extension which are used by another. This happens to work, but isn't guaranteed to. This diff moves these functions to a separate C++ library rule to fix this. Test Plan: CI Differential Revision: D42552515 Pull Request resolved: https://github.com/pytorch/pytorch/pull/92325 Approved by: https://github.com/kit1980, https://github.com/Skylion007
43 lines
1.2 KiB
C++
43 lines
1.2 KiB
C++
#pragma once
|
|
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl.h>
|
|
|
|
//#include <Python.h>
|
|
|
|
namespace caffe2 {
|
|
namespace python {
|
|
class C10_EXPORT BlobFetcherBase {
|
|
public:
|
|
struct FetchedBlob {
|
|
pybind11::object obj;
|
|
bool copied;
|
|
};
|
|
virtual ~BlobFetcherBase();
|
|
virtual pybind11::object Fetch(const Blob& blob) = 0;
|
|
};
|
|
|
|
C10_DECLARE_TYPED_REGISTRY(
|
|
BlobFetcherRegistry,
|
|
TypeIdentifier,
|
|
BlobFetcherBase,
|
|
std::unique_ptr);
|
|
#define REGISTER_BLOB_FETCHER(id, ...) \
|
|
C10_REGISTER_TYPED_CLASS(BlobFetcherRegistry, id, __VA_ARGS__)
|
|
inline unique_ptr<BlobFetcherBase> CreateFetcher(TypeIdentifier id) {
|
|
return BlobFetcherRegistry()->Create(id);
|
|
}
|
|
|
|
Workspace* GetCurrentWorkspace();
|
|
void SetCurrentWorkspace(Workspace* workspace);
|
|
Workspace* NewWorkspace();
|
|
Workspace* GetWorkspaceByName(const std::string& name);
|
|
std::string GetCurrentWorkspaceName();
|
|
void InsertWorkspace(const std::string& name, std::unique_ptr<Workspace> ws);
|
|
void SwitchWorkspaceInternal(const std::string& name, bool create_if_missing);
|
|
void ResetWorkspace(Workspace* workspace);
|
|
void GetWorkspaceNames(std::vector<std::string>& names);
|
|
void ClearWorkspaces();
|
|
} // namespace python
|
|
} // namespace caffe2
|