Skip to content

Commit a8ba9d4

Browse files
alonre24alonre24
authored
Merge branch 'master' into Create_tensor_through_gears-refactor_and_test
2 parents 076b70a + b3d44c3 commit a8ba9d4

File tree

9 files changed

+55
-45
lines changed

9 files changed

+55
-45
lines changed

‎get_deps.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ MKL=mkl
7070
ONNXRUNTIME=onnxruntime
7171

7272
######################################################################################## DLPACK
73-
DLPACK_VERSION="v0.3"
73+
DLPACK_VERSION="v0.4"
7474
if [[ $WITH_DLPACK != 0 ]]; then
7575
[[ $FORCE == 1 ]] && rm -rf $DLPACK
7676

‎src/backends/onnxruntime.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,7 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
172172

173173
ret = RAI_TensorNew();
174174

175-
DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0};
175+
DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = 0};
176176

177177
OrtTensorTypeAndShapeInfo *info;
178178
status = ort->GetTensorTypeAndShape(v, &info);
@@ -243,7 +243,7 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
243243
// TODO: use manager_ctx to ensure ORT tensor doesn't get deallocated
244244
// This applies to outputs
245245

246-
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx,
246+
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device,
247247
#ifdef RAI_COPY_RUN_OUTPUT
248248
.data = data,
249249
#else

‎src/backends/tensorflow.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ RAI_Tensor *RAI_TensorCreateFromTFTensor(TF_Tensor *tensor, size_t batch_offset,
9090
long long batch_size) {
9191
RAI_Tensor *ret = RAI_TensorNew();
9292

93-
DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0};
93+
DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = 0};
9494

9595
const size_t ndims = TF_NumDims(tensor);
9696

