operation.h

Go to the documentation of this file.

#pragma once
#include <string>
#include <map>
#include "tensor/tensor.h"

namespace magmadnn {
namespace op {

template <typename T>
class Operation {
public: 
    Operation() : has_been_computed(false) {}
    Operation(std::vector<Operation<T> *> inputs, bool needs_grad=true) : inputs(inputs), needs_grad(needs_grad) {
        for (typename std::vector<Operation<T> *>::iterator vit = inputs.begin(); vit != inputs.end(); vit++) {
            if (needs_grad) {   /* TODO : verify this is necessary */
                (*vit)->add_consumer(this);
            }
            this->_grad_cache.insert( std::make_pair((uintptr_t) (*vit), (Tensor<T> *) NULL) );
        }
    }
    virtual ~Operation() {
        for (unsigned int i = 0; i < inputs.size(); i++)
            delete inputs[i];
        
        /*  TODO : figure out why this peice of code caused SEGFAULTS 
        if (this->output_tensor != NULL) {
            delete this->output_tensor;
        }
        */
    }

    virtual std::vector<unsigned int> get_output_shape() const { return this->output_shape; }

    virtual unsigned int get_output_shape(unsigned int idx) const {
        assert( idx < this->output_shape.size() );
        return this->output_shape[idx];
    }

    virtual unsigned int get_output_size() const {
        unsigned int size = 1;
        for (unsigned int i = 0; i < this->output_shape.size(); i++) size *= this->output_shape[i];
        return size;
    }

    virtual memory_t get_memory_type() const { return this->mem_type; }

    virtual Tensor<T>* eval(bool recompute=true) {
        if (!recompute && this->has_been_computed && this->output_tensor != NULL) {
            return this->output_tensor;
        } else {
            this->has_been_computed = true;
            return _eval(recompute);
        }
    }

    virtual void reset() { this->has_been_computed = false; this->has_grad_been_computed = false; }

    virtual Tensor<T>* grad(Operation<T> *consumer, Operation<T> *var, Tensor<T> *grad, bool recompute=true) {
        if (!recompute) {
            Tensor<T> *ret;
            ret = this->_grad_cache[(uintptr_t)var];
            
            if (ret != NULL) {
                return ret;
            } else {
                return _grad(consumer, var, grad);
            }
        } else {
            return _grad(consumer, var, grad);
        }
    }

    virtual void add_consumer(Operation<T> *consumer) { this->consumers.push_back(consumer); }

    virtual std::vector<Operation<T> *> get_consumers() { return this->consumers; }

    virtual std::vector<Operation<T> *> get_inputs() { return this->inputs; }


    virtual Tensor<T> *get_output_tensor() { return this->output_tensor; }

    virtual Tensor<T> *get_grad_tensor(Operation<T> *wrt) { return this->_grad_cache.find((uintptr_t)wrt)->second; }

    virtual std::string to_string() = 0;

    virtual std::string get_name() { return this->name; }
    
protected:
    virtual Tensor<T> *_eval(bool recompute=true) = 0;

    virtual Tensor<T> *_grad(Operation<T> *consumer, Operation<T> *var, Tensor<T> *grad) = 0;

    std::vector<Operation<T>*> inputs;
    std::vector<Operation<T>*> consumers;
    std::vector<unsigned int> output_shape;
    memory_t mem_type;
    std::map<uintptr_t, Tensor<T> *> _grad_cache;   /* this will cache the tensors for the gradient computation */
    std::string name = "DefaultOpName";

    Tensor<T> *output_tensor; /* the return tensor */

    bool needs_grad;
    bool has_been_computed;
    bool has_grad_been_computed;
};

} // namespace op
} // namespace magmadnn