Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-09-20 09:26:28

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 <complex>
0013 #include <string>
0014 #include <vector>
0015 #include <map>
0016 #include <memory>
0017 #include <regex>
0018 #include <sstream>
0019 #include <iostream>
0020 #include <iomanip>
0021 #include <cassert>
0022 #include <limits>
0023 
0024 namespace TMVA {
0025 namespace Experimental {
0026 namespace SOFIE {
0027 
0028 enum class ETensorType{
0029    UNDEFINED = 0, FLOAT = 1, UINT8 = 2, INT8 = 3, UINT16 = 4, INT16 = 5, INT32 = 6, INT64 = 7, STRING = 8, BOOL = 9, //order sensitive
0030     FLOAT16 = 10, DOUBLE = 11, UINT32 = 12, UINT64 = 13, COMPLEX64 = 14, COMPLEX28 = 15, BFLOAT16 = 16
0031 };
0032 
0033 enum class EActivationType{
0034    UNDEFINED = 0, RELU = 1, SOFTMAX = 2, SIGMOID = 3, LEAKYRELU = 4, TANH = 5, ELU = 6
0035 };
0036 
0037 constexpr size_t GetTypeSize(ETensorType type) {
0038     switch (type) {
0039         case ETensorType::FLOAT:     return sizeof(float);
0040         case ETensorType::DOUBLE:    return sizeof(double);
0041         case ETensorType::UINT8:     return sizeof(uint8_t);
0042         case ETensorType::INT8:      return sizeof(int8_t);
0043         case ETensorType::UINT16:    return sizeof(uint16_t);
0044         case ETensorType::INT16:     return sizeof(int16_t);
0045         case ETensorType::INT32:     return sizeof(int32_t);
0046         case ETensorType::INT64:     return sizeof(int64_t);
0047         case ETensorType::UINT32:    return sizeof(uint32_t);
0048         case ETensorType::UINT64:    return sizeof(uint64_t);
0049         case ETensorType::BOOL:      return sizeof(bool);
0050         case ETensorType::STRING:    return sizeof(std::string);
0051         default: return 0;
0052     }
0053 }
0054 
0055 typedef std::int64_t int_t;
0056 
0057 std::string ConvertTypeToString(ETensorType type);
0058 ETensorType ConvertStringToType(std::string type);
0059 
0060 // find if a string represents a number
0061 bool IsInteger(const std::string & s);
0062 
0063 struct Dim{
0064    bool isParam = false;
0065    size_t dim = 0;
0066    std::string param;
0067 
0068     // default constructor (for I/O)
0069    Dim() {}
0070 
0071    // constructor for a parametric dimension with the option to pass a default dim value
0072    // We use -1 for dim to indicate that the param dimension is an expression (e.g. "d1+d2")
0073    // in case the string represents a number make Dim not parametric
0074    Dim(const std::string & p, size_t d = 0) : isParam(true), dim(d), param(p)
0075    {
0076       if (IsInteger(p)) {
0077             isParam = false;
0078             dim = std::stoi(p);
0079       }
0080    }
0081 
0082    // constructor for a non-parametric dimension
0083    Dim(size_t d) : dim(d) {}
0084 
0085    std::string GetVal() const {
0086       // cast to int64_t for negative shape values
0087       return (isParam) ? param : std::to_string(static_cast<int64_t>(dim));
0088    }
0089 
0090    std::ostream& operator<< (std::ostream& os) const {
0091       os << GetVal();
0092       return os;
0093    }
0094 
0095    bool operator==(const Dim& rhs) const {
0096        return (isParam && rhs.isParam) ? param == rhs.param : dim == rhs.dim;
0097    }
0098    bool operator!=(const Dim& rhs) const {
0099        return !(*this == rhs);
0100    }
0101 };
0102 
0103 //bool operator==(const Dim& lhs, const Dim& rhs);
0104 inline std::ostream & operator<< (std::ostream &os, const Dim &d) {
0105    os << d.GetVal();
0106    return os;
0107 }
0108 
0109 struct InputTensorInfo{
0110    ETensorType type;
0111    std::vector<Dim> shape;
0112 };
0113 
0114 struct TensorInfo{
0115    ETensorType type;
0116    std::vector<size_t> shape;
0117 };
0118 
0119 struct DynamicTensorInfo{
0120    ETensorType type;
0121    std::vector<Dim> shape;
0122 };
0123 
0124 // template traits for Tensor Shape
0125 template <typename T>
0126 struct TensorShape {};
0127 template<>
0128 struct TensorShape<Dim> {
0129    static bool IsDim() { return true; }
0130 };
0131 template<>
0132 struct TensorShape<size_t> {
0133    static bool IsDim() { return false; }
0134 };
0135 
0136 // template traits for Tensor type
0137 template <typename T>
0138 struct TensorType {};
0139 template<>
0140 struct TensorType<float> {
0141    static const std::string Name() { return "float"; }
0142 };
0143 template<>
0144 struct TensorType<double> {
0145    static const std::string Name() { return "double"; }
0146 };
0147 template<>
0148 struct TensorType<int64_t> {
0149    static const std::string Name() { return "int64_t"; }
0150 };
0151 template<>
0152 struct TensorType<int32_t> {
0153    static const std::string Name() { return "int32_t"; }
0154 };
0155 template<>
0156 struct TensorType<uint32_t> {
0157    static const std::string Name() { return "uint32_t"; }
0158 };
0159 template<>
0160 struct TensorType<uint64_t> {
0161    static const std::string Name() { return "uint64_t"; }
0162 };
0163 template<>
0164 struct TensorType<bool> {
0165    static const std::string Name() { return "bool"; }
0166 };
0167 template<>
0168 struct TensorType<int8_t> {
0169    static const std::string Name() { return "int8_t"; }
0170 };
0171 template<>
0172 struct TensorType<uint8_t> {
0173    static const std::string Name() { return "uint8_t"; }
0174 };
0175 
0176 struct TensorMemoryInfo {
0177    std::string_view tensor_name;
0178    size_t tensor_size;
0179 
0180    TensorMemoryInfo split(const std::string_view new_name, size_t new_size) {
0181         if (new_size > tensor_size) {
0182             throw std::invalid_argument("New size exceeds available tensor size.");
0183         }
0184         tensor_size -= new_size;
0185         return TensorMemoryInfo{new_name, new_size};
0186    }
0187 
0188     // Method to merge another struct into this one
0189    void merge(const TensorMemoryInfo& other) {
0190         tensor_size += other.tensor_size;
0191    }
0192 };
0193 
0194 struct MemoryPoolInfo {
0195 
0196    // ordered map with chunk_idx as key and TensorMemoryInfo as value
0197    std::map<size_t, TensorMemoryInfo> total_stack;
0198 
0199    // ordered map with chunk_idx as key and chunk_size as value
0200    std::map<size_t, size_t> available_stack;
0201 };
0202 
0203 std::vector<Dim> ConvertShapeToDim(const std::vector<size_t> & shape);
0204 
0205 std::vector<size_t> ConvertShapeToInt(const std::vector<Dim> & shape);
0206 
0207 std::size_t ConvertShapeToLength(const std::vector<size_t> & shape);
0208 
0209 std::string ConvertShapeToString(const std::vector<size_t> & shape);
0210 std::string ConvertDimShapeToString(const std::vector<Dim> & shape);
0211 
0212 std::string ConvertDimShapeToLength(const std::vector<Dim> & shape);
0213 
0214 
0215 template<class T>
0216 std::string ConvertValToString(T value) {
0217    std::stringstream ret;
0218    if (std::is_floating_point_v<T>)
0219       ret << std::setprecision(std::numeric_limits<T>::max_digits10);
0220    ret << value;
0221    return ret.str();
0222 }
0223 
0224 
0225 // convert list of values in a string taking into account the precision
0226 template<class T>
0227 std::string ConvertValuesToString(size_t n, const T * data, size_t maxprint = -1) {
0228    std::stringstream ret;
0229    ret << "{ ";
0230    for (size_t i = 0; i < std::min(n,maxprint); i++) {
0231       if (std::is_floating_point_v<T>)
0232          ret << std::setprecision(std::numeric_limits<T>::max_digits10) << data[i];
0233       else
0234          // cast in case of boolean (int8)
0235          ret << data[i];
0236 
0237       if (i < n-1) ret << ", ";
0238       if (i < n-1 && i == maxprint-1) ret << "..... ";
0239    }
0240    ret << "}";
0241    return ret.str();
0242 }
0243 template<class T>
0244 std::string ConvertValuesToString(const std::vector<T> & data, size_t maxprint = 5) {
0245   return ConvertValuesToString(data.size(), data.data(), maxprint);
0246 }
0247 
0248 class InitializedTensor {
0249 public:
0250    InitializedTensor() = default;
0251    InitializedTensor(ETensorType type, std::span<std::size_t> shape, std::shared_ptr<void> data, bool typeConstant = false)
0252       : fConstant(typeConstant), fType{type}, fShape{shape.begin(), shape.end()}, fData{data}
0253    {
0254    }
0255 
0256    ETensorType const &type() const { return fType; }
0257    std::vector<std::size_t> const &shape() const { return fShape; }
0258    std::shared_ptr<void> const &sharedptr() const { return fData; }
0259    // query if tensor comes from a Constant operator
0260    bool IsConstantTensor() const { return fConstant;}
0261    // query if tensor needs to be written in a weight file. Constant tensors are not written in a separate file
0262    bool IsWeightTensor() const { return !fConstant && !fIsNotWritable;}
0263    // check if a Tensor is Writable (need to be written in the file or in the generated code (e.g. as a constant tensor)
0264    // if an initialized tensors is used in a constant operator at compile time does not need to be written and can be omitted in
0265    // the generated code
0266    bool IsNotWritable() const { return fIsNotWritable; }
0267    // set not writable initialized tensors - i.e. tensor that must not be written in a file
0268    void SetNotWritable() { fIsNotWritable = true;}
0269    // set writable initialized tensors - i.e. tensor that must be written in a file
0270    void SetWritable() { fIsNotWritable = false;}
0271    // set as constant (needed for non-float initialized tensors)
0272    void SetConstant() { fConstant = true;}
0273 
0274    template <class T = void>
0275    T const *data() const
0276    {
0277       return static_cast<T const *>(fData.get());
0278    }
0279 
0280    void CastSharedToPersistent()
0281    {
0282       // We only calculate fSize here, because it is only used for IO to know
0283       // the size of the persistent data.
0284       fSize = 1;
0285       for (std::size_t item : fShape) {
0286          fSize *= static_cast<int>(item);
0287       }
0288       // get size in bytes
0289       fSize *= GetTypeSize(fType);
0290       fPersistentData = static_cast<char *>(fData.get());
0291    }
0292    void CastPersistentToShared()
0293    {
0294       // If there is no persistent data, do nothing
0295       if (fSize == 0 || fPersistentData == nullptr) {
0296          return;
0297       }
0298 
0299       // Nothing to be done if the pointed-to data is the same
0300       if (fPersistentData == static_cast<char *>(fData.get())) {
0301          return;
0302       }
0303 
0304       // Initialize the shared_ptr
0305       fData = std::shared_ptr<void>{malloc(fSize), free};
0306       std::memcpy(fData.get(), fPersistentData, fSize);
0307 
0308       // Make sure the data read from disk doesn't leak and delete the
0309       // persistent data
0310       delete[] fPersistentData;
0311       fPersistentData = nullptr;
0312       fSize = 0;
0313    }
0314 
0315 private:
0316    bool  fConstant = false;      ///< Flag specifying if tensor is a Constant one (coming from a Constant operator)
0317    bool  fIsNotWritable = false; ///< Flag to indicate that tensor values do not need to be written as weight or generated code
0318    ETensorType fType;               ///< Encodes the type of the data
0319    std::vector<std::size_t> fShape; ///< The shape of the data in terms of elements in each dimension
0320    std::shared_ptr<void> fData;     ///<! Transient shared data
0321    int fSize = 0;                   ///< The size of the persistent data in bytes (not number of elements!)
0322    char *fPersistentData = nullptr; ///<[fSize] Persistent version of the data
0323 };
0324 
0325 template <typename T>
0326 ETensorType GetTemplatedType(T /*obj*/ ){
0327    if (std::is_same<T, float>::value) return ETensorType::FLOAT;
0328    if (std::is_same<T, uint8_t>::value) return ETensorType::UINT8;
0329    if (std::is_same<T, int8_t>::value) return ETensorType::INT8;
0330    if (std::is_same<T, uint16_t>::value) return ETensorType::UINT16;
0331    if (std::is_same<T, int16_t>::value) return ETensorType::INT16;
0332    if (std::is_same<T, int32_t>::value) return ETensorType::INT32;
0333    if (std::is_same<T, int64_t>::value) return ETensorType::INT64;
0334    if (std::is_same<T, std::string>::value) return ETensorType::STRING;
0335    if (std::is_same<T, bool>::value) return ETensorType::BOOL;
0336    //float16 unimplemented
0337    if (std::is_same<T, double>::value) return ETensorType::DOUBLE;
0338    if (std::is_same<T, uint32_t>::value) return ETensorType::UINT32;
0339    if (std::is_same<T, uint64_t>::value) return ETensorType::UINT64;
0340    //complex 64, 28, bfloat 16 unimplemented
0341 }
0342 
0343 namespace UTILITY{
0344 
0345 
0346 
0347 // clean operator and tensor names
0348 std::string Clean_name(std::string input_tensor_name);
0349 
0350 // Check if two shapes are equal
0351 bool AreSameShape(const std::vector<size_t>&, const std::vector<size_t>&);
0352 bool AreSameShape(const std::vector<size_t>&, const std::vector<Dim>&);
0353 bool AreSameShape(const std::vector<Dim>&, const std::vector<Dim>&);
0354 
0355 
0356 // Multidirectional broadcast a list of tensors to the same shape
0357 std::vector<size_t> MultidirectionalBroadcastShape(std::vector<std::vector<size_t>>);
0358 
0359 // Multidirectional broadcast two shapes to the same shape
0360 
0361 std::pair<int, std::vector<size_t>> MultidirectionalBroadcastShape(std::vector<size_t> &, std::vector<size_t> &);
0362 std::vector<size_t> UnidirectionalBroadcastShape(std::vector<size_t> &, std::vector<size_t> &);
0363 
0364 std::pair<int, std::vector<Dim>> MultidirectionalBroadcastShape(std::vector<Dim> &, std::vector<Dim> &);
0365 
0366 
0367 
0368 template<typename T>
0369 T* BroadcastConvBias(const T* data, const size_t channel, const std::vector<size_t>& targetShape) {
0370    size_t size = targetShape.size();
0371    if (targetShape[1] != channel) {
0372       std::stringstream ss;
0373       ss << "TMVA::SOFIE - Error broadcasting Conv Bias of shape {";
0374       ss << std::to_string(channel);
0375       ss << "} to ";
0376       ss << ConvertShapeToString(targetShape);
0377       throw
0378          std::runtime_error(ss.str());
0379    }
0380 
0381    size_t targetLength = ConvertShapeToLength(targetShape);
0382    T* newData = new T[targetLength];
0383 
0384    if (targetLength == channel) {
0385       std::copy(data, data + channel, newData);
0386       return newData;
0387    }
0388 
0389    // cStride = OutDepth * outHeight * outWidth
0390    size_t cStride = 1;
0391    for (size_t i = 2; i < size; i++)
0392       cStride *= targetShape[i];
0393    // Broadcast each element of the bias to a vector of size cStride and concatenate them
0394    // into a vector of size channel * cStride
0395    for (size_t i = 0; i < channel; i++) {
0396       std::fill(newData + i * cStride, newData + (i + 1) * cStride, data[i]);
0397    }
0398    // Broadcast newData[0...channel * cStride) to newData[0...batch * channel * cStride)
0399    size_t batch = targetShape[0];
0400    size_t bStride = channel * cStride;
0401    for (size_t i = 1; i < batch; i++) {
0402       std::copy(newData, newData + bStride, newData + i * bStride);
0403    }
0404    return newData;
0405 }
0406 
0407 // Broadcast a tensor from shape to targetShape according to numpy broadcasting rules
0408 // See more at https://numpy.org/doc/stable/user/basics.broadcasting.html
0409 // and https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md .
0410 template<typename T, class ConstContT = std::span<const T>>
0411 void BroadcastTensor(ConstContT data, const std::vector<size_t>& shape, const std::vector<size_t>& targetShape, T *broadcastedData) {
0412    // Size of the shapes (tensor input here have shapes with same sizes, we have already added the needed ones )
0413    size_t size = shape.size();
0414    // Current length of the broadcasted tensor
0415    size_t curLength = data.size();
0416    // special case when broadcasting last dimensions (initial shapes must be the same)
0417    if (size > 1 && shape.front() == targetShape.front() && shape.back() == 1) {
0418       size_t bsize = targetShape.back();
0419       // compute the size of the data to broadcast
0420       for (int k = int(size)-2; k >=0; k--) {
0421          if (shape[k] != 1) break;
0422          bsize *= targetShape[k];
0423       }
0424       for (size_t i = 0; i < curLength; i++) {
0425          std::fill(broadcastedData + i*bsize, broadcastedData + (i+1)*bsize , data[i]);
0426       }
0427       return;
0428    }
0429 
0430    std::copy(data.begin(), data.end(), broadcastedData);
0431    // Product of the previous dimensions of targetShape
0432    size_t arrayNum = 1;
0433    // New broadcasted data: is this needed?
0434    std::vector<T> newData(ConvertShapeToLength(targetShape));
0435 
0436    for (size_t idx = 0; idx < size; idx++) {
0437       size_t dim = shape[idx];
0438       size_t targetDim = targetShape[idx];
0439       if (dim == 1 && targetDim > 1) {
0440          // Set the new length of the data
0441          size_t newLength = curLength * targetDim;
0442          // View the data as a list of arrayNum arrays of size arrayLength
0443          size_t arrayLength = curLength / arrayNum;
0444          // Broadcast each array dim times
0445          if (arrayLength > 1) {
0446             // If each array has at least two elements
0447             for (size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) {
0448                for (size_t targetIdx = 0; targetIdx < targetDim; targetIdx++) {
0449                   size_t offset = arrayIdx * arrayLength * targetDim + targetIdx * arrayLength;
0450                   std::copy(broadcastedData + arrayIdx * arrayLength,
0451                      broadcastedData + (arrayIdx + 1) * arrayLength,
0452                      newData.begin() + offset);
0453                }
0454             }
0455          } else {
0456             // If each array has one element
0457             for (size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) {
0458                std::fill(newData.begin() + arrayIdx * targetDim,
0459                   newData.begin() + (arrayIdx + 1) * targetDim, broadcastedData[arrayIdx]);
0460             }
0461          }
0462          // Update current length
0463          curLength = newLength;
0464          // Update broadcasted data
0465          std::copy(newData.begin(), newData.begin() + newLength, broadcastedData);
0466       }
0467       // Update the number of arrays
0468       arrayNum *= targetDim;
0469    }
0470 }
0471 
0472 // interface where we allocate a new array for broadcasted data
0473 template<typename T>
0474 T* CreateBroadcastTensor(const T* data, const std::vector<size_t>& shape, const std::vector<size_t>& targetShape, size_t targetLength) {
0475    // newShape is an array of size equal to dimension along which we are broadcasting the tensor
0476    T* broadcastedData = new T[targetLength];
0477    size_t curLength = ConvertShapeToLength(shape);
0478    BroadcastTensor<T>({data, curLength}, shape, targetShape, broadcastedData);
0479    return broadcastedData;
0480 }
0481 // Unidirectional broadcasting shape to targetShape// In unidirectional broadcast - only tensor B can have the shape changed not
0482 // tensor A - otherwise is a multidirectional broadcast
0483 template<typename T>
0484 T* UnidirectionalBroadcast(const T* data, const std::vector<size_t>& shape, const std::vector<size_t>& targetShape) {
0485    // Prepend shape with ones
0486    if (shape.size() < targetShape.size()) {
0487       size_t targetSize = targetShape.size();
0488       std::vector<size_t> newShape(targetSize, 1);
0489       size_t offset = targetSize - shape.size();
0490       std::copy(shape.begin(), shape.end(), newShape.begin() + offset);
0491       return CreateBroadcastTensor(data, newShape, targetShape, ConvertShapeToLength(targetShape));
0492    }
0493    return CreateBroadcastTensor(data, shape, targetShape, ConvertShapeToLength(targetShape));
0494 }
0495 
0496 // Unidirectional broadcasting shape to targetShape using a passed vector to avoid allocations
0497 template<typename T>
0498 void UnidirectionalBroadcast(const T* data, const std::vector<size_t>& shape, const std::vector<size_t>& targetShape, T *broadcastedData) {
0499    size_t curLength = ConvertShapeToLength(shape);
0500    std::span<T> inData(const_cast<T*>(data), curLength);
0501    // Prepend shape with ones
0502    if (shape.size() < targetShape.size()) {
0503       size_t targetSize = targetShape.size();
0504       std::vector<size_t> newShape(targetSize, 1);
0505       size_t offset = targetSize - shape.size();
0506       std::copy(shape.begin(), shape.end(), newShape.begin() + offset);
0507       BroadcastTensor(inData, newShape, targetShape, broadcastedData);
0508    }
0509    BroadcastTensor(inData, shape, targetShape, broadcastedData);
0510 }
0511 
0512 /// compute stride of a tensor given its shape (assume layout is row-major)
0513 std::vector<size_t> ComputeStrideFromShape(const std::vector<size_t> & shape);
0514 std::vector<Dim> ComputeStrideFromShape(const std::vector<Dim> & shape);
0515 
0516 /// function to check if a >> 0 and a < MAX using a single comparison
0517 //// use trick casting to unsigned values so it becomes a single comparison
0518 inline bool is_a_ge_zero_and_a_lt_b(int a, int b) {
0519    return static_cast<unsigned>(a) < static_cast<unsigned>(b);
0520 }
0521 
0522 
0523 /// im2col : efficient function to re-arrange input data of convolution to a matrix
0524 /// that can be used by BLAS
0525 /// Use trick to loop on each element of filtered region first and follow input data layout
0526 /// By doing this reads and writes are of consecutive data in memory and one gains in efficiency
0527 /// The resulting matrix will be already transposed and can be used directly in BLAS
0528 /// since output will be a matrix : (channels*kernel_h*kernel_w , output_h*output_w)
0529 /// Example: with an input matrix
0530 ///    a1 a2 a3
0531 ///    b1 b2 b3    and a 2x2 kernel    (k1,k2,k3,k4) and padding 1 :
0532 ///    c1 c2 c3
0533 ///     outpout will be a matrix (4 x 16)
0534 ///  the routine will follow output order :
0535 //     first all elements which will be operated by k1 then k2 then k3
0536 ///  -> ( 0  0  0  0  0  a1 a2 a3 0  b1 b2 b3  0 c1 c2 c3  )    all elements for k1
0537 ///     ( 0  0  0  0  a1 a2 a3  0 b1 b2 b3  0 c1 c2 c3  0  )     for k2
0538 ///     ( 0  a1 a2 a3 0  b1 b2 b3 0  c1 c2 c3  0  0  0  0  )     for k3
0539 ///     ( a1 a2 a3 0  b1 b2 b3  0 c1 c2 c3  0  0  0  0  0  )     for k4
0540 ///
0541 
0542 template <typename T>
0543 void Im2col(const T *data_im, const int channels, const int height, const int width, const int kernel_h,
0544                 const int kernel_w, const int pad_h, const int pad_w, const int stride_h, const int stride_w,
0545                 const int dilation_h, const int dilation_w, T *data_col)
0546 {
0547    const int output_h = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
0548    const int output_w = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
0549    const int channel_size = height * width;
0550    for (int channel = channels; channel--; data_im += channel_size) {
0551       for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
0552          for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
0553             int input_row = -pad_h + kernel_row * dilation_h;
0554             for (int output_rows = output_h; output_rows; output_rows--) {
0555                if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
0556                   for (int output_cols = output_w; output_cols; output_cols--) {
0557                      *(data_col++) = 0;
0558                   }
0559                } else {
0560                   int input_col = -pad_w + kernel_col * dilation_w;
0561                   for (int output_col = output_w; output_col; output_col--) {
0562                      if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
0563                         *(data_col++) = data_im[input_row * width + input_col];
0564                      } else {
0565                         *(data_col++) = 0;
0566                      }
0567                      input_col += stride_w;
0568                   }
0569                }
0570                input_row += stride_h;
0571             }
0572          }
0573       }
0574    }
0575 }
0576 
0577 /// 3d implementation
0578 template <typename T>
0579 void Im2col_3d(const T *data_im, const int channels,
0580             const int depth, const int height, const int width,
0581             const int kernel_d, const int kernel_h, const int kernel_w,
0582             const int pad_d, const int pad_h, const int pad_w,
0583             const int stride_d, const int stride_h, const int stride_w,
0584             const int dilation_d, const int dilation_h,  const int dilation_w, T *data_col)
0585 {
0586    const int output_h = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
0587    const int output_w = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
0588    const int output_d = (depth + 2 * pad_d - (dilation_d * (kernel_d - 1) + 1)) / stride_d + 1;
0589    const int channel_size = height * width * depth;
0590    // assume data are c x d x h x w
0591    for (int channel = channels; channel--; data_im += channel_size) {
0592       for (int kernel_depth = 0; kernel_depth < kernel_d; kernel_depth++) {
0593          for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
0594             for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
0595                int input_dep = -pad_d + kernel_depth * dilation_d;
0596                for (int output_dep = output_d; output_dep; output_dep--) {
0597                   if (!is_a_ge_zero_and_a_lt_b(input_dep, depth)) {
0598                      for (int output_rows = output_h; output_rows; output_rows--) {
0599                         for (int output_cols = output_w; output_cols; output_cols--) {
0600                            *(data_col++) = 0;
0601                         }
0602                      }
0603                   } else {
0604                      int input_row = -pad_h + kernel_row * dilation_h;
0605                      for (int output_rows = output_h; output_rows; output_rows--) {
0606                         if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
0607                            for (int output_cols = output_w; output_cols; output_cols--) {
0608                               *(data_col++) = 0;
0609                            }
0610                         } else {
0611                            int input_col = -pad_w + kernel_col * dilation_w;
0612                            for (int output_col = output_w; output_col; output_col--) {
0613                               if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
0614                                  *(data_col++) = data_im[input_dep * width * height + input_row * width + input_col];
0615                               } else {
0616                                  *(data_col++) = 0;
0617                               }
0618                               input_col += stride_w;
0619                            }
0620                         }
0621                         input_row += stride_h;
0622                      }
0623                   }
0624                   input_dep += stride_d;
0625                }
0626             }
0627          }
0628       }
0629    }
0630 }
0631 
0632 template <typename Dtype>
0633 void col2im(const Dtype* data_col, const int channels,
0634     const int height, const int width, const int kernel_h, const int kernel_w,
0635     const int pad_h, const int pad_w,
0636     const int stride_h, const int stride_w,
0637     const int dilation_h, const int dilation_w,
0638     Dtype* data_im) {
0639    // note that output data_im needs to be set to zero value!!!!
0640    std::fill(data_im, data_im + height * width * channels, 0.);
0641   //caffe_set(height * width * channels, Dtype(0), data_im);
0642   // data_im must be a zero vector
0643   //const Dtype * data_col_0 = data_col;
0644   const int output_h = (height + 2 * pad_h -
0645     (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
0646   const int output_w = (width + 2 * pad_w -
0647     (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
0648   const int channel_size = height * width;
0649   for (int channel = channels; channel--; data_im += channel_size) {
0650     for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
0651       for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
0652         int input_row = -pad_h + kernel_row * dilation_h;
0653         for (int output_rows = output_h; output_rows; output_rows--) {
0654           if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
0655             data_col += output_w;
0656           } else {
0657             int input_col = -pad_w + kernel_col * dilation_w;
0658             for (int output_col = output_w; output_col; output_col--) {
0659               if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
0660                 //assert(input_row*width+input_col < height * width * channels);
0661                 //assert(data_col - data_col_0 < output_h*output_w*channels);
0662                //  std::cout << "COL2IM: input_row" << "  " << input_row << "  " << input_col
0663                //       << " <---- " << data_col - data_col_0 << " values:  "
0664                //       << data_im[input_row * width + input_col] << " <--- " << *data_col << std::endl;
0665                 data_im[input_row * width + input_col] += *data_col;
0666               }
0667               data_col++;
0668               input_col += stride_w;
0669             }
0670           }
0671           input_row += stride_h;
0672         }
0673       }
0674     }
0675   }
0676   //std::cout << "finishing col2imp" << std::endl;
0677 }
0678 
0679 }  // end namespace UTILITY
0680 
0681 namespace BLAS{
0682 extern "C" void sgemm_(const char * transa, const char * transb, const int * m, const int * n, const int * k,
0683                        const float * alpha, const float * A, const int * lda, const float * B, const int * ldb,
0684                        const float * beta, float * C, const int * ldc);
0685 }//BLAS
0686 
0687 
0688 struct GNN_Data {
0689       RTensor<float> node_data;      // the node feature data, tensor with shape (num_nodes, num_node_features)
0690       RTensor<float> edge_data;      // the edge feature data, tensor with shape (num_edges, num_edge_features)
0691       RTensor<float> global_data;    // the global features, tensor with shape (1, num_global_features)
0692       RTensor<int> edge_index;       // the edge index (receivers and senders for each edge), tensor with shape (2, num_edges)
0693                                      // edge_index[0,:] are the receivers and edge_index[1,:] are the senders
0694 
0695 
0696       // need to have default constructor since RTensor has not one
0697       GNN_Data(): node_data(RTensor<float>({})), edge_data(RTensor<float>({})), global_data(RTensor<float>({})), edge_index(RTensor<int>({})) {}
0698 
0699 };
0700 
0701 template<typename T>
0702 TMVA::Experimental::RTensor<T> Concatenate( TMVA::Experimental::RTensor<T> & t1,  TMVA::Experimental::RTensor<T> & t2, int axis = 0)
0703 {
0704    // concatenate tensor along axis. Shape must be the same except in the dimension of the concatenated axis
0705    if (t1.GetMemoryLayout() != t2.GetMemoryLayout())
0706       throw std::runtime_error("TMVA RTensor Concatenate - tensors have different memory layout");
0707    auto & shape1 = t1.GetShape();
0708    auto & shape2 = t2.GetShape();
0709    if (t1.GetSize()/shape1[axis] != t2.GetSize()/shape2[axis]) {
0710       std::cout << "axis " << axis << " sizes " << t1.GetSize() << " " << t2.GetSize() << "  ";
0711       std::cout << "shape 1 : " << ConvertShapeToString(t1.GetShape());
0712       std::cout << " shape 2 : " << ConvertShapeToString(t2.GetShape()) << std::endl;
0713       throw std::runtime_error("TMVA RTensor Concatenate - tensors have incompatible shapes");
0714    }
0715    std::vector<size_t> outShape = shape1;
0716    outShape[axis] = shape1[axis] + shape2[axis];
0717    TMVA::Experimental::RTensor<T> tout(outShape, t1.GetMemoryLayout());
0718    if (t1.GetMemoryLayout() == TMVA::Experimental::MemoryLayout::ColumnMajor) {
0719       throw std::runtime_error("TMVA RTensor Concatenate is not yet supported for column major tensors");
0720    }
0721 
0722    auto & stride1 = t1.GetStrides();
0723    auto & stride2 = t2.GetStrides();
0724    auto & outStride = tout.GetStrides();
0725 
0726    size_t s1 = (axis > 0) ? stride1[axis-1] : t1.GetSize();  // block size to copy from first tensor
0727    size_t s2 = (axis > 0) ? stride2[axis-1] : t2.GetSize();  // block size to copy from second tensor
0728    size_t sout = (axis > 0) ? outStride[axis-1] : tout.GetSize();
0729    size_t nb = t1.GetSize()/s1;
0730    for (size_t i = 0; i < nb; i++) {
0731       std::copy(t1.GetData() + i*s1, t1.GetData() + (i+1)*s1, tout.GetData() + i * sout );
0732       std::copy(t2.GetData() + i*s2, t2.GetData() + (i+1)*s2, tout.GetData() + i * sout + s1 );
0733    }
0734 
0735    return tout;
0736 }
0737 
0738 
0739 inline GNN_Data Concatenate(GNN_Data & data1, GNN_Data & data2, int axis = 0) {
0740    GNN_Data out;
0741    out.node_data = Concatenate(data1.node_data,data2.node_data, axis);
0742    out.edge_data = Concatenate(data1.edge_data,data2.edge_data, axis);
0743    out.global_data = Concatenate<float>(data1.global_data,data2.global_data, axis-1);
0744    // assume sender/receivers of data1 and data2 are the same
0745    out.edge_index = data1.edge_index.Copy();
0746    return out;
0747 }
0748 
0749 inline GNN_Data Copy(const GNN_Data & data) {
0750    GNN_Data out;
0751    out.node_data = RTensor<float>(data.node_data.GetShape());
0752    out.edge_data = RTensor<float>(data.edge_data.GetShape());
0753    out.global_data = RTensor<float>(data.global_data.GetShape());
0754    out.edge_index = RTensor<int>(data.edge_index.GetShape());
0755    std::copy(data.node_data.GetData(), data.node_data.GetData()+ data.node_data.GetSize(), out.node_data.GetData());
0756    std::copy(data.edge_data.GetData(), data.edge_data.GetData()+ data.edge_data.GetSize(), out.edge_data.GetData());
0757    std::copy(data.global_data.GetData(), data.global_data.GetData()+ data.global_data.GetSize(), out.global_data.GetData());
0758    std::copy(data.edge_index.GetData(), data.edge_index.GetData()+ data.edge_index.GetSize(), out.edge_index.GetData());
0759    return out;
0760 }
0761 
0762 inline void Gemm_Call(float *output, bool transa, bool transb, int m, int n, int k, float alpha, const float *A,
0763                       const float *B, float beta, const float *C)
0764 {
0765    char ct = 't';
0766    char cn = 'n';
0767    const int *lda = transa ? &k : &m;
0768    const int *ldb = transb ? &n : &k;
0769    const int *ldc = &m;
0770    if (C != nullptr) {
0771       std::copy(C, C + m * n, output);
0772    }
0773    TMVA::Experimental::SOFIE::BLAS::sgemm_(transa ? &ct : &cn, transb ? &ct : &cn, &m, &n, &k, &alpha, A, lda, B, ldb,
0774                                            &beta, output, ldc);
0775 }
0776 
0777 inline void Fill(float *output, float value, int size)
0778 {
0779    std::fill(output, output + size, value);
0780 }
0781 
0782 inline void Copy(float *output, float const *input, int size)
0783 {
0784    std::copy(input, input + size, output);
0785 }
0786 
0787 inline void Relu(float *output, float const *input, int size)
0788 {
0789    for (int i = 0; i < size; i++) {
0790       output[i] = (input[i] > 0.0f) ? input[i] : 0.0f;
0791    }
0792 }
0793 // function to read float from the file dealing with inf and nan values
0794 inline float ParseFloatToken (const std::string & s)  {
0795    if (s == "inf")  return  std::numeric_limits<float>::infinity();
0796    if (s == "-inf") return -std::numeric_limits<float>::infinity();
0797    if (s == "nan")  return  std::numeric_limits<float>::quiet_NaN();
0798    return std::stof(s);
0799 }
0800 
0801 template <class T>
0802 void ReadTensorFromStream(std::istream &is, T &target, std::string const &expectedName, std::size_t expectedLength)
0803 {
0804    std::string name;
0805    std::size_t length;
0806    is >> name >> length;
0807    if (name != expectedName) {
0808       std::string err_msg =
0809          "TMVA-SOFIE failed to read the correct tensor name; expected name is " + expectedName + " , read " + name;
0810       throw std::runtime_error(err_msg);
0811    }
0812    if (length != expectedLength) {
0813       std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is " +
0814                             std::to_string(expectedLength) + " , read " + std::to_string(length);
0815       throw std::runtime_error(err_msg);
0816    }
0817    std::string token;
0818    for (size_t i = 0; i < length; ++i) {
0819       is >> token;
0820       target[i] = ParseFloatToken(token);
0821    }
0822    if (is.fail()) {
0823       throw std::runtime_error("TMVA-SOFIE failed to read the values for tensor " + expectedName);
0824    }
0825 }
0826 
0827 //Utility functions to generate code
0828 void EmitNestedLoops(std::stringstream &out, size_t loopRank, const std::vector<Dim> shape);
0829 void CloseNestedLoops(std::stringstream &out, size_t loopRank);
0830 
0831 
0832 // code for the memory greeding allocations
0833 struct TensorLifeInfo {
0834    int begin;   // start time (op index) lifetime
0835    int end;     //  end time lifetime
0836    size_t size; // size of tensors in bytes
0837 };
0838 
0839 struct MemoryResult {
0840   std::size_t total_bytes = 0;  // total memory needed
0841   std::vector<size_t> offsets; // resulted offsets for each tensor
0842 };
0843 
0844 /// Greedy best-fit planner with coalescing free list.
0845 MemoryResult OrganizeMemory(const std::vector<TensorLifeInfo> & tensorsInfo );
0846 
0847 // Simple Dimension classes ans helpers to add constexpr meta info on input
0848 // tensors to the emitted code.
0849 struct SingleDim {
0850    enum class Kind {
0851       Static,
0852       Symbolic
0853    };
0854 
0855    Kind kind;
0856    std::size_t dim;
0857    std::string_view name;
0858 
0859    constexpr SingleDim(std::size_t v) : kind(Kind::Static), dim(v), name() {}
0860    constexpr SingleDim(const char *v) : kind(Kind::Symbolic), dim(0), name(v) {}
0861 };
0862 
0863 struct TensorDims {
0864    const SingleDim *data;
0865    std::size_t size;
0866 
0867    constexpr std::size_t total_size() const
0868    {
0869       std::size_t result = 1;
0870       for (std::size_t i = 0; i < size; ++i) {
0871          result *= data[i].dim;
0872       }
0873       return result;
0874    }
0875 };
0876 
0877 template<class Arr>
0878 constexpr TensorDims makeDims(Arr const &arr)
0879 {
0880    return TensorDims{arr.data(), arr.size()};
0881 }
0882 
0883 } // namespace SOFIE
0884 } // namespace Experimental
0885 } // namespace TMVA
0886 
0887 #endif //TMVA_SOFIE_COMMON