@@ -129,7 +129,7 @@ RAI_Tensor *RAI_TensorCreateFromTFTensor(TF_Tensor *tensor, size_t batch_offset,
129129
// This applies to outputs
130130

131131
ret->tensor = (DLManagedTensor){
132-
.dl_tensor = (DLTensor){.ctx = ctx,
132+
.dl_tensor = (DLTensor){.device = device,
133133
#ifdef RAI_COPY_RUN_OUTPUT
134134
.data = data,
135135
#else

‎src/libtflite_c/tflite_c.cpp‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,11 @@ static DLDataType getDLDataType(const TfLiteTensor *tensor) {
4848
return dtype;
4949
}
5050

51-
static DLContext getDLContext(const TfLiteTensor *tensor, const int64_t &device_id) {
52-
DLContext ctx;
53-
ctx.device_id = device_id;
54-
// if (tensor->.is_cuda()) {
55-
// ctx.device_type = DLDeviceType::kDLGPU;
56-
// } else {
57-
// ctx.device_type = DLDeviceType::kDLCPU;
58-
// }
59-
ctx.device_type = DLDeviceType::kDLCPU;
60-
return ctx;
51+
static DLDevice getDLDevice(const TfLiteTensor *tensor, const int64_t &device_id) {
52+
DLDevice device;
53+
device.device_id = device_id;
54+
device.device_type = DLDeviceType::kDLCPU;
55+
return device;
6156
}
6257

6358
#if 0
@@ -139,10 +134,10 @@ DLManagedTensor *toManagedDLPack(std::shared_ptr<tflite::Interpreter> interprete
139134
DLDataType dtype = getDLDataType(tensor);
140135

141136
int64_t device_id = 0;
142-
DLContext ctx = getDLContext(tensor, device_id);
137+
DLDevice device = getDLDevice(tensor, device_id);
143138

144139
DLTensor dl_tensor = (DLTensor){.data = new uint8_t[tensor->bytes],
145-
.ctx = ctx,
140+
.device = device,
146141
.ndim = output_dims->size,
147142
.dtype = dtype,
148143
.shape = new int64_t[output_dims->size],

‎src/libtorch_c/torch_c.cpp‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -64,15 +64,15 @@ static DLDataType getDLDataType(const at::Tensor &t) {
6464
return dtype;
6565
}
6666

67-
static DLContext getDLContext(const at::Tensor &tensor, const int64_t &device_id) {
68-
DLContext ctx;
69-
ctx.device_id = device_id;
67+
static DLDevice getDLDevice(const at::Tensor &tensor, const int64_t &device_id) {
68+
DLDevice device;
69+
device.device_id = device_id;
7070
if (tensor.is_cuda()) {
71-
ctx.device_type = DLDeviceType::kDLGPU;
71+
device.device_type = DLDeviceType::kDLGPU;
7272
} else {
73-
ctx.device_type = DLDeviceType::kDLCPU;
73+
device.device_type = DLDeviceType::kDLCPU;
7474
}
75-
return ctx;
75+
return device;
7676
}
7777

7878
static at::DeviceType getATenDeviceType(DLDeviceType device_type) {
@@ -145,7 +145,7 @@ at::ScalarType toScalarType(const DLDataType &dtype) {
145145
}
146146

147147
torch::Tensor fromDLPack(const DLTensor *src) {
148-
at::DeviceType device_type = getATenDeviceType(src->ctx.device_type);
148+
at::DeviceType device_type = getATenDeviceType(src->device.device_type);
149149
at::ScalarType stype = toScalarType(src->dtype);
150150
// torch::Device device(device_type, src->ctx.device_id);
151151
torch::Device device(device_type, -1);
@@ -176,7 +176,7 @@ DLManagedTensor *toManagedDLPack(const torch::Tensor &src_) {
176176
if (src.is_cuda()) {
177177
device_id = src.get_device();
178178
}
179-
atDLMTensor->tensor.dl_tensor.ctx = getDLContext(src, device_id);
179+
atDLMTensor->tensor.dl_tensor.device = getDLDevice(src, device_id);
180180
atDLMTensor->tensor.dl_tensor.ndim = src.dim();
181181
atDLMTensor->tensor.dl_tensor.dtype = getDLDataType(src);
182182
atDLMTensor->tensor.dl_tensor.shape = const_cast<int64_t *>(src.sizes().data());

‎src/serialization/RDB/decoder/current/v1/decode_v1.c‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,15 @@ void *RAI_RDBLoadTensor_v1(RedisModuleIO *io) {
1212
int64_t *shape = NULL;
1313
int64_t *strides = NULL;
1414

15-
DLContext ctx;
16-
ctx.device_type = RedisModule_LoadUnsigned(io);
17-
ctx.device_id = RedisModule_LoadUnsigned(io);
15+
DLDevice device;
16+
device.device_type = RedisModule_LoadUnsigned(io);
17+
device.device_id = RedisModule_LoadUnsigned(io);
1818
if (RedisModule_IsIOError(io))
1919
goto cleanup;
2020

2121
// For now we only support CPU tensors (except during model and script run)
22-
assert(ctx.device_type == kDLCPU);
23-
assert(ctx.device_id == 0);
22+
assert(device.device_type == kDLCPU);
23+
assert(device.device_id == 0);
2424

2525
DLDataType dtype;
2626
dtype.bits = RedisModule_LoadUnsigned(io);
@@ -49,7 +49,7 @@ void *RAI_RDBLoadTensor_v1(RedisModuleIO *io) {
4949
goto cleanup;
5050

5151
RAI_Tensor *ret = RAI_TensorNew();
52-
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx,
52+
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device,
5353
.data = data,
5454
.ndim = ndims,
5555
.dtype = dtype,

‎src/serialization/RDB/decoder/previous/v0/decode_v0.c‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,15 @@ void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) {
55
int64_t *shape = NULL;
66
int64_t *strides = NULL;
77

8-
DLContext ctx;
9-
ctx.device_type = RedisModule_LoadUnsigned(io);
10-
ctx.device_id = RedisModule_LoadUnsigned(io);
8+
DLDevice device;
9+
device.device_type = RedisModule_LoadUnsigned(io);
10+
device.device_id = RedisModule_LoadUnsigned(io);
1111
if (RedisModule_IsIOError(io))
1212
goto cleanup;
1313

1414
// For now we only support CPU tensors (except during model and script run)
15-
assert(ctx.device_type == kDLCPU);
16-
assert(ctx.device_id == 0);
15+
assert(device.device_type == kDLCPU);
16+
assert(device.device_id == 0);
1717

1818
DLDataType dtype;
1919
dtype.bits = RedisModule_LoadUnsigned(io);
@@ -42,7 +42,7 @@ void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) {
4242
goto cleanup;
4343

4444
RAI_Tensor *ret = RAI_TensorNew();
45-
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx,
45+
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device,
4646
.data = data,
4747
.ndim = ndims,
4848
.dtype = dtype,

‎src/serialization/RDB/encoder/v1/encode_v1.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ void RAI_RDBSaveTensor_v1(RedisModuleIO *io, void *value) {
55

66
size_t ndim = tensor->tensor.dl_tensor.ndim;
77

8-
RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.ctx.device_type);
9-
RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.ctx.device_id);
8+
RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.device.device_type);
9+
RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.device.device_id);
1010
RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.dtype.bits);
1111
RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.dtype.code);
1212
RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.dtype.lanes);

‎src/tensor.c‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -122,10 +122,25 @@ RAI_Tensor *RAI_TensorCreateWithDLDataType(DLDataType dtype, long long *dims, in
122122
strides[i] *= strides[i + 1] * shape[i + 1];
123123
}
124124

125-
DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0};
126125
void *data = RedisModule_Calloc(len, dtypeSize);
127-
128-
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx,
126+
DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = 0};
127+
void *data = NULL;
128+
switch (tensorAllocMode) {
129+
case TENSORALLOC_ALLOC:
130+
data = RedisModule_Alloc(len * dtypeSize);
131+
break;
132+
case TENSORALLOC_CALLOC:
133+
data = RedisModule_Calloc(len, dtypeSize);
134+
break;
135+
case TENSORALLOC_NONE:
136+
/* shallow copy no alloc */
137+
default:
138+
/* assume TENSORALLOC_NONE
139+
shallow copy no alloc */
140+
break;
141+
}
142+
143+
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device,
129144
.data = data,
130145
.ndim = ndims,
131146
.dtype = dtype,
@@ -170,7 +185,7 @@ RAI_Tensor *_TensorCreateWithDLDataTypeAndRString(DLDataType dtype, size_t dtype
170185
strides[i] *= strides[i + 1] * shape[i + 1];
171186
}
172187

173-
DLContext ctx = (DLContext){.device_type = kDLCPU, .device_id = 0};
188+
DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = 0};
174189
size_t nbytes = len * dtypeSize;
175190

176191
size_t blob_len;
@@ -186,7 +201,7 @@ RAI_Tensor *_TensorCreateWithDLDataTypeAndRString(DLDataType dtype, size_t dtype
186201
RAI_HoldString(NULL, rstr);
187202

188203
RAI_Tensor *ret = RAI_TensorNew();
189-
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.ctx = ctx,
204+
ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device,
190205
.data = data,
191206
.ndim = ndims,
192207
.dtype = dtype,
@@ -327,7 +342,7 @@ RAI_Tensor *RAI_TensorCreateFromDLTensor(DLManagedTensor *dl_tensor) {
327342
RAI_Tensor *ret = RAI_TensorNew();
328343

329344
ret->tensor =
330-
(DLManagedTensor){.dl_tensor = (DLTensor){.ctx = dl_tensor->dl_tensor.ctx,
345+
(DLManagedTensor){.dl_tensor = (DLTensor){.device = dl_tensor->dl_tensor.device,
331346
.data = dl_tensor->dl_tensor.data,
332347
.ndim = dl_tensor->dl_tensor.ndim,
333348
.dtype = dl_tensor->dl_tensor.dtype,

0 commit comments

Comments
 (0)