|
24 | 24 | #include <utility> |
25 | 25 | #include <algorithm> |
26 | 26 |
|
| 27 | +// tc |
| 28 | +#include <tc/core/check.h> |
| 29 | +#include <tc/core/compiler.h> |
| 30 | +#include <tc/core/tc_executor.h> |
| 31 | +#include <tc/core/tensor.h> |
| 32 | +// tc |
27 | 33 |
|
28 | 34 | #define Noaxis 9999 |
29 | 35 |
|
| 36 | +// namespace is already exist in singa |
| 37 | +// aliasing to avoid duplicates |
| 38 | +namespace tclang = lang; |
| 39 | + |
30 | 40 | namespace singa { |
31 | 41 |
|
32 | 42 | Tensor::~Tensor() { |
@@ -1334,4 +1344,109 @@ Tensor Reshape(const Tensor &in, const Shape &s) { |
1334 | 1344 | return out.Reshape(s); |
1335 | 1345 | } |
1336 | 1346 |
|
| 1347 | + |
| 1348 | +/// tc integration start |
| 1349 | +struct SingaDLManagedTensor { |
| 1350 | + Tensor handle; |
| 1351 | + DLManagedTensor tensor; |
| 1352 | +}; |
| 1353 | + |
| 1354 | +void deleter(DLManagedTensor *arg) { |
| 1355 | + delete static_cast<SingaDLManagedTensor *>(arg->manager_ctx); |
| 1356 | +} |
| 1357 | + |
| 1358 | +static DLDataType getDLDataType(const Tensor &t) { |
| 1359 | + DLDataType dtype; |
| 1360 | + dtype.lanes = 1; |
| 1361 | + dtype.bits = SizeOf(t.data_type()) * 8; |
| 1362 | + switch (t.data_type()) { |
| 1363 | + case kFloat32: |
| 1364 | + dtype.code = DLDataTypeCode::kDLFloat; |
| 1365 | + break; |
| 1366 | + default: |
| 1367 | + throw std::logic_error("only kFloat32 is supported for dlpack conversion"); |
| 1368 | + break; |
| 1369 | + } |
| 1370 | + return dtype; |
| 1371 | +} |
| 1372 | + |
| 1373 | +static DLContext getDLContext(const Tensor &tensor, const int64_t &device_id) { |
| 1374 | + DLContext ctx; |
| 1375 | + ctx.device_id = device_id; |
| 1376 | + if (tensor.device()->lang() == kCuda) { |
| 1377 | + ctx.device_type = DLDeviceType::kDLGPU; |
| 1378 | + } else { |
| 1379 | + ctx.device_type = DLDeviceType::kDLCPU; |
| 1380 | + } |
| 1381 | + return ctx; |
| 1382 | +} |
| 1383 | + |
| 1384 | +// This function returns a shared_ptr to memory managed DLpack tensor |
| 1385 | +// constructed out of ATen tensor |
| 1386 | +DLManagedTensor *toDLPack(const Tensor &src) { |
| 1387 | + SingaDLManagedTensor *singaDLManagedTensor(new SingaDLManagedTensor); |
| 1388 | + singaDLManagedTensor->handle = src; |
| 1389 | + singaDLManagedTensor->tensor.manager_ctx = singaDLManagedTensor; |
| 1390 | + singaDLManagedTensor->tensor.deleter = &deleter; |
| 1391 | + singaDLManagedTensor->tensor.dl_tensor.data = src.block()->mutable_data(); |
| 1392 | + int64_t device_id = src.device()->id(); |
| 1393 | + singaDLManagedTensor->tensor.dl_tensor.ctx = getDLContext(src, device_id); |
| 1394 | + singaDLManagedTensor->tensor.dl_tensor.ndim = src.nDim(); |
| 1395 | + singaDLManagedTensor->tensor.dl_tensor.dtype = getDLDataType(src); |
| 1396 | + |
| 1397 | + auto shapeVec = |
| 1398 | + new std::vector<int64_t>(src.shape().begin(), src.shape().end()); |
| 1399 | + singaDLManagedTensor->tensor.dl_tensor.shape = shapeVec->data(); |
| 1400 | + |
| 1401 | + auto strideVec = |
| 1402 | + new std::vector<int64_t>(src.stride().begin(), src.stride().end()); |
| 1403 | + singaDLManagedTensor->tensor.dl_tensor.strides = strideVec->data(); |
| 1404 | + |
| 1405 | + singaDLManagedTensor->tensor.dl_tensor.byte_offset = 0; |
| 1406 | + return &(singaDLManagedTensor->tensor); |
| 1407 | +} |
| 1408 | + |
| 1409 | +// prepare output |
| 1410 | +std::vector<tc::DLTensorUPtr> |
| 1411 | +inferOutputTensorInfo(const std::string &tc, const std::string &entryPoint, |
| 1412 | + const std::vector<Tensor> &inputs) { |
| 1413 | + auto parsedTcs = tc::detail::parse(tc); |
| 1414 | + if (parsedTcs.count(entryPoint) != 1u) { |
| 1415 | + TC_CHECK_GE(parsedTcs.size(), 1u) |
| 1416 | + << "No TC was parsed, should have thrown earlier"; |
| 1417 | + throw tclang::ErrorReport(parsedTcs.begin()->second) |
| 1418 | + << "\nattempting to access undefined entryPoint: " << entryPoint; |
| 1419 | + } |
| 1420 | + auto inputDLTensors = makeDLConstTensors(inputs); |
| 1421 | + return makeDLTensorVector(tc::detail::inferOutputTensorInfo( |
| 1422 | + parsedTcs.at(entryPoint), extractRawPtrs(inputDLTensors))); |
| 1423 | +} |
| 1424 | + |
| 1425 | +std::vector<Tensor> prepareOutputs(const std::string &tc, |
| 1426 | + const std::string &entryPoint, |
| 1427 | + const std::vector<Tensor> &inputs) { |
| 1428 | + std::vector<Tensor> outputs; |
| 1429 | + auto outTensorInfo = inferOutputTensorInfo(tc, entryPoint, inputs); |
| 1430 | + if (outTensorInfo.size() == 0) { |
| 1431 | + return outputs; |
| 1432 | + } |
| 1433 | + TC_CHECK_GE(inputs.size(), 1u) |
| 1434 | + << "NYI: Need >= 1 input tensors to determine " |
| 1435 | + << "backend and prepare ATen outputs. Add an overload with just an ATen " |
| 1436 | + << "backend"; |
| 1437 | + |
| 1438 | + auto dev = inputs[0].device(); |
| 1439 | + auto dtype = inputs[0].data_type(); |
| 1440 | + for (size_t i = 0; i < outTensorInfo.size(); ++i) { |
| 1441 | + tc::TensorInfo info(outTensorInfo[i]); |
| 1442 | + Shape shape(info.shape.begin(), info.shape.end()); |
| 1443 | + |
| 1444 | + Tensor tmp(shape, dev, dtype); |
| 1445 | + outputs.push_back(tmp); |
| 1446 | + } |
| 1447 | + return outputs; |
| 1448 | +} |
| 1449 | +/// tc integration end |
| 1450 | + |
| 1451 | + |
1337 | 1452 | } // namespace singa |
0 commit comments