<commit_before>#include "btree/iteration.hpp"
#include "btree/btree_data_provider.hpp"
#include "perfmon.hpp"

perfmon_counter_t
    leaf_iterators("leaf_iterators"),
    slice_leaves_iterators("slice_leaves_iterators");

leaf_iterator_t::leaf_iterator_t(const leaf_node_t *leaf, int index, buf_lock_t *lock, const boost::shared_ptr<transaction_t>& transaction) :
    leaf(leaf), index(index), lock(lock), transaction(transaction) {

    rassert(leaf != NULL);
    rassert(lock != NULL);
    leaf_iterators++;
}

boost::optional<key_with_data_provider_t> leaf_iterator_t::next() {
    rassert(index >= 0);
    const btree_leaf_pair *pair;
    do {
        if (index >= leaf->npairs) {
            done();
            return boost::none;
        }
        pair = leaf::get_pair_by_index(leaf, index++);
    } while (reinterpret_cast<const btree_value_t *>(pair->value())->expired());

    return boost::make_optional(pair_to_key_with_data_provider(pair));
}

void leaf_iterator_t::prefetch() {
    /* this space intentionally left blank */
}

leaf_iterator_t::~leaf_iterator_t() {
    leaf_iterators--;
    done();
}

void leaf_iterator_t::done() {
    if (lock) {
        on_thread_t th(transaction->home_thread());
        delete lock;
        lock = NULL;
    }
}

key_with_data_provider_t leaf_iterator_t::pair_to_key_with_data_provider(const btree_leaf_pair* pair) {
    on_thread_t th(transaction->home_thread());
    value_data_provider_t *data_provider = value_data_provider_t::create(reinterpret_cast<const btree_value_t *>(pair->value()), transaction.get());
    return key_with_data_provider_t(key_to_str(&pair->key), reinterpret_cast<const btree_value_t *>(pair->value())->mcflags(),
        boost::shared_ptr<data_provider_t>(data_provider));
}

slice_leaves_iterator_t::slice_leaves_iterator_t(const boost::shared_ptr<transaction_t>& transaction, btree_slice_t *slice,
    rget_bound_mode_t left_mode, const btree_key_t *left_key, rget_bound_mode_t right_mode, const btree_key_t *right_key) :
        transaction(transaction), slice(slice),
        left_mode(left_mode), left_key(left_key), right_mode(right_mode), right_key(right_key),
        traversal_state(), started(false), nevermore(false) {
            slice_leaves_iterators++;
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::next() {
    if (nevermore)
        return boost::none;

    boost::optional<leaf_iterator_t*> leaf = !started ? get_first_leaf() : get_next_leaf();
    if (!leaf)
        done();
    return leaf;
}

void slice_leaves_iterator_t::prefetch() {
    // Not implemented yet
}

slice_leaves_iterator_t::~slice_leaves_iterator_t() {
    slice_leaves_iterators--;
    done();
}

void slice_leaves_iterator_t::done() {
    while (!traversal_state.empty()) {
        delete traversal_state.back().lock;
        traversal_state.pop_back();
    }
    nevermore = true;
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::get_first_leaf() {
    on_thread_t mover(slice->home_thread()); // Move to the slice's thread.

    started = true;
    // TODO: Why is this a buf_lock_t pointer?  That's not how buf_lock_t should be used.
    buf_lock_t *buf_lock = new buf_lock_t(transaction.get(), block_id_t(SUPERBLOCK_ID), rwi_read);
    block_id_t root_id = reinterpret_cast<const btree_superblock_t *>(buf_lock->buf()->get_data_read())->root_block;
    rassert(root_id != SUPERBLOCK_ID);

    if (root_id == NULL_BLOCK_ID) {
        delete buf_lock;
        return boost::none;
    }

    if (left_mode == rget_bound_none) {
        boost::optional<leaf_iterator_t*> leftmost_leaf = get_leftmost_leaf(root_id);
        delete buf_lock;
        return leftmost_leaf;
    }

    {
        buf_lock_t tmp(transaction.get(), root_id, rwi_read);
        buf_lock->swap(tmp);
    }

    // read root node
    const node_t *node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());
    rassert(node != NULL);

    while (node::is_internal(node)) {
        const internal_node_t *i_node = reinterpret_cast<const internal_node_t *>(node);

        // push the i_node onto the traversal_state stack
        int index = internal_node::impl::get_offset_index(i_node, left_key);
        rassert(index >= 0);
        if (index >= i_node->npairs) {
            // this subtree has all the keys smaller than 'left_key', move to the leaf in the next subtree
            delete buf_lock;
            return get_next_leaf();
        }
        traversal_state.push_back(internal_node_state(i_node, index, buf_lock));

        block_id_t child_id = get_child_id(i_node, index);

        buf_lock = new buf_lock_t(transaction.get(), child_id, rwi_read);
        node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());
    }

    rassert(node != NULL);
    rassert(node::is_leaf(node));
    rassert(buf_lock != NULL);

    const leaf_node_t *l_node = reinterpret_cast<const leaf_node_t *>(node);
    int index = leaf::impl::get_offset_index(l_node, left_key);

    if (index < l_node->npairs) {
        return boost::make_optional(new leaf_iterator_t(l_node, index, buf_lock, transaction));
    } else {
        // there's nothing more in this leaf, we might as well move to the next leaf
        delete buf_lock;
        return get_next_leaf();
    }
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::get_next_leaf() {
    while (traversal_state.size() > 0) {
        internal_node_state state = traversal_state.back();
        traversal_state.pop_back();

        rassert(state.index >= 0 && state.index < state.node->npairs);
        ++state.index;
        if (state.index < state.node->npairs) {
            traversal_state.push_back(state);

            // get next child and find its leftmost leaf
            block_id_t child_id = get_child_id(state.node, state.index);
            return get_leftmost_leaf(child_id);
        } else {
            delete state.lock;  // this also releases the memory used by state.node
        }
    }
    return boost::none;
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::get_leftmost_leaf(block_id_t node_id) {
    on_thread_t mover(slice->home_thread()); // Move to the slice's thread.

    // TODO: Why is there a buf_lock_t pointer?  This is not how
    // buf_lock_t works.  Just use a buf_t pointer then.
    buf_lock_t *buf_lock = new buf_lock_t(transaction.get(), node_id, rwi_read);
    const node_t *node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());

    while (node::is_internal(node)) {
        const internal_node_t *i_node = reinterpret_cast<const internal_node_t *>(node);
        rassert(i_node->npairs > 0);

        // push the i_node onto the traversal_state stack
        const int leftmost_child_index = 0;
        traversal_state.push_back(internal_node_state(i_node, leftmost_child_index, buf_lock));

        block_id_t child_id = get_child_id(i_node, leftmost_child_index);

        buf_lock = new buf_lock_t(transaction.get(), child_id, rwi_read);
        node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());
    }
    rassert(node != NULL);
    rassert(node::is_leaf(node));
    rassert(buf_lock != NULL);

    const leaf_node_t *l_node = reinterpret_cast<const leaf_node_t *>(node);
    return boost::make_optional(new leaf_iterator_t(l_node, 0, buf_lock, transaction));
}

block_id_t slice_leaves_iterator_t::get_child_id(const internal_node_t *i_node, int index) const {
    block_id_t child_id = internal_node::get_pair_by_index(i_node, index)->lnode;
    rassert(child_id != NULL_BLOCK_ID);
    rassert(child_id != SUPERBLOCK_ID);

    return child_id;
}

slice_keys_iterator_t::slice_keys_iterator_t(const boost::shared_ptr<transaction_t>& transaction_, btree_slice_t *slice_,
        rget_bound_mode_t left_mode, const store_key_t &left_key, rget_bound_mode_t right_mode, const store_key_t &right_key) :
    transaction(transaction_), slice(slice_),
    left_mode(left_mode), left_key(left_key), right_mode(right_mode), right_key(right_key),
    left_str(key_to_str(left_key)), right_str(key_to_str(right_key)),
    no_more_data(false), active_leaf(NULL), leaves_iterator(NULL) { }

