tensor.hpp

#pragma once 
#include <dlprim/context.hpp>
#include <dlprim/shape.hpp>
#include <memory>
namespace dlprim {

    class Tensor;    
    class TensorSpecs {
    public:
        TensorSpecs(Shape const &s=Shape(),DataType d=float_data,bool trainable = true) :
            shape_(s),
            dtype_(d)
        {
            is_trainable_ = trainable && is_floating_point_data_type(d);
        }
        
        bool operator==(TensorSpecs const &other) const
        {
            return shape_ == other.shape_ && dtype_ == other.dtype_ && is_trainable_ == other.is_trainable_;
        }
        
        bool operator!=(TensorSpecs const &other) const
        {
            return !(*this == other);
        }

        Shape const &shape() const
        {
            return shape_;
        }

        void shape(Shape const &s)
        {
            shape_=s;
        }

        bool is_trainable() const
        {
            return is_trainable_;
        }

        void freeze()
        {
            is_trainable_ = false;
        }

        void is_trainable(bool v)
        {
            is_trainable_ = v;
        }

        size_t memory_size() const
        {
            return shape_.total_size() * size_of_data_type(dtype_);
        }


        DataType dtype() const
        {
            return dtype_;
        }
    private:
        friend class Tensor;
        Shape shape_;
        DataType dtype_;
        bool is_trainable_;
    };


    class Tensor {
    public:
       
        Tensor(Context &ctx,Shape const &s,DataType d=float_data,bool is_trainable=true);
        
        Tensor(cl::Buffer const &buffer,cl_ulong offset,Shape const &s,DataType d=float_data,bool is_trainable=true);

        Tensor();
        Tensor(Tensor const &) = default;
        Tensor &operator=(Tensor const &) = default;
        Tensor(Tensor &&) = default;
        Tensor &operator=(Tensor &&) = default;
        ~Tensor() {}
        
        TensorSpecs const &specs() const
        {
            return *specs_;
        }

        Shape const &shape() const
        {
            return specs_->shape();
        }

        bool is_trainable() const
        {
            return specs_->is_trainable();
        }

        size_t memory_size() const
        {
            return shape().total_size() * size_of_data_type(dtype());
        }


        DataType dtype() const
        {
            return specs_->dtype();
        }
        

        void reshape(Shape const &ns);
        
        cl::Buffer &device_buffer() 
        { 
            return buffer_;
        }
        cl_ulong device_offset() 
        {
            return offset_; 
        }
        void *host_data();


        Tensor workspace_as_type(DataType d=float_data) const
        {
            size_t size = memory_size() / size_of_data_type(d);
            return sub_tensor(0,Shape(size),d);
        }
        
        Tensor sub_tensor_target_offset(size_t offset,Shape const &s,DataType d=float_data,bool trainable = true) const
        {
            size_t bytes = offset * size_of_data_type(d);
            int this_sizeof = size_of_data_type(dtype());
            DLPRIM_CHECK(bytes % this_sizeof == 0);
            return sub_tensor(bytes / this_sizeof,s,d,trainable);
        }
        Tensor sub_tensor(size_t offset,Shape const &s,DataType d=float_data,bool trainable = true) const;

        Tensor alias() const
        {
            return sub_tensor(0,shape(),dtype(),is_trainable());
        }
        Tensor alias(Shape const &new_shape) const
        {
            Tensor t = alias();
            t.reshape(new_shape);
            return t;
        }

        template<typename T>
        T *data()
        {
            DLPRIM_CHECK(TypeTraits<T>::data_type == dtype());
            return static_cast<T*>(host_data());
        }

        void to_device(ExecutionContext const &c,void *host_memory,bool sync=true);
        void to_host(ExecutionContext const &c,void *host_memory,bool sync=true);
        void to_device(ExecutionContext const &c,bool sync=true);
        void to_host(ExecutionContext const &c,bool sync=true);

        void set_arg(cl::Kernel &k,int &pos)
        {
            k.setArg(pos++,device_buffer());
            k.setArg(pos++,device_offset());
        }

    private:
                struct HostMem;
        std::shared_ptr<TensorSpecs> specs_;
        std::shared_ptr<HostMem> host_;
        bool cpu_tensor_;
        int offset_;
        cl::Buffer buffer_;
        size_t capacity_;
        size_t full_capacity_;
    };

    struct TensorAndGradient {
        bool requires_gradient=true; 
        float accumulate_gradient=0.0; 
        Tensor data; 
        Tensor diff; 
    };

    inline std::ostream &operator<<(std::ostream &out,TensorSpecs const &ts)
    {
        out << '[' << ts.shape() << ",dtype=" << data_type_to_string(ts.dtype()) << ']';
        return out;
    }

    inline std::ostream &operator<<(std::ostream &out,Tensor const &ts)
    {
        out << ts.specs();
        return out;
    }

}