| Line |
Branch |
Exec |
Source |
| 1 |
|
|
#include <torch/extension.h> |
| 2 |
|
|
#include "optimization/bagpipe/cache_metadata.h" |
| 3 |
|
|
|
| 4 |
|
|
#include <torch/custom_class.h> |
| 5 |
|
|
#include <algorithm> |
| 6 |
|
|
#include <tuple> |
| 7 |
|
|
|
| 8 |
|
|
namespace recstore::framework { |
| 9 |
|
|
namespace { |
| 10 |
|
|
|
| 11 |
|
|
// Tensor validation and conversion only; the policy lives in optimization/. |
| 12 |
|
|
class BagPipeMetadataBinding : public torch::CustomClassHolder { |
| 13 |
|
|
public: |
| 14 |
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
✗ |
explicit BagPipeMetadataBinding(int64_t capacity) : metadata_(capacity) {} |
| 15 |
|
|
|
| 16 |
|
✗ |
int64_t size() const { return metadata_.Size(); } |
| 17 |
|
|
|
| 18 |
|
✗ |
torch::Tensor previous( |
| 19 |
|
|
int64_t step, const torch::Tensor& local, const torch::Tensor& active) { |
| 20 |
|
✗ |
return Tensor(metadata_.Previous( |
| 21 |
0/2
✗ Branch 5 not taken.
✗ Branch 6 not taken.
|
✗ |
step, Keys(local), local.numel(), Keys(active), active.numel())); |
| 22 |
|
|
} |
| 23 |
|
|
|
| 24 |
|
✗ |
void clear_plan() { metadata_.ClearPlan(); } |
| 25 |
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
✗ |
torch::Tensor updated_ids() const { return Tensor(metadata_.UpdatedIds()); } |
| 26 |
|
|
|
| 27 |
|
✗ |
torch::Tensor prefetch(const torch::Tensor& ids, |
| 28 |
|
|
const torch::Tensor& previous, |
| 29 |
|
|
int64_t committed) const { |
| 30 |
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
✗ |
TORCH_CHECK(ids.numel() == previous.numel(), |
| 31 |
|
|
"BagPipe version count mismatch"); |
| 32 |
|
|
return Tensor( |
| 33 |
0/2
✗ Branch 5 not taken.
✗ Branch 6 not taken.
|
✗ |
metadata_.Prefetch(Keys(ids), Keys(previous), ids.numel(), committed)); |
| 34 |
|
|
} |
| 35 |
|
|
|
| 36 |
|
|
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> |
| 37 |
|
✗ |
lookup(const torch::Tensor& ids) { |
| 38 |
0/6
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
|
✗ |
auto result = metadata_.Lookup(Keys(ids), ids.numel()); |
| 39 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
return {Tensor(result.hit_positions), |
| 40 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
Tensor(result.hit_slots), |
| 41 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
Tensor(result.missing_positions)}; |
| 42 |
|
✗ |
} |
| 43 |
|
|
|
| 44 |
|
|
std::tuple<torch::Tensor, torch::Tensor, int64_t> |
| 45 |
|
✗ |
commit(const torch::Tensor& active, |
| 46 |
|
|
const torch::Tensor& local, |
| 47 |
|
|
const torch::Tensor& exclusive, |
| 48 |
|
|
bool ps) { |
| 49 |
0/30
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
✗ Branch 18 not taken.
✗ Branch 19 not taken.
✗ Branch 21 not taken.
✗ Branch 22 not taken.
✗ Branch 23 not taken.
✗ Branch 24 not taken.
✗ Branch 25 not taken.
✗ Branch 26 not taken.
✗ Branch 28 not taken.
✗ Branch 29 not taken.
✗ Branch 31 not taken.
✗ Branch 32 not taken.
✗ Branch 33 not taken.
✗ Branch 34 not taken.
✗ Branch 35 not taken.
✗ Branch 36 not taken.
|
✗ |
TORCH_CHECK( |
| 50 |
|
|
exclusive.device().is_cpu() && exclusive.dim() == 1 && |
| 51 |
|
|
exclusive.scalar_type() == torch::kBool && |
| 52 |
|
|
exclusive.is_contiguous() && exclusive.numel() == local.numel(), |
| 53 |
|
|
"BagPipe exclusive mask must match local IDs"); |
| 54 |
|
|
auto result = metadata_.Commit( |
| 55 |
|
|
Keys(active), |
| 56 |
|
✗ |
active.numel(), |
| 57 |
|
|
Keys(local), |
| 58 |
|
✗ |
local.numel(), |
| 59 |
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
✗ |
exclusive.data_ptr<bool>(), |
| 60 |
0/10
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
|
✗ |
ps); |
| 61 |
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
✗ |
return {Tensor(result.slots), Tensor(result.rows), result.invalidated}; |
| 62 |
|
✗ |
} |
| 63 |
|
|
|
| 64 |
|
|
std::tuple<torch::Tensor, torch::Tensor, int64_t> |
| 65 |
|
✗ |
admit(const torch::Tensor& ids) { |
| 66 |
0/6
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
|
✗ |
auto result = metadata_.Admit(Keys(ids), ids.numel()); |
| 67 |
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
✗ |
return {Tensor(result.slots), Tensor(result.rows), result.evicted}; |
| 68 |
|
✗ |
} |
| 69 |
|
|
|
| 70 |
|
|
private: |
| 71 |
|
✗ |
static const int64_t* Keys(const torch::Tensor& ids) { |
| 72 |
0/22
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
✗ Branch 18 not taken.
✗ Branch 19 not taken.
✗ Branch 21 not taken.
✗ Branch 22 not taken.
✗ Branch 23 not taken.
✗ Branch 24 not taken.
✗ Branch 25 not taken.
✗ Branch 26 not taken.
|
✗ |
TORCH_CHECK(ids.device().is_cpu() && ids.dim() == 1 && |
| 73 |
|
|
ids.scalar_type() == torch::kInt64 && ids.is_contiguous(), |
| 74 |
|
|
"BagPipe metadata requires contiguous CPU int64 IDs"); |
| 75 |
|
✗ |
return ids.data_ptr<int64_t>(); |
| 76 |
|
|
} |
| 77 |
|
|
|
| 78 |
|
✗ |
static torch::Tensor Tensor(const std::vector<int64_t>& values) { |
| 79 |
|
|
auto result = |
| 80 |
0/4
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
|
✗ |
torch::empty({static_cast<int64_t>(values.size())}, torch::kInt64); |
| 81 |
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
|
✗ |
std::copy(values.begin(), values.end(), result.data_ptr<int64_t>()); |
| 82 |
|
✗ |
return result; |
| 83 |
|
✗ |
} |
| 84 |
|
|
|
| 85 |
|
|
optimization::bagpipe::CacheMetadata metadata_; |
| 86 |
|
|
}; |
| 87 |
|
|
|
| 88 |
|
|
} // namespace |
| 89 |
|
|
|
| 90 |
|
16 |
void RegisterBagpipeMetadataOps(torch::Library& m) { |
| 91 |
2/4
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 16 times.
✗ Branch 6 not taken.
|
32 |
m.class_<BagPipeMetadataBinding>("BagPipeCacheMetadata") |
| 92 |
2/4
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
✓ Branch 5 taken 16 times.
✗ Branch 6 not taken.
|
32 |
.def(torch::init<int64_t>()) |
| 93 |
3/6
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 16 times.
✗ Branch 10 not taken.
|
48 |
.def("size", &BagPipeMetadataBinding::size) |
| 94 |
3/6
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 16 times.
✗ Branch 10 not taken.
|
48 |
.def("previous", &BagPipeMetadataBinding::previous) |
| 95 |
3/6
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 16 times.
✗ Branch 10 not taken.
|
48 |
.def("clear_plan", &BagPipeMetadataBinding::clear_plan) |
| 96 |
3/6
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 16 times.
✗ Branch 10 not taken.
|
48 |
.def("updated_ids", &BagPipeMetadataBinding::updated_ids) |
| 97 |
3/6
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 16 times.
✗ Branch 10 not taken.
|
48 |
.def("prefetch", &BagPipeMetadataBinding::prefetch) |
| 98 |
3/6
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 16 times.
✗ Branch 10 not taken.
|
48 |
.def("lookup", &BagPipeMetadataBinding::lookup) |
| 99 |
3/6
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 16 times.
✗ Branch 10 not taken.
|
48 |
.def("commit", &BagPipeMetadataBinding::commit) |
| 100 |
3/6
✓ Branch 3 taken 16 times.
✗ Branch 4 not taken.
✓ Branch 7 taken 16 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
|
32 |
.def("admit", &BagPipeMetadataBinding::admit); |
| 101 |
|
16 |
} |
| 102 |
|
|
|
| 103 |
|
|
} // namespace recstore::framework |
| 104 |
|
|
|