chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add AddVector and GetVector to user cache.

Browse Source 
PiperOrigin-RevId: 601485067
Change-Id: I1a31f15c2495de4f057316022dcddf4df39afd64
This commit is contained in:
Kyle Bayes
2024-01-25 10:12:29 -08:00
committed by Copybara-Service
parent a1e5ede266
commit 4f53d9a0d7
3 changed files with 102 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/user_cache.cc
+58 -15
View File
@@ -26,35 +26,54 @@
#include <utility>
#include <vector>
// adds a block of data for the asset and returns number of bytes stored
template<typename T>
std::size_t mjCAsset::Add(const std::string& name, const std::vector<T>& v) {
// copies a block of data into the asset and returns number of bytes stored
// loading data into an asset should happen in a single thread
template<typename T> std::size_t mjCAsset::Add(const std::string& name,
const T* data, std::size_t n) {
auto [it, inserted] = blocks_.insert({name, mjCAssetData()});
if (!inserted) {
return 0;
}
std::size_t n = v.size() * sizeof(T);
const uint8_t* ptr = reinterpret_cast<const uint8_t*>(v.data());
std::size_t nbytes = n * sizeof(T);
const uint8_t* ptr = reinterpret_cast<const uint8_t*>(data);
mjCAssetData& block = it->second;
block.bytes = std::make_shared<uint8_t[]>(n);
std::copy(ptr, ptr + n, block.bytes.get());
block.nbytes = n;
nbytes_ += n;
return n;
block.bytes = std::make_shared<uint8_t[]>(nbytes);
std::copy(ptr, ptr + nbytes, block.bytes.get());
block.nbytes = nbytes;
nbytes_ += nbytes;
return nbytes;
}
template std::size_t mjCAsset::Add(const std::string& name,
const std::vector<int>& v);
const uint8_t* data, std::size_t n);
template std::size_t mjCAsset::Add(const std::string& name,
const std::vector<float>& v);
const int* data, std::size_t n);
template std::size_t mjCAsset::Add(const std::string& name,
const std::vector<double>& v);
const float* data, std::size_t n);
template std::size_t mjCAsset::Add(const std::string& name,
const double* data, std::size_t n);
// fetches a block of data, sets n to size of data
// copies a vector into the asset and returns number of bytes stored
// loading data into an asset should happen in a single thread
template<typename T> std::size_t mjCAsset::AddVector(const std::string& name,
const std::vector<T>& v) {
return Add<T>(name, v.data(), v.size());
}
template std::size_t mjCAsset::AddVector(const std::string& name,
const std::vector<uint8_t>& v);
template std::size_t mjCAsset::AddVector(const std::string& name,
const std::vector<int>& v);
template std::size_t mjCAsset::AddVector(const std::string& name,
const std::vector<float>& v);
template std::size_t mjCAsset::AddVector(const std::string& name,
const std::vector<double>& v);
// returns a pointer to a block of data, sets n to size of data
template<typename T>
const T* mjCAsset::Get(const std::string& name, std::size_t* n) const {
auto it = blocks_.find(name);
@@ -76,6 +95,8 @@ const T* mjCAsset::Get(const std::string& name, std::size_t* n) const {
return reinterpret_cast<T*>(data.bytes.get());
}
template
const uint8_t* mjCAsset::Get(const std::string& name, std::size_t* n) const;
template
const int* mjCAsset::Get(const std::string& name, std::size_t* n) const;
template
@@ -85,6 +106,28 @@ const double* mjCAsset::Get(const std::string& name, std::size_t* n) const;
// copies a block of data into a vector
template<typename T> std::optional<std::vector<T>>
mjCAsset::GetVector(const std::string& name) const {
std::size_t n;
const T* ptr = Get<T>(name, &n);
if (ptr == nullptr) {
return std::nullopt;
}
return std::vector<T>(ptr, ptr + n);
}
template std::optional<std::vector<uint8_t>>
mjCAsset::GetVector(const std::string& name) const;
template std::optional<std::vector<int>>
mjCAsset::GetVector(const std::string& name) const;
template std::optional<std::vector<float>>
mjCAsset::GetVector(const std::string& name) const;
template std::optional<std::vector<double>>
mjCAsset::GetVector(const std::string& name) const;
// replaces blocks data in asset
void mjCAsset::ReplaceBlocks(
const std::unordered_map<std::string, mjCAssetData>& blocks,
src/user/user_cache.h
+19 -6
View File
@@ -56,17 +56,30 @@ class mjCAsset {
mjCAsset(const mjCAsset& other) = default;
mjCAsset& operator=(const mjCAsset& other) = default;
// adds a block of data for the asset and returns number of bytes stored
// loading data into an asset should happen single thread
// copies a block of data into the asset and returns number of bytes stored
// loading data into an asset should happen in a single thread
template<typename T> std::size_t Add(const std::string& name,
const std::vector<T>& v);
const T* data, std::size_t n);
// fetches a block of data, sets n to size of data
// TODO(kylebayes): The C++ span utility doesn't seem to be supported by
// Google C++ coding standards. For now, we fallback to an C style API.
// copies a vector into the asset and returns number of bytes stored
// loading data into an asset should happen in a single thread
template<typename T> std::size_t AddVector(const std::string& name,
const std::vector<T>& v);
// returns a pointer to a block of data, sets n to size of data
template<typename T>
const T* Get(const std::string& name, std::size_t* n) const;
// copies a block of data into a vector
template<typename T>
std::optional<std::vector<T>> GetVector(const std::string& name) const;
// returns true if a block of data by the given name is stored in the asset
bool HasData(const std::string& name) const {
return blocks_.find(name) != blocks_.end();
}
private:
mjCAsset() = default;
test/user/user_cache_test.cc
+25 -36
View File
@@ -66,8 +66,8 @@ TEST(CacheTest, AddSuccessTest) {
std::vector<double> v2 = {1.0, 2.0, 3.0};
mjCAsset asset("file.xml", "foo.obj", "now");
std::size_t nbytes1 = asset.Add("v1", v1);
std::size_t nbytes2 = asset.Add("v2", v2);
std::size_t nbytes1 = asset.AddVector("v1", v1);
std::size_t nbytes2 = asset.AddVector("v2", v2);
cache.Insert(asset);
ASSERT_EQ(nbytes1, 12);
@@ -81,8 +81,8 @@ TEST(CacheTest, AddFailureTest) {
std::vector<double> v2 = {1.0, 2.0, 3.0};
mjCAsset asset("file.xml", "foo.obj", "now");
asset.Add("v1", v1);
std::size_t nbytes = asset.Add("v1", v2);
asset.AddVector("v1", v1);
std::size_t nbytes = asset.AddVector("v1", v2);
cache.Insert(asset);
ASSERT_EQ(nbytes, 0);
@@ -95,18 +95,16 @@ TEST(CacheTest, InsertReplaceTest) {
std::vector<double> v2 = {1.0, 2.0, 3.0};
mjCAsset asset("file.xml", "foo.obj", "now");
asset.Add("v", v1);
asset.AddVector("v", v1);
cache.Insert(asset);
mjCAsset asset2("file.xml", "foo.obj", "nower");
asset2.Add("v", v2);
asset2.AddVector("v", v2);
bool inserted = cache.Insert(asset2);
EXPECT_TRUE(inserted);
mjCAsset asset3 = *(cache.Get("foo.obj"));
std::size_t n = 0;
const double* ptr = asset3.Get<double>("v", &n);
std::vector<double> v3 = std::vector<double>(ptr, ptr + n);
std::vector<double> v3 = asset3.GetVector<double>("v").value();
EXPECT_THAT(v3, ElementsAreArray(v2));
ASSERT_EQ(cache.Size(), 24);
@@ -118,8 +116,8 @@ TEST(CacheTest, MoveInsertNewTest) {
std::vector<double> v2 = {1.0, 2.0, 3.0};
mjCAsset asset("file.xml", "foo.obj", "now");
std::size_t nbytes1 = asset.Add("v1", v1);
std::size_t nbytes2 = asset.Add("v2", v2);
std::size_t nbytes1 = asset.AddVector("v1", v1);
std::size_t nbytes2 = asset.AddVector("v2", v2);
cache.Insert(std::move(asset));
ASSERT_EQ(nbytes1, 12);
@@ -133,18 +131,16 @@ TEST(CacheTest, MoveInsertReplaceTest) {
std::vector<double> v2 = {1.0, 2.0, 3.0};
mjCAsset asset("file.xml", "foo.obj", "now");
asset.Add("v", v1);
asset.AddVector("v", v1);
cache.Insert(std::move(asset));
mjCAsset asset2("file.xml", "foo.obj", "nower");
asset2.Add("v", v2);
asset2.AddVector("v", v2);
bool inserted = cache.Insert(std::move(asset2));
EXPECT_TRUE(inserted);
mjCAsset asset3 = *(cache.Get("foo.obj"));
std::size_t n = 0;
const double* ptr = asset3.Get<double>("v", &n);
std::vector<double> v3 = std::vector<double>(ptr, ptr + n);
std::vector<double> v3 = asset3.GetVector<double>("v").value();
EXPECT_THAT(v3, ElementsAreArray(v2));
ASSERT_EQ(cache.Size(), 24);
@@ -153,40 +149,33 @@ TEST(CacheTest, MoveInsertReplaceTest) {
TEST(CacheTest, GetSuccessTest) {
mjCCache cache(kMaxSize);
std::vector<int> v1 = {1, 2, 3};
std::vector<double> v2 = {1.0, 2.0, 3.0};
std::vector<double> v2 = {4.0, 5.0, 6.0};
mjCAsset asset("file.xml", "foo.obj", "now");
asset.Add("v1", v1);
asset.Add("v2", v2);
asset.AddVector("v1", v1);
asset.AddVector("v2", v2);
cache.Insert(asset);
std::size_t n = 0;
mjCAsset asset2 = *(cache.Get("foo.obj"));
const int* ptr1 = asset2.Get<int>("v1", &n);
EXPECT_EQ(n, 3);
std::vector<int> v3 = std::vector<int>(ptr1, ptr1 + n);
std::vector<int> v3 = asset2.GetVector<int>("v1").value();
EXPECT_THAT(v3, ElementsAreArray(v1));
const double* ptr2 = asset2.Get<double>("v2", &n);
EXPECT_EQ(n, 3);
std::vector<double> v4 = std::vector<double>(ptr2, ptr2 + n);
EXPECT_THAT(v4, ElementsAreArray(v3));
std::vector<double> v4 = asset2.GetVector<double>("v2").value();
EXPECT_THAT(v4, ElementsAreArray(v2));
}
TEST(CacheTest, GetFailueTest) {
mjCCache cache(kMaxSize);
std::vector<int> v = {1, 2, 3};
mjCAsset asset("file.xml", "foo.obj", "now");
asset.Add("v", v);
asset.AddVector("v", v);
cache.Insert(asset);
mjCAsset asset2 = *(cache.Get("foo.obj"));
EXPECT_EQ(cache.Get("bar.obj").has_value(), false);
std::size_t n = 0;
const int* ptr = asset.Get<int>("v2", &n);
EXPECT_THAT(ptr, IsNull());
auto v2 = asset2.GetVector<int>("v2");
EXPECT_EQ(v2.has_value(), false);
}
// Trim cache based off of access count
@@ -197,8 +186,8 @@ TEST(CacheTest, LimitTest1) {
mjCAsset asset1 = mjCAsset("file.xml", "foo.obj", "now");
mjCAsset asset2 = mjCAsset("file.xml", "bar.obj", "now");
asset1.Add("foo.obj", v);
asset2.Add("bar.obj", v);
asset1.AddVector("foo.obj", v);
asset2.AddVector("bar.obj", v);
cache.Insert(asset1);
cache.Insert(asset2);
@@ -223,8 +212,8 @@ TEST(CacheTest, LimitTest2) {
std::vector<int> v = {1, 2, 3};
mjCAsset asset1("file.xml", "foo.obj", "now");
mjCAsset asset2("file.xml", "bar.obj", "now");
asset1.Add("v", v);
asset2.Add("v", v);
asset1.AddVector("v", v);
asset2.AddVector("v", v);
cache.Insert(asset1);
cache.Insert(asset2);