Skip to content

Commit 361bf44

Browse files
mzientmzient
committed
Tensor and TensorList C API wrappers.
Signed-off-by: Michał Zientkiewicz <[email protected]>
1 parent e9db30f commit 361bf44

File tree

3 files changed

+1389
-0
lines changed

3 files changed

+1389
-0
lines changed

‎dali/c_api_2/data_objects.cc

Lines changed: 362 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,362 @@
1+
// Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "dali/c_api_2/data_objects.h"
16+
#include "dali/c_api_2/error_handling.h"
17+
18+
namespace dali::c_api {
19+
20+
RefCountedPtr<ITensor> ITensor::Create(daliBufferPlacement_t placement) {
21+
Validate(placement);
22+
switch (placement.device_type) {
23+
case DALI_STORAGE_CPU:
24+
{
25+
auto tl = std::make_shared<Tensor<CPUBackend>>();
26+
tl->set_pinned(placement.pinned);
27+
if (placement.pinned)
28+
tl->set_device_id(placement.device_id);
29+
return Wrap(std::move(tl));
30+
}
31+
case DALI_STORAGE_GPU:
32+
{
33+
auto tl = std::make_shared<Tensor<GPUBackend>>();
34+
tl->set_pinned(placement.pinned);
35+
tl->set_device_id(placement.device_id);
36+
return Wrap(std::move(tl));
37+
}
38+
default:
39+
assert(!"Unreachable code");
40+
return {};
41+
}
42+
}
43+
44+
RefCountedPtr<ITensorList> ITensorList::Create(daliBufferPlacement_t placement) {
45+
Validate(placement);
46+
switch (placement.device_type) {
47+
case DALI_STORAGE_CPU:
48+
{
49+
auto tl = std::make_shared<TensorList<CPUBackend>>();
50+
tl->set_pinned(placement.pinned);
51+
if (placement.pinned)
52+
tl->set_device_id(placement.device_id);
53+
return Wrap(std::move(tl));
54+
}
55+
case DALI_STORAGE_GPU:
56+
{
57+
auto tl = std::make_shared<TensorList<GPUBackend>>();
58+
tl->set_pinned(placement.pinned);
59+
tl->set_device_id(placement.device_id);
60+
return Wrap(std::move(tl));
61+
}
62+
default:
63+
assert(!"Unreachable code");
64+
return {};
65+
}
66+
}
67+
68+
ITensor *ToPointer(daliTensor_h handle) {
69+
if (!handle)
70+
throw NullHandle("Tensor");
71+
return static_cast<ITensor *>(handle);
72+
}
73+
74+
ITensorList *ToPointer(daliTensorList_h handle) {
75+
if (!handle)
76+
throw NullHandle("TensorList");
77+
return static_cast<ITensorList *>(handle);
78+
}
79+
80+
} // namespace dali::c_api
81+
82+
using namespace dali::c_api; // NOLINT
83+
84+
template <typename T>
85+
std::optional<T> ToOptional(const T *nullable) {
86+
if (nullable == nullptr)
87+
return std::nullopt;
88+
else
89+
return *nullable;
90+
}
91+
92+
//////////////////////////////////////////////////////////////////////////////
93+
// Tensor
94+
//////////////////////////////////////////////////////////////////////////////
95+
96+
daliResult_t daliTensorCreate(daliTensor_h *out, daliBufferPlacement_t placement) {
97+
DALI_PROLOG();
98+
if (!out)
99+
throw std::invalid_argument("The output parameter must not be NULL.");
100+
auto t = dali::c_api::ITensor::Create(placement);
101+
*out = t.release(); // no throwing allowed after this line!
102+
DALI_EPILOG();
103+
}
104+
105+
daliResult_t daliTensorIncRef(daliTensor_h t, int *new_ref) {
106+
DALI_PROLOG();
107+
auto *ptr = ToPointer(t);
108+
int r = ptr->IncRef();
109+
if (new_ref)
110+
*new_ref = r;
111+
DALI_EPILOG();
112+
}
113+
114+
daliResult_t daliTensorDecRef(daliTensor_h t, int *new_ref) {
115+
DALI_PROLOG();
116+
auto *ptr = ToPointer(t);
117+
int r = ptr->DecRef();
118+
if (new_ref)
119+
*new_ref = r;
120+
DALI_EPILOG();
121+
}
122+
123+
daliResult_t daliTensorRefCount(daliTensor_h t, int *ref) {
124+
DALI_PROLOG();
125+
auto *ptr = ToPointer(t);
126+
int r = ptr->RefCount();
127+
if (!ref)
128+
throw std::invalid_argument("The output parameter must not be NULL.");
129+
*ref = r;
130+
DALI_EPILOG();
131+
}
132+
133+
daliResult_t daliTensorAttachBuffer(
134+
daliTensor_h tensor,
135+
int ndim,
136+
const int64_t *shape,
137+
daliDataType_t dtype,
138+
const char *layout,
139+
void *data,
140+
daliDeleter_t deleter) {
141+
DALI_PROLOG();
142+
auto *ptr = ToPointer(tensor);
143+
ptr->AttachBuffer(ndim, shape, dtype, layout, data, deleter);
144+
DALI_EPILOG();
145+
}
146+
147+
daliResult_t daliTensorResize(
148+
daliTensor_h tensor,
149+
int ndim,
150+
const int64_t *shape,
151+
daliDataType_t dtype,
152+
const char *layout) {
153+
DALI_PROLOG();
154+
auto *ptr = ToPointer(tensor);
155+
ptr->Resize(ndim, shape, dtype, layout);
156+
DALI_EPILOG();
157+
}
158+
159+
daliResult_t daliTensorSetLayout(
160+
daliTensor_h tensor,
161+
const char *layout) {
162+
DALI_PROLOG();
163+
auto *ptr = ToPointer(tensor);
164+
ptr->SetLayout(layout);
165+
DALI_EPILOG();
166+
}
167+
168+
daliResult_t daliTensorGetLayout(
169+
daliTensor_h tensor,
170+
const char **layout) {
171+
DALI_PROLOG();
172+
auto *ptr = ToPointer(tensor);
173+
if (!layout)
174+
throw std::invalid_argument("The output parameter `layout` must not be be NULL");
175+
*layout = ptr->GetLayout();
176+
DALI_EPILOG();
177+
}
178+
179+
daliResult_t daliTensorGetStream(
180+
daliTensor_h tensor,
181+
cudaStream_t *out_stream) {
182+
DALI_PROLOG();
183+
auto *ptr = ToPointer(tensor);
184+
if (!out_stream)
185+
throw std::invalid_argument("The output parameter `out_stream` must not be NULL");
186+
auto str = ptr->GetStream();
187+
*out_stream = str.has_value() ? *str : cudaStream_t(-1);
188+
return str.has_value() ? DALI_SUCCESS : DALI_NO_DATA;
189+
DALI_EPILOG();
190+
}
191+
192+
daliResult_t daliTensorSetStream(
193+
daliTensor_h tensor,
194+
const cudaStream_t *stream,
195+
daliBool synchronize) {
196+
DALI_PROLOG();
197+
auto *ptr = ToPointer(tensor);
198+
ptr->SetStream(ToOptional(stream), synchronize);
199+
DALI_EPILOG();
200+
}
201+
202+
daliResult_t daliTensorGetDesc(
203+
daliTensor_h tensor,
204+
daliTensorDesc_t *out_desc) {
205+
DALI_PROLOG();
206+
auto *ptr = ToPointer(tensor);
207+
if (!out_desc)
208+
throw std::invalid_argument("The output parameter `out_desc` must not be NULL.");
209+
*out_desc = ptr->GetDesc();
210+
DALI_EPILOG();
211+
}
212+
213+
//////////////////////////////////////////////////////////////////////////////
214+
// TensorList
215+
//////////////////////////////////////////////////////////////////////////////
216+
217+
daliResult_t daliTensorListCreate(daliTensorList_h *out, daliBufferPlacement_t placement) {
218+
DALI_PROLOG();
219+
if (!out)
220+
throw std::invalid_argument("The output parameter must not be NULL.");
221+
auto tl = dali::c_api::ITensorList::Create(placement);
222+
*out = tl.release(); // no throwing allowed after this line!
223+
DALI_EPILOG();
224+
}
225+
226+
daliResult_t daliTensorListIncRef(daliTensorList_h tl, int *new_ref) {
227+
DALI_PROLOG();
228+
auto *ptr = ToPointer(tl);
229+
int r = ptr->IncRef();
230+
if (new_ref)
231+
*new_ref = r;
232+
DALI_EPILOG();
233+
}
234+
235+
daliResult_t daliTensorListDecRef(daliTensorList_h tl, int *new_ref) {
236+
DALI_PROLOG();
237+
auto *ptr = ToPointer(tl);
238+
int r = ptr->DecRef();
239+
if (new_ref)
240+
*new_ref = r;
241+
DALI_EPILOG();
242+
}
243+
244+
daliResult_t daliTensorListRefCount(daliTensorList_h tl, int *ref) {
245+
DALI_PROLOG();
246+
auto *ptr = ToPointer(tl);
247+
if (!ref)
248+
throw std::invalid_argument("The output pointer must not be NULL.");
249+
int r = ptr->RefCount();
250+
*ref = r;
251+
DALI_EPILOG();
252+
}
253+
254+
daliResult_t daliTensorListAttachBuffer(
255+
daliTensorList_h tensor_list,
256+
int num_samples,
257+
int ndim,
258+
const int64_t *shapes,
259+
daliDataType_t dtype,
260+
const char *layout,
261+
void *data,
262+
const ptrdiff_t *sample_offsets,
263+
daliDeleter_t deleter) {
264+
DALI_PROLOG();
265+
auto *ptr = ToPointer(tensor_list);
266+
ptr->AttachBuffer(num_samples, ndim, shapes, dtype, layout, data, sample_offsets, deleter);
267+
DALI_EPILOG();
268+
}
269+
270+
daliResult_t daliTensorListAttachSamples(
271+
daliTensorList_h tensor_list,
272+
int num_samples,
273+
int ndim,
274+
daliDataType_t dtype,
275+
const char *layout,
276+
const daliTensorDesc_t *samples,
277+
const daliDeleter_t *sample_deleters) {
278+
DALI_PROLOG();
279+
auto *ptr = ToPointer(tensor_list);
280+
ptr->AttachSamples(num_samples, ndim, dtype, layout, samples, sample_deleters);
281+
DALI_EPILOG();
282+
}
283+
284+
daliResult_t daliTensorListResize(
285+
daliTensorList_h tensor_list,
286+
int num_samples,
287+
int ndim,
288+
const int64_t *shapes,
289+
daliDataType_t dtype,
290+
const char *layout) {
291+
DALI_PROLOG();
292+
auto *ptr = ToPointer(tensor_list);
293+
ptr->Resize(num_samples, ndim, shapes, dtype, layout);
294+
DALI_EPILOG();
295+
}
296+
297+
daliResult_t daliTensorListSetLayout(
298+
daliTensorList_h tensor_list,
299+
const char *layout) {
300+
DALI_PROLOG();
301+
auto *ptr = ToPointer(tensor_list);
302+
ptr->SetLayout(layout);
303+
DALI_EPILOG();
304+
}
305+
306+
daliResult_t daliTensorListGetLayout(
307+
daliTensorList_h tensor_list,
308+
const char **layout) {
309+
DALI_PROLOG();
310+
auto *ptr = ToPointer(tensor_list);
311+
if (!layout)
312+
throw std::invalid_argument("The output parameter `layout` must not be be NULL");
313+
*layout = ptr->GetLayout();
314+
DALI_EPILOG();
315+
}
316+
317+
daliResult_t daliTensorListGetStream(
318+
daliTensorList_h tensor_list,
319+
cudaStream_t *out_stream) {
320+
DALI_PROLOG();
321+
auto *ptr = ToPointer(tensor_list);
322+
if (!out_stream)
323+
throw std::invalid_argument("The output parameter `out_stream` must not be NULL");
324+
auto str = ptr->GetStream();
325+
*out_stream = str.has_value() ? *str : cudaStream_t(-1);
326+
return str.has_value() ? DALI_SUCCESS : DALI_NO_DATA;
327+
DALI_EPILOG();
328+
}
329+
330+
daliResult_t daliTensorListSetStream(
331+
daliTensorList_h tensor_list,
332+
const cudaStream_t *stream,
333+
daliBool synchronize) {
334+
DALI_PROLOG();
335+
auto *ptr = ToPointer(tensor_list);
336+
ptr->SetStream(ToOptional(stream), synchronize);
337+
DALI_EPILOG();
338+
}
339+
340+
daliResult_t daliTensorListGetTensorDesc(
341+
daliTensorList_h tensor_list,
342+
daliTensorDesc_t *out_tensor,
343+
int sample_idx) {
344+
DALI_PROLOG();
345+
auto *ptr = ToPointer(tensor_list);
346+
if (!out_tensor)
347+
throw std::invalid_argument("The output parameter `out_tensor` must not be NULL.");
348+
*out_tensor = ptr->GetTensorDesc(sample_idx);
349+
DALI_EPILOG();
350+
}
351+
352+
daliResult_t daliTensorListViewAsTensor(
353+
daliTensorList_h tensor_list,
354+
daliTensor_h *out_tensor) {
355+
DALI_PROLOG();
356+
auto *ptr = ToPointer(tensor_list);
357+
if (!out_tensor)
358+
throw std::invalid_argument("The output parameter `out_tensor` must not be NULL.");
359+
auto t = ptr->ViewAsTensor();
360+
*out_tensor = t.release(); // no throwing allowed after this line
361+
DALI_EPILOG();
362+
}

0 commit comments

Comments
 (0)