File indexing completed on 2025-01-18 10:11:10
0001 #ifndef TMVA_SOFIE_SOFIE_COMMON
0002 #define TMVA_SOFIE_SOFIE_COMMON
0003
0004 #include "TMVA/RTensor.hxx"
0005
0006 #include "ROOT/RSpan.hxx"
0007
0008 #include <stdexcept>
0009 #include <type_traits>
0010 #include <cstdint>
0011 #include <cstring>
0012 #include <string>
0013 #include <vector>
0014 #include <memory>
0015 #include <regex>
0016 #include <sstream>
0017 #include <iostream>
0018
0019 namespace TMVA{
0020 namespace Experimental{
0021 namespace SOFIE{
0022
0023
0024
0025 enum class ETensorType{
0026 UNDEFINED = 0, FLOAT = 1, UNINT8 = 2, INT8 = 3, UINT16 = 4, INT16 = 5, INT32 = 6, INT64 = 7, STRING = 8, BOOL = 9,
0027 FLOAT16 = 10, DOUBLE = 11, UINT32 = 12, UINT64 = 13, COMPLEX64 = 14, COMPLEX28 = 15, BFLOAT16 = 16
0028 };
0029
0030 typedef std::int64_t int_t;
0031
0032 std::string ConvertTypeToString(ETensorType type);
0033 ETensorType ConvertStringToType(std::string type);
0034
0035 struct Dim{
0036 bool isParam = false;
0037 size_t dim = 0;
0038 std::string param;
0039
0040
0041 Dim() {}
0042
0043
0044 Dim(const std::string & p, size_t d = 0) : isParam(true), dim(d), param(p) {}
0045
0046
0047 Dim(size_t d) : dim(d) {}
0048
0049 std::string GetVal() const {
0050 return (isParam) ? param : std::to_string(dim);
0051 }
0052 };
0053
0054
0055 struct InputTensorInfo{
0056 ETensorType type;
0057 std::vector<Dim> shape;
0058 };
0059
0060 struct TensorInfo{
0061 ETensorType type;
0062 std::vector<size_t> shape;
0063 };
0064
0065 struct DynamicTensorInfo{
0066 ETensorType type;
0067 std::vector<Dim> shape;
0068 };
0069
0070 std::vector<Dim> ConvertShapeToDim(std::vector<size_t> shape);
0071
0072 std::vector<size_t> ConvertShapeToInt(std::vector<Dim> shape);
0073
0074 std::size_t ConvertShapeToLength(std::vector<size_t> shape);
0075
0076 std::string ConvertShapeToString(std::vector<size_t> shape);
0077 std::string ConvertDynamicShapeToString(std::vector<Dim> shape);
0078
0079
0080
0081
0082 std::string ConvertDynamicShapeToLength(std::vector<Dim> shape);
0083
0084 class InitializedTensor {
0085 public:
0086 InitializedTensor() = default;
0087 InitializedTensor(ETensorType type, std::span<std::size_t> shape, std::shared_ptr<void> data)
0088 : fType{type}, fShape{shape.begin(), shape.end()}, fData{data}
0089 {
0090 }
0091
0092 ETensorType const &type() const { return fType; }
0093 std::vector<std::size_t> const &shape() const { return fShape; }
0094 std::shared_ptr<void> const &sharedptr() const { return fData; }
0095
0096 template <class T = void>
0097 T const *data() const
0098 {
0099 return static_cast<T const *>(fData.get());
0100 }
0101
0102 void CastSharedToPersistent()
0103 {
0104
0105
0106 fSize = 1;
0107 for (std::size_t item : fShape) {
0108 fSize *= static_cast<int>(item);
0109 }
0110 switch (fType) {
0111 case ETensorType::FLOAT: fSize *= sizeof(float); break;
0112 case ETensorType::DOUBLE: fSize *= sizeof(double); break;
0113 case ETensorType::INT32: fSize *= sizeof(int32_t); break;
0114 case ETensorType::INT64: fSize *= sizeof(int64_t); break;
0115 case ETensorType::BOOL: fSize *= sizeof(bool); break;
0116 default:
0117 throw std::runtime_error("TMVA::SOFIE doesn't yet supports serialising data-type " +
0118 ConvertTypeToString(fType));
0119 }
0120 fPersistentData = static_cast<char *>(fData.get());
0121 }
0122 void CastPersistentToShared()
0123 {
0124
0125 if (fSize == 0 || fPersistentData == nullptr) {
0126 return;
0127 }
0128
0129
0130 if (fPersistentData == static_cast<char *>(fData.get())) {
0131 return;
0132 }
0133
0134
0135 fData = std::shared_ptr<void>{malloc(fSize), free};
0136 std::memcpy(fData.get(), fPersistentData, fSize);
0137
0138
0139
0140 delete[] fPersistentData;
0141 fPersistentData = nullptr;
0142 fSize = 0;
0143 }
0144
0145 private:
0146 ETensorType fType;
0147 std::vector<std::size_t> fShape;
0148 std::shared_ptr<void> fData;
0149 int fSize = 0;
0150 char *fPersistentData = nullptr;
0151 };
0152
0153 template <typename T>
0154 ETensorType GetTemplatedType(T ){
0155 if (std::is_same<T, float>::value) return ETensorType::FLOAT;
0156 if (std::is_same<T, uint8_t>::value) return ETensorType::UNINT8;
0157 if (std::is_same<T, int8_t>::value) return ETensorType::INT8;
0158 if (std::is_same<T, uint16_t>::value) return ETensorType::UINT16;
0159 if (std::is_same<T, int16_t>::value) return ETensorType::INT16;
0160 if (std::is_same<T, int32_t>::value) return ETensorType::INT32;
0161 if (std::is_same<T, int64_t>::value) return ETensorType::INT64;
0162 if (std::is_same<T, std::string>::value) return ETensorType::STRING;
0163 if (std::is_same<T, bool>::value) return ETensorType::BOOL;
0164
0165 if (std::is_same<T, double>::value) return ETensorType::DOUBLE;
0166 if (std::is_same<T, uint32_t>::value) return ETensorType::UINT32;
0167 if (std::is_same<T, uint64_t>::value) return ETensorType::UINT64;
0168
0169 }
0170
0171 namespace UTILITY{
0172
0173 bool AreSameShape(const std::vector<size_t>&, const std::vector<size_t>&);
0174 bool AreSameShape(const std::vector<size_t>&, const std::vector<Dim>&);
0175 bool AreSameShape(const std::vector<Dim>&, const std::vector<Dim>&);
0176
0177
0178
0179 std::vector<size_t> MultidirectionalBroadcastShape(std::vector<std::vector<size_t>>);
0180
0181
0182 std::vector<size_t> UnidirectionalBroadcastShape(std::vector<size_t>, std::vector<size_t>);
0183
0184 std::string Clean_name(std::string input_tensor_name);
0185
0186 template<typename T>
0187 T* BroadcastConvBias(const T* data, const size_t channel, const std::vector<size_t>& targetShape) {
0188 size_t size = targetShape.size();
0189 if (targetShape[1] != channel) {
0190 std::stringstream ss;
0191 ss << "TMVA::SOFIE - Error broadcasting Conv Bias of shape {";
0192 ss << std::to_string(channel);
0193 ss << "} to ";
0194 ss << ConvertShapeToString(targetShape);
0195 throw
0196 std::runtime_error(ss.str());
0197 }
0198
0199 size_t targetLength = ConvertShapeToLength(targetShape);
0200 T* newData = new T[targetLength];
0201
0202 if (targetLength == channel) {
0203 std::copy(data, data + channel, newData);
0204 return newData;
0205 }
0206
0207
0208 size_t cStride = 1;
0209 for (size_t i = 2; i < size; i++)
0210 cStride *= targetShape[i];
0211
0212
0213 for (size_t i = 0; i < channel; i++) {
0214 std::fill(newData + i * cStride, newData + (i + 1) * cStride, data[i]);
0215 }
0216
0217 size_t batch = targetShape[0];
0218 size_t bStride = channel * cStride;
0219 for (size_t i = 1; i < batch; i++) {
0220 std::copy(newData, newData + bStride, newData + i * bStride);
0221 }
0222 return newData;
0223 }
0224
0225
0226
0227
0228 template<typename T>
0229 T* BroadcastTensor(const T* data, const std::vector<size_t>& shape, const std::vector<size_t>& targetShape) {
0230
0231 size_t size = shape.size();
0232
0233 size_t curLength = ConvertShapeToLength(shape);
0234 size_t targetLength = ConvertShapeToLength(targetShape);
0235
0236 T* broadcastedData = new T[targetLength];
0237 std::copy(data, data + curLength, broadcastedData);
0238
0239 size_t arrayNum = 1;
0240
0241 std::vector<T> newData(targetLength);
0242
0243 for (size_t idx = 0; idx < size; idx++) {
0244 size_t dim = shape[idx];
0245 size_t targetDim = targetShape[idx];
0246 if (dim == 1 && targetDim > 1) {
0247
0248 size_t newLength = curLength * targetDim;
0249
0250 size_t arrayLength = curLength / arrayNum;
0251
0252 if (arrayLength > 1) {
0253
0254 for (size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) {
0255 for (size_t targetIdx = 0; targetIdx < targetDim; targetIdx++) {
0256 size_t offset = arrayIdx * arrayLength * targetDim + targetIdx * arrayLength;
0257 std::copy(broadcastedData + arrayIdx * arrayLength,
0258 broadcastedData + (arrayIdx + 1) * arrayLength,
0259 newData.begin() + offset);
0260 }
0261 }
0262 } else {
0263
0264 for (size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) {
0265 std::fill(newData.begin() + arrayIdx * targetDim,
0266 newData.begin() + (arrayIdx + 1) * targetDim, broadcastedData[arrayIdx]);
0267 }
0268 }
0269
0270 curLength = newLength;
0271
0272 std::copy(newData.begin(), newData.begin() + newLength, broadcastedData);
0273 }
0274
0275 arrayNum *= targetDim;
0276 }
0277 return broadcastedData;
0278 }
0279
0280
0281 template<typename T>
0282 T* UnidirectionalBroadcast(const T* data, const std::vector<size_t>& shape, const std::vector<size_t>& targetShape) {
0283
0284 if (shape.size() < targetShape.size()) {
0285 size_t targetSize = targetShape.size();
0286 std::vector<size_t> newShape(targetSize, 1);
0287 size_t offset = targetSize - shape.size();
0288 std::copy(shape.begin(), shape.end(), newShape.begin() + offset);
0289 return BroadcastTensor<T>(data, newShape, targetShape);
0290 }
0291 return BroadcastTensor<T>(data, shape, targetShape);
0292 }
0293
0294
0295 std::vector<size_t> ComputeStrideFromShape(const std::vector<size_t> & shape);
0296 std::vector<Dim> ComputeStrideFromShape(const std::vector<Dim> & shape);
0297
0298
0299
0300 inline bool is_a_ge_zero_and_a_lt_b(int a, int b) {
0301 return static_cast<unsigned>(a) < static_cast<unsigned>(b);
0302 }
0303
0304
0305
0306
0307
0308
0309
0310
0311
0312
0313
0314
0315
0316
0317
0318
0319
0320
0321
0322
0323
0324 template <typename T>
0325 void Im2col(const T *data_im, const int channels, const int height, const int width, const int kernel_h,
0326 const int kernel_w, const int pad_h, const int pad_w, const int stride_h, const int stride_w,
0327 const int dilation_h, const int dilation_w, T *data_col)
0328 {
0329 const int output_h = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
0330 const int output_w = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
0331 const int channel_size = height * width;
0332 for (int channel = channels; channel--; data_im += channel_size) {
0333 for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
0334 for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
0335 int input_row = -pad_h + kernel_row * dilation_h;
0336 for (int output_rows = output_h; output_rows; output_rows--) {
0337 if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
0338 for (int output_cols = output_w; output_cols; output_cols--) {
0339 *(data_col++) = 0;
0340 }
0341 } else {
0342 int input_col = -pad_w + kernel_col * dilation_w;
0343 for (int output_col = output_w; output_col; output_col--) {
0344 if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
0345 *(data_col++) = data_im[input_row * width + input_col];
0346 } else {
0347 *(data_col++) = 0;
0348 }
0349 input_col += stride_w;
0350 }
0351 }
0352 input_row += stride_h;
0353 }
0354 }
0355 }
0356 }
0357 }
0358
0359
0360 template <typename T>
0361 void Im2col_3d(const T *data_im, const int channels,
0362 const int depth, const int height, const int width,
0363 const int kernel_d, const int kernel_h, const int kernel_w,
0364 const int pad_d, const int pad_h, const int pad_w,
0365 const int stride_d, const int stride_h, const int stride_w,
0366 const int dilation_d, const int dilation_h, const int dilation_w, T *data_col)
0367 {
0368 const int output_h = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
0369 const int output_w = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
0370 const int output_d = (depth + 2 * pad_d - (dilation_d * (kernel_d - 1) + 1)) / stride_d + 1;
0371 const int channel_size = height * width * depth;
0372
0373 for (int channel = channels; channel--; data_im += channel_size) {
0374 for (int kernel_depth = 0; kernel_depth < kernel_d; kernel_depth++) {
0375 for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
0376 for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
0377 int input_dep = -pad_d + kernel_depth * dilation_d;
0378 for (int output_dep = output_d; output_dep; output_dep--) {
0379 if (!is_a_ge_zero_and_a_lt_b(input_dep, depth)) {
0380 for (int output_rows = output_h; output_rows; output_rows--) {
0381 for (int output_cols = output_w; output_cols; output_cols--) {
0382 *(data_col++) = 0;
0383 }
0384 }
0385 } else {
0386 int input_row = -pad_h + kernel_row * dilation_h;
0387 for (int output_rows = output_h; output_rows; output_rows--) {
0388 if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
0389 for (int output_cols = output_w; output_cols; output_cols--) {
0390 *(data_col++) = 0;
0391 }
0392 } else {
0393 int input_col = -pad_w + kernel_col * dilation_w;
0394 for (int output_col = output_w; output_col; output_col--) {
0395 if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
0396 *(data_col++) = data_im[input_dep * width * height + input_row * width + input_col];
0397 } else {
0398 *(data_col++) = 0;
0399 }
0400 input_col += stride_w;
0401 }
0402 }
0403 input_row += stride_h;
0404 }
0405 }
0406 input_dep += stride_d;
0407 }
0408 }
0409 }
0410 }
0411 }
0412 }
0413
0414 template <typename Dtype>
0415 void col2im(const Dtype* data_col, const int channels,
0416 const int height, const int width, const int kernel_h, const int kernel_w,
0417 const int pad_h, const int pad_w,
0418 const int stride_h, const int stride_w,
0419 const int dilation_h, const int dilation_w,
0420 Dtype* data_im) {
0421
0422 std::fill(data_im, data_im + height * width * channels, 0.);
0423
0424
0425
0426 const int output_h = (height + 2 * pad_h -
0427 (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
0428 const int output_w = (width + 2 * pad_w -
0429 (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
0430 const int channel_size = height * width;
0431 for (int channel = channels; channel--; data_im += channel_size) {
0432 for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
0433 for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
0434 int input_row = -pad_h + kernel_row * dilation_h;
0435 for (int output_rows = output_h; output_rows; output_rows--) {
0436 if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
0437 data_col += output_w;
0438 } else {
0439 int input_col = -pad_w + kernel_col * dilation_w;
0440 for (int output_col = output_w; output_col; output_col--) {
0441 if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
0442
0443
0444
0445
0446
0447 data_im[input_row * width + input_col] += *data_col;
0448 }
0449 data_col++;
0450 input_col += stride_w;
0451 }
0452 }
0453 input_row += stride_h;
0454 }
0455 }
0456 }
0457 }
0458
0459 }
0460
0461
0462
0463 }
0464
0465 namespace BLAS{
0466 extern "C" void sgemm_(const char * transa, const char * transb, const int * m, const int * n, const int * k,
0467 const float * alpha, const float * A, const int * lda, const float * B, const int * ldb,
0468 const float * beta, float * C, const int * ldc);
0469 }
0470
0471
0472 struct GNN_Data {
0473 RTensor<float> node_data;
0474 RTensor<float> edge_data;
0475 RTensor<float> global_data;
0476 RTensor<int> edge_index;
0477
0478
0479
0480
0481 GNN_Data(): node_data(RTensor<float>({})), edge_data(RTensor<float>({})), global_data(RTensor<float>({})), edge_index(RTensor<int>({})) {}
0482
0483 };
0484
0485 template<typename T>
0486 TMVA::Experimental::RTensor<T> Concatenate( TMVA::Experimental::RTensor<T> & t1, TMVA::Experimental::RTensor<T> & t2, int axis = 0)
0487 {
0488
0489 if (t1.GetMemoryLayout() != t2.GetMemoryLayout())
0490 throw std::runtime_error("TMVA RTensor Concatenate - tensors have different memory layout");
0491 auto & shape1 = t1.GetShape();
0492 auto & shape2 = t2.GetShape();
0493 if (t1.GetSize()/shape1[axis] != t2.GetSize()/shape2[axis]) {
0494 std::cout << "axis " << axis << " sizes " << t1.GetSize() << " " << t2.GetSize() << " ";
0495 std::cout << "shape 1 : " << ConvertShapeToString(t1.GetShape());
0496 std::cout << " shape 2 : " << ConvertShapeToString(t2.GetShape()) << std::endl;
0497 throw std::runtime_error("TMVA RTensor Concatenate - tensors have incompatible shapes");
0498 }
0499 std::vector<size_t> outShape = shape1;
0500 outShape[axis] = shape1[axis] + shape2[axis];
0501 TMVA::Experimental::RTensor<T> tout(outShape, t1.GetMemoryLayout());
0502 if (t1.GetMemoryLayout() == TMVA::Experimental::MemoryLayout::ColumnMajor) {
0503 throw std::runtime_error("TMVA RTensor Concatenate is not yet supported for column major tensors");
0504 }
0505
0506 auto & stride1 = t1.GetStrides();
0507 auto & stride2 = t2.GetStrides();
0508 auto & outStride = tout.GetStrides();
0509
0510 size_t s1 = (axis > 0) ? stride1[axis-1] : t1.GetSize();
0511 size_t s2 = (axis > 0) ? stride2[axis-1] : t2.GetSize();
0512 size_t sout = (axis > 0) ? outStride[axis-1] : tout.GetSize();
0513 size_t nb = t1.GetSize()/s1;
0514 for (size_t i = 0; i < nb; i++) {
0515 std::copy(t1.GetData() + i*s1, t1.GetData() + (i+1)*s1, tout.GetData() + i * sout );
0516 std::copy(t2.GetData() + i*s2, t2.GetData() + (i+1)*s2, tout.GetData() + i * sout + s1 );
0517 }
0518
0519 return tout;
0520 }
0521
0522
0523 inline GNN_Data Concatenate(GNN_Data & data1, GNN_Data & data2, int axis = 0) {
0524 GNN_Data out;
0525 out.node_data = Concatenate(data1.node_data,data2.node_data, axis);
0526 out.edge_data = Concatenate(data1.edge_data,data2.edge_data, axis);
0527 out.global_data = Concatenate<float>(data1.global_data,data2.global_data, axis-1);
0528
0529 out.edge_index = data1.edge_index.Copy();
0530 return out;
0531 }
0532
0533 inline GNN_Data Copy(const GNN_Data & data) {
0534 GNN_Data out;
0535 out.node_data = RTensor<float>(data.node_data.GetShape());
0536 out.edge_data = RTensor<float>(data.edge_data.GetShape());
0537 out.global_data = RTensor<float>(data.global_data.GetShape());
0538 out.edge_index = RTensor<int>(data.edge_index.GetShape());
0539 std::copy(data.node_data.GetData(), data.node_data.GetData()+ data.node_data.GetSize(), out.node_data.GetData());
0540 std::copy(data.edge_data.GetData(), data.edge_data.GetData()+ data.edge_data.GetSize(), out.edge_data.GetData());
0541 std::copy(data.global_data.GetData(), data.global_data.GetData()+ data.global_data.GetSize(), out.global_data.GetData());
0542 std::copy(data.edge_index.GetData(), data.edge_index.GetData()+ data.edge_index.GetSize(), out.edge_index.GetData());
0543 return out;
0544 }
0545
0546 }
0547 }
0548 }
0549
0550 #endif