import torch
import libgpucache

# args: strtegy, capacity, key_size(only support int32 and int64), value_size(dim * sizeof(elem), elem is decided by the passing tensor dtype of NewLRUCache), device_id, dim
cfg = libgpucache.CacheConfig(libgpucache.CacheConfig.LRU,65536,4,128,4096,0,32)
t = torch.empty([1],dtype=torch.float32)
cache = libgpucache.NewLRUCache(t,cfg)

keys = torch.tensor([1,2,3],dtype=torch.int32,device="cuda:0")
values, find_mask = cache.Get(3,keys)
print("--------------------------------")
print("get result")
print(values)
print(find_mask)
print("--------------------------------")
values = torch.arange(0,3*32,dtype=torch.float32).reshape(3,-1).to("cuda:0")
cache.Put(3,keys,values)
values, find_mask = cache.Get(3,keys)
print("--------------------------------")
print("after put get result")
values, find_mask = cache.Get(3,keys)
print(values)
print(find_mask)
print("--------------------------------")