slice_keys_iterator_t::~slice_keys_iterator_t() {
    done();
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::next() {
    if (no_more_data)
        return boost::none;

    boost::optional<key_with_data_provider_t> result = active_leaf == NULL ? get_first_value() : get_next_value();
    if (!result) {
        done();
    }
    return result;
}

void slice_keys_iterator_t::prefetch() {
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::get_first_value() {
    leaves_iterator = new slice_leaves_iterator_t(transaction, slice, left_mode, left_key.key(), right_mode, right_key.key());

    // get the first leaf with our left key (or something greater)
    boost::optional<leaf_iterator_t*> first = leaves_iterator->next();
    if (!first)
        return boost::none;

    active_leaf = first.get();
    rassert(active_leaf != NULL);

    // get a value from the leaf
    boost::optional<key_with_data_provider_t> first_pair = active_leaf->next();
    if (!first_pair)
        return first_pair;  // returns empty result

    key_with_data_provider_t pair = first_pair.get();

    // skip the left key if left_mode == rget_bound_mode_open
    int compare_result = pair.key.compare(left_str);
    rassert(compare_result >= 0);
    if (left_mode == rget_bound_open && compare_result == 0) {
        return get_next_value();
    } else {
        return validate_return_value(pair);
    }
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::get_next_value() {
    rassert(leaves_iterator != NULL);
    rassert(active_leaf != NULL);

    // get a value from the leaf
    boost::optional<key_with_data_provider_t> current_pair = active_leaf->next();
    if (!current_pair) {
        delete active_leaf;
        active_leaf = NULL;

        boost::optional<leaf_iterator_t*> leaf = leaves_iterator->next();
        if (!leaf)
            return boost::none;

        active_leaf = leaf.get();
        current_pair = active_leaf->next();
        if (!current_pair)
            return boost::none;
    }

    return validate_return_value(current_pair.get());
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::validate_return_value(key_with_data_provider_t &pair) const {
    if (right_mode == rget_bound_none)
        return boost::make_optional(pair);

    int compare_result = pair.key.compare(right_str);
    if (compare_result < 0 || (right_mode == rget_bound_closed && compare_result == 0)) {
        return boost::make_optional(pair);
    } else {
        return boost::none;
    }
}

void slice_keys_iterator_t::done() {
    if (active_leaf) {
        delete active_leaf;
        active_leaf = NULL;
    }
    if (leaves_iterator) {
        delete leaves_iterator;
        leaves_iterator = NULL;
    }
    no_more_data = true;
}

<commit_msg>Made constructor arguments named appropriately.<commit_after>#include "btree/iteration.hpp"
#include "btree/btree_data_provider.hpp"
#include "perfmon.hpp"

perfmon_counter_t
    leaf_iterators("leaf_iterators"),
    slice_leaves_iterators("slice_leaves_iterators");

leaf_iterator_t::leaf_iterator_t(const leaf_node_t *_leaf, int _index, buf_lock_t *_lock, const boost::shared_ptr<transaction_t>& _transaction) :
    leaf(_leaf), index(_index), lock(_lock), transaction(_transaction) {

    rassert(leaf != NULL);
    rassert(lock != NULL);
    leaf_iterators++;
}

boost::optional<key_with_data_provider_t> leaf_iterator_t::next() {
    rassert(index >= 0);
    const btree_leaf_pair *pair;
    do {
        if (index >= leaf->npairs) {
            done();
            return boost::none;
        }
        pair = leaf::get_pair_by_index(leaf, index++);
    } while (reinterpret_cast<const btree_value_t *>(pair->value())->expired());

    return boost::make_optional(pair_to_key_with_data_provider(pair));
}

void leaf_iterator_t::prefetch() {
    /* this space intentionally left blank */
}

leaf_iterator_t::~leaf_iterator_t() {
    leaf_iterators--;
    done();
}

void leaf_iterator_t::done() {
    if (lock) {
        on_thread_t th(transaction->home_thread());
        delete lock;
        lock = NULL;
    }
}

key_with_data_provider_t leaf_iterator_t::pair_to_key_with_data_provider(const btree_leaf_pair* pair) {
    on_thread_t th(transaction->home_thread());
    value_data_provider_t *data_provider = value_data_provider_t::create(reinterpret_cast<const btree_value_t *>(pair->value()), transaction.get());
    return key_with_data_provider_t(key_to_str(&pair->key), reinterpret_cast<const btree_value_t *>(pair->value())->mcflags(),
        boost::shared_ptr<data_provider_t>(data_provider));
}

slice_leaves_iterator_t::slice_leaves_iterator_t(const boost::shared_ptr<transaction_t>& transaction, btree_slice_t *slice,
    rget_bound_mode_t left_mode, const btree_key_t *left_key, rget_bound_mode_t right_mode, const btree_key_t *right_key) :
        transaction(transaction), slice(slice),
        left_mode(left_mode), left_key(left_key), right_mode(right_mode), right_key(right_key),
        traversal_state(), started(false), nevermore(false) {
            slice_leaves_iterators++;
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::next() {
    if (nevermore)
        return boost::none;

    boost::optional<leaf_iterator_t*> leaf = !started ? get_first_leaf() : get_next_leaf();
    if (!leaf)
        done();
    return leaf;
}

void slice_leaves_iterator_t::prefetch() {
    // Not implemented yet
}

slice_leaves_iterator_t::~slice_leaves_iterator_t() {
    slice_leaves_iterators--;
    done();
}

void slice_leaves_iterator_t::done() {
    while (!traversal_state.empty()) {
        delete traversal_state.back().lock;
        traversal_state.pop_back();
    }
    nevermore = true;
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::get_first_leaf() {
    on_thread_t mover(slice->home_thread()); // Move to the slice's thread.

    started = true;
    // TODO: Why is this a buf_lock_t pointer?  That's not how buf_lock_t should be used.
    buf_lock_t *buf_lock = new buf_lock_t(transaction.get(), block_id_t(SUPERBLOCK_ID), rwi_read);
    block_id_t root_id = reinterpret_cast<const btree_superblock_t *>(buf_lock->buf()->get_data_read())->root_block;
    rassert(root_id != SUPERBLOCK_ID);

    if (root_id == NULL_BLOCK_ID) {
        delete buf_lock;
        return boost::none;
    }

    if (left_mode == rget_bound_none) {
        boost::optional<leaf_iterator_t*> leftmost_leaf = get_leftmost_leaf(root_id);
        delete buf_lock;
        return leftmost_leaf;
    }

    {
        buf_lock_t tmp(transaction.get(), root_id, rwi_read);
        buf_lock->swap(tmp);
    }

    // read root node
    const node_t *node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());
    rassert(node != NULL);

    while (node::is_internal(node)) {
        const internal_node_t *i_node = reinterpret_cast<const internal_node_t *>(node);

        // push the i_node onto the traversal_state stack
        int index = internal_node::impl::get_offset_index(i_node, left_key);
        rassert(index >= 0);
        if (index >= i_node->npairs) {
            // this subtree has all the keys smaller than 'left_key', move to the leaf in the next subtree
            delete buf_lock;
            return get_next_leaf();
        }
        traversal_state.push_back(internal_node_state(i_node, index, buf_lock));

        block_id_t child_id = get_child_id(i_node, index);

        buf_lock = new buf_lock_t(transaction.get(), child_id, rwi_read);
        node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());
    }

    rassert(node != NULL);
    rassert(node::is_leaf(node));
    rassert(buf_lock != NULL);

    const leaf_node_t *l_node = reinterpret_cast<const leaf_node_t *>(node);
    int index = leaf::impl::get_offset_index(l_node, left_key);

    if (index < l_node->npairs) {
        return boost::make_optional(new leaf_iterator_t(l_node, index, buf_lock, transaction));
    } else {
        // there's nothing more in this leaf, we might as well move to the next leaf
        delete buf_lock;
        return get_next_leaf();
    }
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::get_next_leaf() {
    while (traversal_state.size() > 0) {
        internal_node_state state = traversal_state.back();
        traversal_state.pop_back();

        rassert(state.index >= 0 && state.index < state.node->npairs);
        ++state.index;
        if (state.index < state.node->npairs) {
            traversal_state.push_back(state);

            // get next child and find its leftmost leaf
            block_id_t child_id = get_child_id(state.node, state.index);
            return get_leftmost_leaf(child_id);
        } else {
            delete state.lock;  // this also releases the memory used by state.node
        }
    }
    return boost::none;
}

boost::optional<leaf_iterator_t*> slice_leaves_iterator_t::get_leftmost_leaf(block_id_t node_id) {
    on_thread_t mover(slice->home_thread()); // Move to the slice's thread.

    // TODO: Why is there a buf_lock_t pointer?  This is not how
    // buf_lock_t works.  Just use a buf_t pointer then.
    buf_lock_t *buf_lock = new buf_lock_t(transaction.get(), node_id, rwi_read);
    const node_t *node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());

    while (node::is_internal(node)) {
        const internal_node_t *i_node = reinterpret_cast<const internal_node_t *>(node);
        rassert(i_node->npairs > 0);

        // push the i_node onto the traversal_state stack
        const int leftmost_child_index = 0;
        traversal_state.push_back(internal_node_state(i_node, leftmost_child_index, buf_lock));

        block_id_t child_id = get_child_id(i_node, leftmost_child_index);

        buf_lock = new buf_lock_t(transaction.get(), child_id, rwi_read);
        node = reinterpret_cast<const node_t *>(buf_lock->buf()->get_data_read());
    }
    rassert(node != NULL);
    rassert(node::is_leaf(node));
    rassert(buf_lock != NULL);

    const leaf_node_t *l_node = reinterpret_cast<const leaf_node_t *>(node);
    return boost::make_optional(new leaf_iterator_t(l_node, 0, buf_lock, transaction));
}

block_id_t slice_leaves_iterator_t::get_child_id(const internal_node_t *i_node, int index) const {
    block_id_t child_id = internal_node::get_pair_by_index(i_node, index)->lnode;
    rassert(child_id != NULL_BLOCK_ID);
    rassert(child_id != SUPERBLOCK_ID);

    return child_id;
}

slice_keys_iterator_t::slice_keys_iterator_t(const boost::shared_ptr<transaction_t>& transaction_, btree_slice_t *slice_,
        rget_bound_mode_t left_mode, const store_key_t &left_key, rget_bound_mode_t right_mode, const store_key_t &right_key) :
    transaction(transaction_), slice(slice_),
    left_mode(left_mode), left_key(left_key), right_mode(right_mode), right_key(right_key),
    left_str(key_to_str(left_key)), right_str(key_to_str(right_key)),
    no_more_data(false), active_leaf(NULL), leaves_iterator(NULL) { }

slice_keys_iterator_t::~slice_keys_iterator_t() {
    done();
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::next() {
    if (no_more_data)
        return boost::none;

    boost::optional<key_with_data_provider_t> result = active_leaf == NULL ? get_first_value() : get_next_value();
    if (!result) {
        done();
    }
    return result;
}

void slice_keys_iterator_t::prefetch() {
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::get_first_value() {
    leaves_iterator = new slice_leaves_iterator_t(transaction, slice, left_mode, left_key.key(), right_mode, right_key.key());

    // get the first leaf with our left key (or something greater)
    boost::optional<leaf_iterator_t*> first = leaves_iterator->next();
    if (!first)
        return boost::none;

    active_leaf = first.get();
    rassert(active_leaf != NULL);

    // get a value from the leaf
    boost::optional<key_with_data_provider_t> first_pair = active_leaf->next();
    if (!first_pair)
        return first_pair;  // returns empty result

    key_with_data_provider_t pair = first_pair.get();

    // skip the left key if left_mode == rget_bound_mode_open
    int compare_result = pair.key.compare(left_str);
    rassert(compare_result >= 0);
    if (left_mode == rget_bound_open && compare_result == 0) {
        return get_next_value();
    } else {
        return validate_return_value(pair);
    }
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::get_next_value() {
    rassert(leaves_iterator != NULL);
    rassert(active_leaf != NULL);

    // get a value from the leaf
    boost::optional<key_with_data_provider_t> current_pair = active_leaf->next();
    if (!current_pair) {
        delete active_leaf;
        active_leaf = NULL;

        boost::optional<leaf_iterator_t*> leaf = leaves_iterator->next();
        if (!leaf)
            return boost::none;

        active_leaf = leaf.get();
        current_pair = active_leaf->next();
        if (!current_pair)
            return boost::none;
    }

    return validate_return_value(current_pair.get());
}

boost::optional<key_with_data_provider_t> slice_keys_iterator_t::validate_return_value(key_with_data_provider_t &pair) const {
    if (right_mode == rget_bound_none)
        return boost::make_optional(pair);

    int compare_result = pair.key.compare(right_str);
    if (compare_result < 0 || (right_mode == rget_bound_closed && compare_result == 0)) {
        return boost::make_optional(pair);
    } else {
        return boost::none;
    }
}

void slice_keys_iterator_t::done() {
    if (active_leaf) {
        delete active_leaf;
        active_leaf = NULL;
    }
    if (leaves_iterator) {
        delete leaves_iterator;
        leaves_iterator = NULL;
    }
    no_more_data = true;
}

<|endoftext|>


<commit_before>#include "app.hpp"

namespace rack {

Port::~Port() {
	gRackWidget->wireContainer->removeAllWires(this);
}

void Port::draw(NVGcontext *vg) {
	WireWidget *activeWire = gRackWidget->wireContainer->activeWire;
	if (activeWire) {
		// Dim the Port if the active wire cannot plug into this Port
		if (type == INPUT ? activeWire->inputPort : activeWire->outputPort)
			nvgGlobalAlpha(vg, 0.5);
	}
}

void Port::onMouseDownOpaque(int button) {
	if (button == 1) {
		gRackWidget->wireContainer->removeTopWire(this);

		// HACK
		// Update hovered*Port of active wire if applicable
		onDragEnter(NULL);
	}
}

void Port::onDragEnd() {
	gRackWidget->wireContainer->commitActiveWire();
}

void Port::onDragStart() {
	// Try to grab wire on top of stack
	WireWidget *wire = gRackWidget->wireContainer->getTopWire(this);

	if (wire) {
		// Disconnect existing wire
		if (type == INPUT)
			wire->inputPort = NULL;
		else
			wire->outputPort = NULL;
		wire->updateWire();
	}
	else {
		// Create a new wire
		wire = new WireWidget();
		if (type == INPUT)
			wire->inputPort = this;
		else
			wire->outputPort = this;
	}
	gRackWidget->wireContainer->setActiveWire(wire);
}

void Port::onDragDrop(Widget *origin) {
}

void Port::onDragEnter(Widget *origin) {
	// Reject ports if this is an input port and something is already plugged into it
	if (type == INPUT) {
		WireWidget *topWire = gRackWidget->wireContainer->getTopWire(this);
		if (topWire)
			return;
	}

	WireWidget *activeWire = gRackWidget->wireContainer->activeWire;
	if (activeWire) {
		if (type == INPUT)
			activeWire->hoveredInputPort = this;
		else
			activeWire->hoveredOutputPort = this;
	}
}

void Port::onDragLeave(Widget *origin) {
	WireWidget *activeWire = gRackWidget->wireContainer->activeWire;
	if (activeWire) {
		if (type == INPUT)
			activeWire->hoveredInputPort = NULL;
		else
			activeWire->hoveredOutputPort = NULL;
	}
}

} // namespace rack
<commit_msg>Use Ctrl/Cmd+Drag on an output to create a new cable<commit_after>#include "app.hpp"
#include "gui.hpp"

namespace rack {

Port::~Port() {
	gRackWidget->wireContainer->removeAllWires(this);
}

void Port::draw(NVGcontext *vg) {
	WireWidget *activeWire = gRackWidget->wireContainer->activeWire;
	if (activeWire) {
		// Dim the Port if the active wire cannot plug into this Port
		if (type == INPUT ? activeWire->inputPort : activeWire->outputPort)
			nvgGlobalAlpha(vg, 0.5);
	}
}

void Port::onMouseDownOpaque(int button) {
	if (button == 1) {
		gRackWidget->wireContainer->removeTopWire(this);

		// HACK
		// Update hovered*Port of active wire if applicable
		onDragEnter(NULL);
	}
}

void Port::onDragEnd() {
	gRackWidget->wireContainer->commitActiveWire();
}

void Port::onDragStart() {
	// Try to grab wire on top of stack
	WireWidget *wire = gRackWidget->wireContainer->getTopWire(this);

	if (wire && !guiIsModPressed()) {
		// Disconnect existing wire
		if (type == INPUT)
			wire->inputPort = NULL;
		else
			wire->outputPort = NULL;
		wire->updateWire();
	}
	else {
		// Create a new wire
		wire = new WireWidget();
		if (type == INPUT)
			wire->inputPort = this;
		else
			wire->outputPort = this;
	}
	gRackWidget->wireContainer->setActiveWire(wire);
}

void Port::onDragDrop(Widget *origin) {
}

void Port::onDragEnter(Widget *origin) {
	// Reject ports if this is an input port and something is already plugged into it
	if (type == INPUT) {
		WireWidget *topWire = gRackWidget->wireContainer->getTopWire(this);
		if (topWire)
			return;
	}

	WireWidget *activeWire = gRackWidget->wireContainer->activeWire;
	if (activeWire) {
		if (type == INPUT)
			activeWire->hoveredInputPort = this;
		else
			activeWire->hoveredOutputPort = this;
	}
}

void Port::onDragLeave(Widget *origin) {
	WireWidget *activeWire = gRackWidget->wireContainer->activeWire;
	if (activeWire) {
		if (type == INPUT)
			activeWire->hoveredInputPort = NULL;
		else
			activeWire->hoveredOutputPort = NULL;
	}
}

} // namespace rack
<|endoftext|>


<commit_before>#include <btBulletDynamicsCommon.h>
#include <BulletCollision/CollisionDispatch/btGhostObject.h>
#include "../Physics/GlmConversion.hpp"
#include "RigidBody.hpp"

#ifdef USINGMEMTRACK
#include <MemTrackInclude.hpp>
#endif

namespace Component {
    RigidBody::~RigidBody() {
        Destroy();
    }

    Json::Value RigidBody::Save() const {
        Json::Value component;
        component["mass"] = mass;
        component["friction"] = friction;
        component["rollingFriction"] = rollingFriction;
        component["spinningFriction"] = spinningFriction;
        component["cor"] = restitution;
        component["linearDamping"] = linearDamping;
        component["angularDamping"] = angularDamping;
        component["kinematic"] = kinematic;
        return component;
    }

    bool RigidBody::IsKinematic() const {
        return kinematic;
    }

    float RigidBody::GetFriction() const {
        return friction;
    }

    float RigidBody::GetRollingFriction() const {
        return rollingFriction;
    }

    float RigidBody::GetSpinningFriction() const {
        return spinningFriction;
    }

    float RigidBody::GetRestitution() const {
        return restitution;
    }

    float RigidBody::GetLinearDamping() const {
        return linearDamping;
    }

    float RigidBody::GetAngularDamping() const {
        return angularDamping;
    }

    btRigidBody* RigidBody::GetBulletRigidBody() {
        return rigidBody;
    }

    void RigidBody::NewBulletRigidBody(float mass) {
        Destroy();

        this->mass = mass;

        // Motion states inform us of movement caused by physics so that we can
        // deal with those changes as needed.
        btDefaultMotionState* motionState = new btDefaultMotionState(btTransform(btQuaternion(0, 0, 0, 1), btVector3(0, 0, 0)));

        // Bullet treats zero mass as infinite, resulting in immovable objects.
        btRigidBody::btRigidBodyConstructionInfo constructionInfo(mass, motionState, nullptr, btVector3(0, 0, 0));
        rigidBody = new btRigidBody(constructionInfo);
        rigidBody->setUserPointer(this);

        ghostObject = new btGhostObject();
    }

    void RigidBody::Destroy() {
        if (rigidBody) {
            delete rigidBody->getMotionState();
            delete rigidBody;
        }

        delete ghostObject;
    }

    glm::vec3 RigidBody::GetPosition() const {
        btTransform trans;
        if (IsKinematic())
            rigidBody->getMotionState()->getWorldTransform(trans);
        else
            trans = rigidBody->getWorldTransform();

        return Physics::btToGlm(trans.getOrigin());
    }

    void RigidBody::SetPosition(const glm::vec3& pos) {
        if (IsKinematic()) {
            btTransform trans;
            rigidBody->getMotionState()->getWorldTransform(trans);
            trans.setOrigin(Physics::glmToBt(pos));
            rigidBody->getMotionState()->setWorldTransform(trans);
        } else {
            btTransform trans = rigidBody->getWorldTransform();
            trans.setOrigin(Physics::glmToBt(pos));
            rigidBody->setWorldTransform(trans);
        }
    }

    glm::quat RigidBody::GetOrientation() const {
        btTransform trans;
        if (IsKinematic())
            rigidBody->getMotionState()->getWorldTransform(trans);
        else
            trans = rigidBody->getWorldTransform();

        return Physics::btToGlm(trans.getRotation());
    }

    void RigidBody::SetOrientation(const glm::quat& rotation) {
        if (IsKinematic()) {
            btTransform trans;
            rigidBody->getMotionState()->getWorldTransform(trans);
            trans.setRotation(Physics::glmToBt(rotation));
            rigidBody->getMotionState()->setWorldTransform(trans);
        } else {
            btTransform trans = rigidBody->getWorldTransform();
            trans.setRotation(Physics::glmToBt(rotation));
            rigidBody->setWorldTransform(trans);
        }
    }

    float RigidBody::GetMass() {
        return mass;
    }

    void RigidBody::SetMass(float mass) {
        // Bullet provides a method on the shape that we can use to calculate
        // inertia.
        btVector3 inertia;
        rigidBody->getCollisionShape()->calculateLocalInertia(mass, inertia);
        rigidBody->setMassProps(mass, inertia);
        this->mass = mass;
    }

    void RigidBody::SetFriction(float friction) {
        rigidBody->setFriction(friction);
        this->friction = friction;
    }

    void RigidBody::SetRollingFriction(float friction) {
        rigidBody->setRollingFriction(friction);
        this->rollingFriction = friction;
    }

    void RigidBody::SetSpinningFriction(float friction) {
        rigidBody->setSpinningFriction(friction);
        this->spinningFriction = friction;
    }

    void RigidBody::SetRestitution(float cor) {
        rigidBody->setRestitution(cor);
        this->restitution = cor;
    }

    void RigidBody::SetLinearDamping(float damping) {
        rigidBody->setDamping(damping, angularDamping);
        this->linearDamping = damping;
    }

    void RigidBody::SetAngularDamping(float damping) {
        rigidBody->setDamping(linearDamping, damping);
        this->angularDamping = damping;
    }

    void RigidBody::MakeKinematic() {
        rigidBody->setCollisionFlags(rigidBody->getCollisionFlags() | btCollisionObject::CF_KINEMATIC_OBJECT);
        kinematic = true;
    }

    void RigidBody::MakeDynamic() {
        rigidBody->setCollisionFlags(rigidBody->getCollisionFlags() & ~btCollisionObject::CF_KINEMATIC_OBJECT);
        kinematic = false;
    }

    void RigidBody::SetGhost(bool ghost) {
        this->ghost = ghost;
    }

    bool RigidBody::GetForceTransformSync() const {
        return forceTransformSync;
    }

    void RigidBody::SetForceTransformSync(bool sync) {
        forceTransformSync = sync;
    }

    bool RigidBody::GetHaltMovement() const {
        return haltMovement;
    }

    void RigidBody::SetHaltMovement(bool halt) {
        haltMovement = halt;
    }
}
<commit_msg>Make cached ghost state false when making a rigid body kinematic or dynamic.<commit_after>#include <btBulletDynamicsCommon.h>
#include <BulletCollision/CollisionDispatch/btGhostObject.h>
#include "../Physics/GlmConversion.hpp"
#include "RigidBody.hpp"

#ifdef USINGMEMTRACK
#include <MemTrackInclude.hpp>
#endif

namespace Component {
    RigidBody::~RigidBody() {
        Destroy();
    }

    Json::Value RigidBody::Save() const {
        Json::Value component;
        component["mass"] = mass;
        component["friction"] = friction;
        component["rollingFriction"] = rollingFriction;
        component["spinningFriction"] = spinningFriction;
        component["cor"] = restitution;
        component["linearDamping"] = linearDamping;
        component["angularDamping"] = angularDamping;
        component["kinematic"] = kinematic;
        return component;
    }

    bool RigidBody::IsKinematic() const {
        return kinematic;
    }

    float RigidBody::GetFriction() const {
        return friction;
    }

    float RigidBody::GetRollingFriction() const {
        return rollingFriction;
    }

    float RigidBody::GetSpinningFriction() const {
        return spinningFriction;
    }

    float RigidBody::GetRestitution() const {
        return restitution;
    }

    float RigidBody::GetLinearDamping() const {
        return linearDamping;
    }

    float RigidBody::GetAngularDamping() const {
        return angularDamping;
    }

    btRigidBody* RigidBody::GetBulletRigidBody() {
        return rigidBody;
    }

    void RigidBody::NewBulletRigidBody(float mass) {
        Destroy();

        this->mass = mass;

        // Motion states inform us of movement caused by physics so that we can
        // deal with those changes as needed.
        btDefaultMotionState* motionState = new btDefaultMotionState(btTransform(btQuaternion(0, 0, 0, 1), btVector3(0, 0, 0)));

        // Bullet treats zero mass as infinite, resulting in immovable objects.
        btRigidBody::btRigidBodyConstructionInfo constructionInfo(mass, motionState, nullptr, btVector3(0, 0, 0));
        rigidBody = new btRigidBody(constructionInfo);
        rigidBody->setUserPointer(this);

        ghostObject = new btGhostObject();
    }

    void RigidBody::Destroy() {
        if (rigidBody) {
            delete rigidBody->getMotionState();
            delete rigidBody;
        }

        delete ghostObject;
    }

    glm::vec3 RigidBody::GetPosition() const {
        btTransform trans;
        if (IsKinematic())
            rigidBody->getMotionState()->getWorldTransform(trans);
        else
            trans = rigidBody->getWorldTransform();

        return Physics::btToGlm(trans.getOrigin());
    }

    void RigidBody::SetPosition(const glm::vec3& pos) {
        if (IsKinematic()) {
            btTransform trans;
            rigidBody->getMotionState()->getWorldTransform(trans);
            trans.setOrigin(Physics::glmToBt(pos));
            rigidBody->getMotionState()->setWorldTransform(trans);
        } else {
            btTransform trans = rigidBody->getWorldTransform();
            trans.setOrigin(Physics::glmToBt(pos));
            rigidBody->setWorldTransform(trans);
        }
    }

    glm::quat RigidBody::GetOrientation() const {
        btTransform trans;
        if (IsKinematic())
            rigidBody->getMotionState()->getWorldTransform(trans);
        else
            trans = rigidBody->getWorldTransform();

        return Physics::btToGlm(trans.getRotation());
    }

    void RigidBody::SetOrientation(const glm::quat& rotation) {
        if (IsKinematic()) {
            btTransform trans;
            rigidBody->getMotionState()->getWorldTransform(trans);
            trans.setRotation(Physics::glmToBt(rotation));
            rigidBody->getMotionState()->setWorldTransform(trans);
        } else {
            btTransform trans = rigidBody->getWorldTransform();
            trans.setRotation(Physics::glmToBt(rotation));
            rigidBody->setWorldTransform(trans);
        }
    }

    float RigidBody::GetMass() {
        return mass;
    }

    void RigidBody::SetMass(float mass) {
        // Bullet provides a method on the shape that we can use to calculate
        // inertia.
        btVector3 inertia;
        rigidBody->getCollisionShape()->calculateLocalInertia(mass, inertia);
        rigidBody->setMassProps(mass, inertia);
        this->mass = mass;
    }

    void RigidBody::SetFriction(float friction) {
        rigidBody->setFriction(friction);
        this->friction = friction;
    }

    void RigidBody::SetRollingFriction(float friction) {
        rigidBody->setRollingFriction(friction);
        this->rollingFriction = friction;
    }

    void RigidBody::SetSpinningFriction(float friction) {
        rigidBody->setSpinningFriction(friction);
        this->spinningFriction = friction;
    }

    void RigidBody::SetRestitution(float cor) {
        rigidBody->setRestitution(cor);
        this->restitution = cor;
    }

    void RigidBody::SetLinearDamping(float damping) {
        rigidBody->setDamping(damping, angularDamping);
        this->linearDamping = damping;
    }

    void RigidBody::SetAngularDamping(float damping) {
        rigidBody->setDamping(linearDamping, damping);
        this->angularDamping = damping;
    }

    void RigidBody::MakeKinematic() {
        rigidBody->setCollisionFlags(rigidBody->getCollisionFlags() | btCollisionObject::CF_KINEMATIC_OBJECT);
        kinematic = true;
        ghost = false;
    }

    void RigidBody::MakeDynamic() {
        rigidBody->setCollisionFlags(rigidBody->getCollisionFlags() & ~btCollisionObject::CF_KINEMATIC_OBJECT);
        kinematic = false;
        ghost = false;
    }

    void RigidBody::SetGhost(bool ghost) {
        this->ghost = ghost;
    }

    bool RigidBody::GetForceTransformSync() const {
        return forceTransformSync;
    }

    void RigidBody::SetForceTransformSync(bool sync) {
        forceTransformSync = sync;
    }

    bool RigidBody::GetHaltMovement() const {
        return haltMovement;
    }

    void RigidBody::SetHaltMovement(bool halt) {
        haltMovement = halt;
    }
}
<|endoftext|>


<commit_before>/*
    This file is part of KAddressbook.
    Copyright (c) 2000 Cornelius Schumacher <schumacher@kde.org>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

    As a special exception, permission is given to link this program
    with any edition of Qt, and distribute the resulting executable,
    without including the source code for Qt in the source distribution.
*/

#include <qlayout.h>

#include <kaction.h>
#include <kapplication.h>
#include <kdebug.h>
#include <kiconloader.h>
#include <kinstance.h>
#include <klocale.h>
#include <kparts/genericfactory.h>
#include <kparts/statusbarextension.h>
#include <kstatusbar.h>

#include "kabcore.h"
#include "kabprefs.h"
#include "kaddressbookiface.h"

#include "kaddressbook_part.h"

typedef KParts::GenericFactory< KAddressbookPart > KAddressbookFactory;
K_EXPORT_COMPONENT_FACTORY( libkaddressbookpart, KAddressbookFactory )

KAddressbookPart::KAddressbookPart( QWidget *parentWidget, const char *widgetName,
                                    QObject *parent, const char *name,
                                    const QStringList & )
  : DCOPObject( "KAddressBookIface" ), KParts::ReadOnlyPart( parent, name )
{
  kdDebug(5720) << "KAddressbookPart()" << endl;
  kdDebug(5720) << "  InstanceName: " << kapp->instanceName() << endl;

  setInstance( KAddressbookFactory::instance() );

  kdDebug(5720) << "KAddressbookPart()..." << endl;
  kdDebug(5720) << "  InstanceName: " << kapp->instanceName() << endl;

  // create a canvas to insert our widget
  QWidget *canvas = new QWidget( parentWidget, widgetName );
  canvas->setFocusPolicy( QWidget::ClickFocus );
  setWidget( canvas );

  QVBoxLayout *topLayout = new QVBoxLayout( canvas );

  KGlobal::iconLoader()->addAppDir( "kaddressbook" );

  mCore = new KABCore( this, true, canvas );
  mCore->restoreSettings();
  topLayout->addWidget( mCore->widget() );

  KParts::StatusBarExtension *statusBar = new KParts::StatusBarExtension( this );
  mCore->setStatusBar( statusBar->statusBar() );

  setXMLFile( "kaddressbook_part.rc" );
}

KAddressbookPart::~KAddressbookPart()
{
  mCore->save();
  mCore->saveSettings();

  KABPrefs::instance()->writeConfig();
  closeURL();
}

KAboutData *KAddressbookPart::createAboutData()
{
  return KABCore::createAboutData();
}

void KAddressbookPart::addEmail( QString addr )
{
  mCore->addEmail( addr );
}

void KAddressbookPart::importVCard( const QString& vCardURL )
{
  mCore->importVCard( vCardURL );
}

ASYNC KAddressbookPart::showContactEditor( QString uid )
{
  mCore->editContact( uid );
}

void KAddressbookPart::newContact()
{
  mCore->newContact();
}

QString KAddressbookPart::getNameByPhone( QString phone )
{
  return mCore->getNameByPhone( phone );
}

void KAddressbookPart::save()
{
  mCore->save();
}

void KAddressbookPart::exit()
{
  mCore->queryClose();

  delete this;
}

bool KAddressbookPart::openURL( const KURL &url )
{
  kdDebug(5720) << "KAddressbookPart:openFile()" << endl;

  mCore->widget()->show();

  if ( !url.isEmpty() )
    mCore->importVCard( url );

  emit setWindowCaption( url.prettyURL() );

  return true;
}

bool KAddressbookPart::openFile()
{
  return false;
}

bool KAddressbookPart::handleCommandLine()
{
  return mCore->handleCommandLine( this );
}

void KAddressbookPart::guiActivateEvent( KParts::GUIActivateEvent *e )
{
  kdDebug(5720) << "KAddressbookPart::guiActivateEvent" << endl;
  KParts::ReadOnlyPart::guiActivateEvent( e );
}

#include "kaddressbook_part.moc"
<commit_msg>SVN_SILENT remove debug statements<commit_after>/*
    This file is part of KAddressbook.
    Copyright (c) 2000 Cornelius Schumacher <schumacher@kde.org>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

    As a special exception, permission is given to link this program
    with any edition of Qt, and distribute the resulting executable,
    without including the source code for Qt in the source distribution.
*/

#include <qlayout.h>

#include <kaction.h>
#include <kapplication.h>
#include <kdebug.h>
#include <kiconloader.h>
#include <kinstance.h>
#include <klocale.h>
#include <kparts/genericfactory.h>
#include <kparts/statusbarextension.h>
#include <kstatusbar.h>

#include "kabcore.h"
#include "kabprefs.h"
#include "kaddressbookiface.h"

#include "kaddressbook_part.h"

typedef KParts::GenericFactory< KAddressbookPart > KAddressbookFactory;
K_EXPORT_COMPONENT_FACTORY( libkaddressbookpart, KAddressbookFactory )

KAddressbookPart::KAddressbookPart( QWidget *parentWidget, const char *widgetName,
                                    QObject *parent, const char *name,
                                    const QStringList & )
  : DCOPObject( "KAddressBookIface" ), KParts::ReadOnlyPart( parent, name )
{
  setInstance( KAddressbookFactory::instance() );

  // create a canvas to insert our widget
  QWidget *canvas = new QWidget( parentWidget, widgetName );
  canvas->setFocusPolicy( QWidget::ClickFocus );
  setWidget( canvas );

  QVBoxLayout *topLayout = new QVBoxLayout( canvas );

  KGlobal::iconLoader()->addAppDir( "kaddressbook" );

  mCore = new KABCore( this, true, canvas );
  mCore->restoreSettings();
  topLayout->addWidget( mCore->widget() );

  KParts::StatusBarExtension *statusBar = new KParts::StatusBarExtension( this );
  mCore->setStatusBar( statusBar->statusBar() );

  setXMLFile( "kaddressbook_part.rc" );
}

KAddressbookPart::~KAddressbookPart()
{
  mCore->save();
  mCore->saveSettings();

  KABPrefs::instance()->writeConfig();
  closeURL();
}

KAboutData *KAddressbookPart::createAboutData()
{
  return KABCore::createAboutData();
}

void KAddressbookPart::addEmail( QString addr )
{
  mCore->addEmail( addr );
}

void KAddressbookPart::importVCard( const QString& vCardURL )
{
  mCore->importVCard( vCardURL );
}

ASYNC KAddressbookPart::showContactEditor( QString uid )
{
  mCore->editContact( uid );
}

void KAddressbookPart::newContact()
{
  mCore->newContact();
}

QString KAddressbookPart::getNameByPhone( QString phone )
{
  return mCore->getNameByPhone( phone );
}

void KAddressbookPart::save()
{
  mCore->save();
}

void KAddressbookPart::exit()
{
  mCore->queryClose();

  delete this;
}

bool KAddressbookPart::openURL( const KURL &url )
{
  kdDebug(5720) << "KAddressbookPart:openFile()" << endl;

  mCore->widget()->show();

  if ( !url.isEmpty() )
    mCore->importVCard( url );

  emit setWindowCaption( url.prettyURL() );

  return true;
}

bool KAddressbookPart::openFile()
{
  return false;
}

bool KAddressbookPart::handleCommandLine()
{
  return mCore->handleCommandLine( this );
}

void KAddressbookPart::guiActivateEvent( KParts::GUIActivateEvent *e )
{
  kdDebug(5720) << "KAddressbookPart::guiActivateEvent" << endl;
  KParts::ReadOnlyPart::guiActivateEvent( e );
}

#include "kaddressbook_part.moc"
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/launch_dimensions.h"

#include <ostream>
#include <string>

#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/core/platform/logging.h"

namespace xla {
namespace gpu {

std::ostream& operator<<(std::ostream& out,
                         const LaunchDimensions& launch_dims) {
  LaunchDimensions::Dim3D block_counts = launch_dims.block_counts();
  LaunchDimensions::Dim3D thread_counts = launch_dims.thread_counts_per_block();
  out << absl::StrFormat("[block: {%d, %d, %d}, thread: {%d, %d, %d}] last_dim_aligned_for_vectorization: %d",
                         block_counts.x, block_counts.y, block_counts.z,
                         thread_counts.x, thread_counts.y, thread_counts.z,
                         launch_dims.last_dim_aligned_for_vectorization());
  return out;
}

static int64 ThreadsPerBlockLimit(GpuDeviceInfo gpu_device_info) {
  int64 threads_per_block = gpu_device_info.threads_per_block_limit;
  if (threads_per_block <= 0) {
    static std::atomic<int64> log_count{0};
    if (log_count.fetch_add(1) < 8) {
      LOG(WARNING) << "Attempting to calculate launch dimensions for GPU "
                      "without full information about its capabilities.  "
                      "StreamExecutor's PopulateDeviceDescription should be "
                      "updated for this device.";
    }
    threads_per_block = gpu_device_info.threads_per_warp;
    if (threads_per_block == 0) {
      // Fall back to *something* if we can't even get num threads per warp.
      threads_per_block = 32;
    }
  }
  return threads_per_block;
}

// Calculates the launch dimensions used to invoke `hlo`.
StatusOr<LaunchDimensions> CalculateLaunchDimensions(
    const Shape& shape, GpuDeviceInfo gpu_device_info,
    LaunchDimensionsConfig dim_config) {
  int64 num_elements = ShapeUtil::ElementsIn(shape);
  if (num_elements <= 1) {
    return LaunchDimensions();
  }

  CHECK_EQ(num_elements % dim_config.unroll_factor, 0);
  num_elements = num_elements / dim_config.unroll_factor;

  // Since we don't do any inter-warp communication, we're free to choose any
  // block size we want, subject to hardware constraints.  We choose the largest
  // block size allowed, as empirically, this is a performance win on almost
  // (but not all) benchmarks.
  //
  // My guess is that using a larger block size encourages ptxas to decrease
  // per-thread register usage, thus allowing for higher occupancy, but I
  // haven't verified this.
  //
  // TODO(jlebar): Investigate this further, and tune this heuristic so we can
  // run faster on the few benchmarks where smaller block size helps.
  int64 threads_per_block = ThreadsPerBlockLimit(gpu_device_info);
  // We unroll kernels to make use of vectorized loads/stores. This means we
  // need more registers to hold intermediate values. Reduce the number of
  // blocks per thread to increase the number of registers available to ptxas.
  // Make sure we still have a multiple of 32.
  int64 threads_per_block_row_vectorized = shape.dimensions().back() / dim_config.unroll_factor;
  if (dim_config.row_vectorized &&
      shape.dimensions().back() % dim_config.unroll_factor == 0 &&
      // If the row size is a multiple of 256, then use the old code
      // path that use a block size of 256. This give small speed up on V100.
      // Vectorization of the row load was already happening.
      (shape.dimensions().back() % 256) != 0 &&
      // Do not trigger the row vectorized codepath if this create too
      // small block size as this hurt performance.
      (threads_per_block_row_vectorized >= 128 &&
       threads_per_block_row_vectorized <= gpu_device_info.threads_per_block_limit)) {
    threads_per_block = threads_per_block_row_vectorized;
    VLOG(2) << "Update # of threads per block to ("
              << threads_per_block << ") to be row_vectorized.";
  } else {
    threads_per_block =
      RoundUpToNearest(threads_per_block / dim_config.unroll_factor, int64{32});
    if (num_elements < threads_per_block) {
      threads_per_block = num_elements;
      VLOG(2) << "Update # of threads per block to the element count ("
              << threads_per_block << ") because the latter is smaller.";
    }
  }

  int64 block_count = CeilOfRatio(num_elements, threads_per_block);
  if (dim_config.few_waves) {
    int64 capped_threads_per_block = std::min<int64>(threads_per_block, 128);
    int64 capped_block_count =
        gpu_device_info.core_count *
        (gpu_device_info.threads_per_core_limit / capped_threads_per_block);
    // Do not increase the number of blocks. This can happens for
    // small num_elements.
    if (capped_block_count < block_count) {
      threads_per_block = capped_threads_per_block;
      block_count = capped_block_count;
    }
  }

  if (gpu_device_info.block_dim_limit_x > 0 &&
      block_count >= gpu_device_info.block_dim_limit_x) {
    return tensorflow::errors::Unimplemented(
        "Kernel launch needs more blocks (", block_count,
        ") than allowed by hardware (", gpu_device_info.block_dim_limit_x,
        ").");
  }

  VLOG(2) << absl::StrFormat(
      "Initialized the block count to ceil(# of elements / threads per "
      "block) = ceil(%d/%d) = %d",
      num_elements, threads_per_block, block_count);

  bool last_dim_aligned_for_vectorization = false;
  if (dim_config.row_vectorized && shape.rank() > 1) {
    last_dim_aligned_for_vectorization = threads_per_block * dim_config.unroll_factor ==
        shape.dimensions().back();
  }

  return LaunchDimensions(block_count, threads_per_block, last_dim_aligned_for_vectorization);
}

}  // namespace gpu
}  // namespace xla
<commit_msg>Move comment to its right place.<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/launch_dimensions.h"

#include <ostream>
#include <string>

#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/core/platform/logging.h"

namespace xla {
namespace gpu {

std::ostream& operator<<(std::ostream& out,
                         const LaunchDimensions& launch_dims) {
  LaunchDimensions::Dim3D block_counts = launch_dims.block_counts();
  LaunchDimensions::Dim3D thread_counts = launch_dims.thread_counts_per_block();
  out << absl::StrFormat("[block: {%d, %d, %d}, thread: {%d, %d, %d}] last_dim_aligned_for_vectorization: %d",
                         block_counts.x, block_counts.y, block_counts.z,
                         thread_counts.x, thread_counts.y, thread_counts.z,
                         launch_dims.last_dim_aligned_for_vectorization());
  return out;
}

static int64 ThreadsPerBlockLimit(GpuDeviceInfo gpu_device_info) {
  int64 threads_per_block = gpu_device_info.threads_per_block_limit;
  if (threads_per_block <= 0) {
    static std::atomic<int64> log_count{0};
    if (log_count.fetch_add(1) < 8) {
      LOG(WARNING) << "Attempting to calculate launch dimensions for GPU "
                      "without full information about its capabilities.  "
                      "StreamExecutor's PopulateDeviceDescription should be "
                      "updated for this device.";
    }
    threads_per_block = gpu_device_info.threads_per_warp;
    if (threads_per_block == 0) {
      // Fall back to *something* if we can't even get num threads per warp.
      threads_per_block = 32;
    }
  }
  return threads_per_block;
}

// Calculates the launch dimensions used to invoke `hlo`.
StatusOr<LaunchDimensions> CalculateLaunchDimensions(
    const Shape& shape, GpuDeviceInfo gpu_device_info,
    LaunchDimensionsConfig dim_config) {
  int64 num_elements = ShapeUtil::ElementsIn(shape);
  if (num_elements <= 1) {
    return LaunchDimensions();
  }

  CHECK_EQ(num_elements % dim_config.unroll_factor, 0);
  num_elements = num_elements / dim_config.unroll_factor;

  // Since we don't do any inter-warp communication, we're free to choose any
  // block size we want, subject to hardware constraints.  We choose the largest
  // block size allowed, as empirically, this is a performance win on almost
  // (but not all) benchmarks.
  //
  // My guess is that using a larger block size encourages ptxas to decrease
  // per-thread register usage, thus allowing for higher occupancy, but I
  // haven't verified this.
  //
  // TODO(jlebar): Investigate this further, and tune this heuristic so we can
  // run faster on the few benchmarks where smaller block size helps.
  int64 threads_per_block = ThreadsPerBlockLimit(gpu_device_info);
  int64 threads_per_block_row_vectorized = shape.dimensions().back() / dim_config.unroll_factor;
  if (dim_config.row_vectorized &&
      shape.dimensions().back() % dim_config.unroll_factor == 0 &&
      // If the row size is a multiple of 256, then use the old code
      // path that use a block size of 256. This give small speed up on V100.
      // Vectorization of the row load was already happening.
      (shape.dimensions().back() % 256) != 0 &&
      // Do not trigger the row vectorized codepath if this create too
      // small block size as this hurt performance.
      (threads_per_block_row_vectorized >= 128 &&
       threads_per_block_row_vectorized <= gpu_device_info.threads_per_block_limit)) {
    threads_per_block = threads_per_block_row_vectorized;
    VLOG(2) << "Update # of threads per block to ("
              << threads_per_block << ") to be row_vectorized.";
  } else {
    // We unroll kernels to make use of vectorized loads/stores. This means we
    // need more registers to hold intermediate values. Reduce the number of
    // blocks per thread to increase the number of registers available to ptxas.
    // Make sure we still have a multiple of 32.
    threads_per_block =
      RoundUpToNearest(threads_per_block / dim_config.unroll_factor, int64{32});
    if (num_elements < threads_per_block) {
      threads_per_block = num_elements;
      VLOG(2) << "Update # of threads per block to the element count ("
              << threads_per_block << ") because the latter is smaller.";
    }
  }

  int64 block_count = CeilOfRatio(num_elements, threads_per_block);
  if (dim_config.few_waves) {
    int64 capped_threads_per_block = std::min<int64>(threads_per_block, 128);
    int64 capped_block_count =
        gpu_device_info.core_count *
        (gpu_device_info.threads_per_core_limit / capped_threads_per_block);
    // Do not increase the number of blocks. This can happens for
    // small num_elements.
    if (capped_block_count < block_count) {
      threads_per_block = capped_threads_per_block;
      block_count = capped_block_count;
    }
  }

  if (gpu_device_info.block_dim_limit_x > 0 &&
      block_count >= gpu_device_info.block_dim_limit_x) {
    return tensorflow::errors::Unimplemented(
        "Kernel launch needs more blocks (", block_count,
        ") than allowed by hardware (", gpu_device_info.block_dim_limit_x,
        ").");
  }

  VLOG(2) << absl::StrFormat(
      "Initialized the block count to ceil(# of elements / threads per "
      "block) = ceil(%d/%d) = %d",
      num_elements, threads_per_block, block_count);

  bool last_dim_aligned_for_vectorization = false;
  if (dim_config.row_vectorized && shape.rank() > 1) {
    last_dim_aligned_for_vectorization = threads_per_block * dim_config.unroll_factor ==
        shape.dimensions().back();
  }

  return LaunchDimensions(block_count, threads_per_block, last_dim_aligned_for_vectorization);
}

}  // namespace gpu
}  // namespace xla
<|endoftext|>


<commit_before>/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/tpu/kernels/tpu_embedding_enqueue_ops.h"

#include <string>
#include <vector>

#include "tensorflow/c/tf_tensor_internal.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/tpu/c_api_decl.h"
#include "tensorflow/stream_executor/tpu/status_helper.h"

namespace tensorflow {

Status ValidateCombiners(absl::Span<const std::string> combiners) {
  for (const std::string& combiner : combiners) {
    if (combiner != "sum" && combiner != "mean" && combiner != "sqrtn") {
      return errors::InvalidArgument(
          "Invalid combiner: only \"sum\", \"mean\", and "
          "\"sqrtn\" are supported.");
    }
  }
  return Status::OK();
}

Status GetValidatedModeOverride(const string& mode_override,
                                tpu::TPUEmbeddingConfiguration::Mode* mode) {
  if (mode_override == "train") {
    *mode = tpu::TPUEmbeddingConfiguration::TRAINING;
  } else if (mode_override == "inference") {
    *mode = tpu::TPUEmbeddingConfiguration::INFERENCE;
  } else if (mode_override == "unspecified") {
    *mode = tpu::TPUEmbeddingConfiguration::UNSPECIFIED;
  } else {
    return errors::InvalidArgument("Unsupported value ", mode_override,
                                   " specified for mode_override.");
  }
  return Status::OK();
}

namespace {

// T1: The type of the sample_indices op input.
// T2: The type of the embedding_indices op input.
// T3: The type of the aggregation_weights op input.
template <typename T1, typename T2, typename T3>
class EnqueueTPUEmbeddingArbitraryTensorBatchOp : public OpKernel {
 public:
  explicit EnqueueTPUEmbeddingArbitraryTensorBatchOp(OpKernelConstruction* ctx)
      : OpKernel(ctx) {
    if (ctx->HasAttr("device_ordinal")) {
      OP_REQUIRES_OK(ctx, ctx->GetAttr("device_ordinal", &device_ordinal_));
      device_ordinal_set_ = true;
    }

    std::vector<std::string> combiners;
    OP_REQUIRES_OK(ctx, ctx->GetAttr("combiners", &combiners));
    OP_REQUIRES_OK(ctx, ValidateCombiners(combiners));
    TpuEmbedding_TensorBatchFixedState_Create_Params fixed_state_create_params;
    StatusHelper status;
    fixed_state_create_params.status = status.c_status;

    std::vector<char*> c_str_combiners;
    c_str_combiners.reserve(combiners.size());
    for (std::string combiner : combiners) {
      c_str_combiners.push_back(&combiner.front());
    }
    fixed_state_create_params.combiners_size = c_str_combiners.size();
    fixed_state_create_params.combiners = c_str_combiners.data();
    fixed_state_ = tpu::OpsApiFn()->TpuEmbeddingTensorBatchFixedState_CreateFn(
        &fixed_state_create_params);

    OP_REQUIRES_OK(ctx, status.status());
  }

  void Compute(OpKernelContext* ctx) override {
    VLOG(2) << "EnqueueTPUEmbeddingArbitraryTensorBatchOp::Compute";

    OpInputList sample_indices_or_row_splits_list;
    OP_REQUIRES_OK(ctx, ctx->input_list("sample_indices_or_row_splits",
                                        &sample_indices_or_row_splits_list));

    OpInputList embedding_indices_list;
    OP_REQUIRES_OK(
        ctx, ctx->input_list("embedding_indices", &embedding_indices_list));

    OpInputList aggregation_weights_list;
    OP_REQUIRES_OK(
        ctx, ctx->input_list("aggregation_weights", &aggregation_weights_list));

    const Tensor* mode_override;
    OP_REQUIRES_OK(ctx, ctx->input("mode_override", &mode_override));

    const string& mode_value = mode_override->scalar<tstring>()();
    tpu::TPUEmbeddingConfiguration::Mode mode;
    OP_REQUIRES_OK(ctx, GetValidatedModeOverride(mode_value, &mode));

    // Extract device_ordinal from input when op did not has it as attribute.
    if (!device_ordinal_set_) {
      const Tensor* device_ordinal_tensor;
      OP_REQUIRES_OK(ctx, ctx->input("device_ordinal", &device_ordinal_tensor));
      device_ordinal_ = device_ordinal_tensor->flat<int32>()(0);
      device_ordinal_set_ = true;
    }

    const int num_input_features = sample_indices_or_row_splits_list.size();

    if (num_input_features != embedding_indices_list.size() ||
        num_input_features != aggregation_weights_list.size()) {
      VLOG(0) << "EnqueueTPUEmbeddingArbitraryTensorBatchOp::Compute failed"
              << "All three lists must have num_input_features but "
              << "len(sample_indices_or_row_splits_list) == "
              << num_input_features << ", len(embedding_indices_list) == "
              << embedding_indices_list.size()
              << ", len(aggregation_weights_list) == "
              << aggregation_weights_list.size();
    }

    std::vector<TF_Tensor*> sample_indices_or_row_splits_tensors(
        num_input_features);
    std::vector<TF_Tensor*> embedding_indices_tensors(num_input_features);
    std::vector<TF_Tensor*> aggregation_weights_tensors(num_input_features);

    for (int i = 0; i < num_input_features; ++i) {
      Status tf_status;
      sample_indices_or_row_splits_tensors[i] =
          TF_TensorFromTensor(sample_indices_or_row_splits_list[i], &tf_status);
      OP_REQUIRES_OK(ctx, tf_status);
      embedding_indices_tensors[i] =
          TF_TensorFromTensor(embedding_indices_list[i], &tf_status);
      OP_REQUIRES_OK(ctx, tf_status);
      aggregation_weights_tensors[i] =
          TF_TensorFromTensor(aggregation_weights_list[i], &tf_status);
      OP_REQUIRES_OK(ctx, tf_status);
    }

    TpuEmbeddingEngine_EnqueueTensorBatch_Params params;
    params.sample_indices_tensors = sample_indices_or_row_splits_tensors.data();
    params.sample_indices_tensors_size = num_input_features;
    params.embedding_indices_tensors = embedding_indices_tensors.data();
    params.embedding_indices_tensors_size = num_input_features;
    params.aggregation_weights_tensors = aggregation_weights_tensors.data();
    params.aggregation_weights_tensors_size = num_input_features;

    StatusHelper status;
    params.status = status.c_status;
    params.status = status.c_status;
    params.fixed_state = fixed_state_;
    params.local_device_ordinal = device_ordinal_;
    params.mode = mode;

    tpu::OpsApiFn()->TpuEmbeddingEngine_EnqueueTensorBatchFn(&params);
    OP_REQUIRES_OK(ctx, status.status());

    VLOG(2) << "EnqueueTPUEmbeddingArbitraryTensorBatchOp::Compute done";
  }

  ~EnqueueTPUEmbeddingArbitraryTensorBatchOp() override {
    tpu::OpsApiFn()->TpuEmbeddingTensorBatchFixedState_DestroyFn(fixed_state_);
  }

 private:
  int device_ordinal_ = -1;
  bool device_ordinal_set_ = false;
  TpuEmbedding_TensorBatchFixedState* fixed_state_;

  TF_DISALLOW_COPY_AND_ASSIGN(EnqueueTPUEmbeddingArbitraryTensorBatchOp);
};

#ifdef LIBTPU_ON_GCE
// Arbitrary tensor batch.
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, double>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, double>);

// Arbitrary tensor batch dynamic op having device tensor as input.
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, double>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, double>);
#endif  // LIBTPU_ON_GCE
}  // namespace
}  // namespace tensorflow
<commit_msg>Deallocate TPU embedding enqueue Tensors<commit_after>/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/tpu/kernels/tpu_embedding_enqueue_ops.h"

#include <string>
#include <vector>

#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/c/tf_tensor_internal.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/tpu/c_api_decl.h"
#include "tensorflow/stream_executor/tpu/status_helper.h"

namespace tensorflow {

Status ValidateCombiners(absl::Span<const std::string> combiners) {
  for (const std::string& combiner : combiners) {
    if (combiner != "sum" && combiner != "mean" && combiner != "sqrtn") {
      return errors::InvalidArgument(
          "Invalid combiner: only \"sum\", \"mean\", and "
          "\"sqrtn\" are supported.");
    }
  }
  return Status::OK();
}

Status GetValidatedModeOverride(const string& mode_override,
                                tpu::TPUEmbeddingConfiguration::Mode* mode) {
  if (mode_override == "train") {
    *mode = tpu::TPUEmbeddingConfiguration::TRAINING;
  } else if (mode_override == "inference") {
    *mode = tpu::TPUEmbeddingConfiguration::INFERENCE;
  } else if (mode_override == "unspecified") {
    *mode = tpu::TPUEmbeddingConfiguration::UNSPECIFIED;
  } else {
    return errors::InvalidArgument("Unsupported value ", mode_override,
                                   " specified for mode_override.");
  }
  return Status::OK();
}

namespace {

// Deallocates all tensors in `tf_tensors`.
void DeleteTensors(std::vector<TF_Tensor*>& tf_tensors) {
  for (TF_Tensor* tf_tensor : tf_tensors) {
    TF_DeleteTensor(tf_tensor);
  }
  tf_tensors.clear();
}

// T1: The type of the sample_indices op input.
// T2: The type of the embedding_indices op input.
// T3: The type of the aggregation_weights op input.
template <typename T1, typename T2, typename T3>
class EnqueueTPUEmbeddingArbitraryTensorBatchOp : public OpKernel {
 public:
  explicit EnqueueTPUEmbeddingArbitraryTensorBatchOp(OpKernelConstruction* ctx)
      : OpKernel(ctx) {
    if (ctx->HasAttr("device_ordinal")) {
      OP_REQUIRES_OK(ctx, ctx->GetAttr("device_ordinal", &device_ordinal_));
      device_ordinal_set_ = true;
    }

    std::vector<std::string> combiners;
    OP_REQUIRES_OK(ctx, ctx->GetAttr("combiners", &combiners));
    OP_REQUIRES_OK(ctx, ValidateCombiners(combiners));
    TpuEmbedding_TensorBatchFixedState_Create_Params fixed_state_create_params;
    StatusHelper status;
    fixed_state_create_params.status = status.c_status;

    std::vector<char*> c_str_combiners;
    c_str_combiners.reserve(combiners.size());
    for (std::string combiner : combiners) {
      c_str_combiners.push_back(&combiner.front());
    }
    fixed_state_create_params.combiners_size = c_str_combiners.size();
    fixed_state_create_params.combiners = c_str_combiners.data();
    fixed_state_ = tpu::OpsApiFn()->TpuEmbeddingTensorBatchFixedState_CreateFn(
        &fixed_state_create_params);

    OP_REQUIRES_OK(ctx, status.status());
  }

  void Compute(OpKernelContext* ctx) override {
    VLOG(2) << "EnqueueTPUEmbeddingArbitraryTensorBatchOp::Compute";

    OpInputList sample_indices_or_row_splits_list;
    OP_REQUIRES_OK(ctx, ctx->input_list("sample_indices_or_row_splits",
                                        &sample_indices_or_row_splits_list));

    OpInputList embedding_indices_list;
    OP_REQUIRES_OK(
        ctx, ctx->input_list("embedding_indices", &embedding_indices_list));

    OpInputList aggregation_weights_list;
    OP_REQUIRES_OK(
        ctx, ctx->input_list("aggregation_weights", &aggregation_weights_list));

    const Tensor* mode_override;
    OP_REQUIRES_OK(ctx, ctx->input("mode_override", &mode_override));

    const string& mode_value = mode_override->scalar<tstring>()();
    tpu::TPUEmbeddingConfiguration::Mode mode;
    OP_REQUIRES_OK(ctx, GetValidatedModeOverride(mode_value, &mode));

    // Extract device_ordinal from input when op did not has it as attribute.
    if (!device_ordinal_set_) {
      const Tensor* device_ordinal_tensor;
      OP_REQUIRES_OK(ctx, ctx->input("device_ordinal", &device_ordinal_tensor));
      device_ordinal_ = device_ordinal_tensor->flat<int32>()(0);
      device_ordinal_set_ = true;
    }

    const int num_input_features = sample_indices_or_row_splits_list.size();

    std::vector<TF_Tensor*> sample_indices_or_row_splits_tensors(
        num_input_features);
    std::vector<TF_Tensor*> embedding_indices_tensors(num_input_features);
    std::vector<TF_Tensor*> aggregation_weights_tensors(num_input_features);

    for (int i = 0; i < num_input_features; ++i) {
      Status tf_status;
      sample_indices_or_row_splits_tensors[i] =
          TF_TensorFromTensor(sample_indices_or_row_splits_list[i], &tf_status);
      OP_REQUIRES_OK(ctx, tf_status);
      embedding_indices_tensors[i] =
          TF_TensorFromTensor(embedding_indices_list[i], &tf_status);
      OP_REQUIRES_OK(ctx, tf_status);
      aggregation_weights_tensors[i] =
          TF_TensorFromTensor(aggregation_weights_list[i], &tf_status);
      OP_REQUIRES_OK(ctx, tf_status);
    }

    TpuEmbeddingEngine_EnqueueTensorBatch_Params params;
    params.sample_indices_tensors = sample_indices_or_row_splits_tensors.data();
    params.sample_indices_tensors_size = num_input_features;
    params.embedding_indices_tensors = embedding_indices_tensors.data();
    params.embedding_indices_tensors_size = num_input_features;
    params.aggregation_weights_tensors = aggregation_weights_tensors.data();
    params.aggregation_weights_tensors_size = num_input_features;

    StatusHelper status;
    params.status = status.c_status;
    params.status = status.c_status;
    params.fixed_state = fixed_state_;
    params.local_device_ordinal = device_ordinal_;
    params.mode = mode;

    tpu::OpsApiFn()->TpuEmbeddingEngine_EnqueueTensorBatchFn(&params);
    OP_REQUIRES_OK(ctx, status.status());

    DeleteTensors(sample_indices_or_row_splits_tensors);
    DeleteTensors(embedding_indices_tensors);
    DeleteTensors(aggregation_weights_tensors);

    VLOG(2) << "EnqueueTPUEmbeddingArbitraryTensorBatchOp::Compute done";
  }

  ~EnqueueTPUEmbeddingArbitraryTensorBatchOp() override {
    tpu::OpsApiFn()->TpuEmbeddingTensorBatchFixedState_DestroyFn(fixed_state_);
  }

 private:
  int device_ordinal_ = -1;
  bool device_ordinal_set_ = false;
  TpuEmbedding_TensorBatchFixedState* fixed_state_;

  TF_DISALLOW_COPY_AND_ASSIGN(EnqueueTPUEmbeddingArbitraryTensorBatchOp);
};

#ifdef LIBTPU_ON_GCE
// Arbitrary tensor batch.
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, double>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("EnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, double>);

// Arbitrary tensor batch dynamic op having device tensor as input.
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<float>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, float>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int32, double>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int32>("T1")
        .TypeConstraint<int64>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int32, int64, double>);
REGISTER_KERNEL_BUILDER(
    Name("DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")
        .TypeConstraint<int64>("T1")
        .TypeConstraint<int32>("T2")
        .TypeConstraint<double>("T3")
        .Device(DEVICE_CPU),
    EnqueueTPUEmbeddingArbitraryTensorBatchOp<int64, int32, double>);
#endif  // LIBTPU_ON_GCE
}  // namespace
}  // namespace tensorflow
<|endoftext|>