<commit_before>/*
*	Copyright 2009 Griffin Software
*
*	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 "MetalScrollPCH.h"
#include "Connect.h"
#include "MetalBar.h"
#include "ColorChip.h"

extern CAddInModule _AtlModule;

CComPtr<EnvDTE80::DTE2>				g_dte;
CComPtr<EnvDTE::AddIn>				g_addInInstance;
CComPtr<IVsTextManager>				g_textMgr;

// {454DB430-837D-4435-B7C1-F503B93EDA04} is the GUID for the marker service we register in MetalScrollPackage.rgs.
static const GUID g_markerServGUID = { 0x454DB430, 0x837D, 0x4435, { 0xB7, 0xC1, 0xF5, 0x03, 0xB9, 0x3E, 0xDA, 0x04 } };
// {FF447164-72F9-42ed-A767-C737B86BBD0B} is the GUID for the marker type.
static const GUID g_markerTypeGUID = { 0xFF447164, 0x72F9, 0x42ed, { 0xA7, 0x67, 0xC7, 0x37, 0xB8, 0x6B, 0xBD, 0x0B } };

CConnect::CConnect()
{
	m_textMgrEventsCookie = 0;
}

bool CConnect::GetTextManagerEventsPlug(IConnectionPoint** connPt)
{
	CComQIPtr<IConnectionPointContainer> connPoints = g_textMgr;
	if(!connPoints)
		return false;

	HRESULT hr = connPoints->FindConnectionPoint(__uuidof(IVsTextManagerEvents), connPt);
	return SUCCEEDED(hr) && *connPt;
}

bool CConnect::GetTextViewEventsPlug(IConnectionPoint** connPt, IVsTextView* view)
{
	CComQIPtr<IConnectionPointContainer> connPoints = view;
	if(!connPoints)
		return false;

	HRESULT hr = connPoints->FindConnectionPoint(__uuidof(IVsTextViewEvents), connPt);
	return SUCCEEDED(hr) && *connPt;
}

STDMETHODIMP CConnect::OnConnection(IDispatch* application, ext_ConnectMode /*connectMode*/, IDispatch* addInInst, SAFEARRAY** /*custom*/)
{
	HRESULT hr = application->QueryInterface(__uuidof(EnvDTE80::DTE2), (void**)&g_dte);
	if(FAILED(hr))
		return hr;

	hr = addInInst->QueryInterface(__uuidof(EnvDTE::AddIn), (void**)&g_addInInstance);
	if(FAILED(hr))
		return hr;

	CComQIPtr<IServiceProvider> sp = g_dte;
	if(!sp)
		return E_NOINTERFACE;

	hr = sp->QueryService(SID_SVsTextManager, IID_IVsTextManager, (void**)&g_textMgr);
	if(FAILED(hr) || !g_textMgr)
		return E_NOINTERFACE;

	CComPtr<IConnectionPoint> connPt;
	if(!GetTextManagerEventsPlug(&connPt))
		return E_NOINTERFACE;

	hr = connPt->Advise((IVsTextManagerEvents*)this, &m_textMgrEventsCookie);
	if(FAILED(hr))
		return hr;

	CComPtr<IProfferService> profferServ;
	hr = sp->QueryService(SID_SProfferService, IID_IProfferService, (void**)&profferServ);
	if(FAILED(hr) || !profferServ)
		return E_NOINTERFACE;

	DWORD cookie;
	hr = profferServ->ProfferService(g_markerServGUID, (IServiceProvider*)this, &cookie);
	if(FAILED(hr))
		return hr;

	MetalBar::ReadSettings();
	ColorChip::Register();

	return S_OK;
}

STDMETHODIMP CConnect::OnDisconnection(ext_DisconnectMode /*removeMode*/, SAFEARRAY** /*custom*/)
{
	CComPtr<IConnectionPoint> textMgrEventsPlug;
	if(m_textMgrEventsCookie && GetTextManagerEventsPlug(&textMgrEventsPlug))
	{
		textMgrEventsPlug->Unadvise(m_textMgrEventsCookie);
		m_textMgrEventsCookie = 0;
	}

	MetalBar::RemoveAllBars();
	ColorChip::Unregister();

	// Give up the global pointers.
	g_textMgr = 0;
	g_addInInstance = 0;
	g_dte = 0;
	return S_OK;
}

void STDMETHODCALLTYPE CConnect::OnRegisterView(IVsTextView* view)
{
	// Unfortunately, the window hasn't been created at this point yet, so we can't get the HWND
	// here. Register an even handler to catch SetFocus(), and get the HWND from there. We'll remove
	// the handler after the first SetFocus() as we don't care about getting more events once we
	// have the HWND.
	CComPtr<IConnectionPoint> connPt;
	if(!GetTextViewEventsPlug(&connPt, view))
		return;
	DWORD cookie;
	connPt->Advise((IVsTextViewEvents*)this, &cookie);
}

struct ScrollbarHandles
{
	HWND		vert;
	HWND		horiz;
};

static BOOL CALLBACK FindScrollbars(HWND hwnd, LPARAM param)
{
	ScrollbarHandles* handles = (ScrollbarHandles*)param;

	char className[64];
	GetClassNameA(hwnd, className, _countof(className));
	if(strcmp(className, "ScrollBar") == 0)
	{
		LONG styles = GetWindowLong(hwnd, GWL_STYLE);
		if(styles & SBS_VERT)
			handles->vert = hwnd;
		else
			handles->horiz = hwnd;
	}

	return TRUE;
}

static LRESULT FAR PASCAL ScrollBarProc(HWND hwnd, UINT message, WPARAM wparam, LPARAM lparam)
{
	MetalBar* bar = (MetalBar*)GetWindowLongPtr(hwnd, GWL_USERDATA);
	if(!bar)
		return 0;

	return bar->WndProc(hwnd, message, wparam, lparam);
}

void CConnect::HookScrollbar(IVsTextView* view)
{
	HWND editorHwnd = view->GetWindowHandle();
	if(!editorHwnd)
		return;

	ScrollbarHandles bars = { 0 };
	EnumChildWindows(GetParent(editorHwnd), FindScrollbars, (LPARAM)&bars);
	if(!bars.horiz || !bars.vert)
		return;

	if(GetWindowLongPtr(bars.vert, GWL_USERDATA) != 0)
		return;

	WNDPROC oldProc = (WNDPROC)SetWindowLongPtr(bars.vert, GWLP_WNDPROC, (::LONG_PTR)ScrollBarProc);
	MetalBar* bar = new MetalBar(bars.vert, editorHwnd, bars.horiz, oldProc, view);
	SetWindowLongPtr(bars.vert, GWL_USERDATA, (::LONG_PTR)bar);
}

void STDMETHODCALLTYPE CConnect::OnSetFocus(IVsTextView* view)
{
	HookScrollbar(view);

	// Remove ourselves from the event list. Since we can't store the cookie we got
	// when we registered the event handler, we'll have to scan the event list looking
	// for our pointer.
	CComPtr<IConnectionPoint> connPt;
	if(!GetTextViewEventsPlug(&connPt, view))
		return;

	CComPtr<IEnumConnections> enumerator;
	HRESULT hr = connPt->EnumConnections(&enumerator);
	if(FAILED(hr) || !enumerator)
		return;

	CONNECTDATA connData;
	ULONG numRet;
	bool found = false;
	while( SUCCEEDED(enumerator->Next(1, &connData, &numRet)) && (numRet > 0) )
	{
		if(connData.pUnk == (IVsTextViewEvents*)this)
		{
			connPt->Unadvise(connData.dwCookie);
			found = true;
		}
		connData.pUnk->Release();
	}
}

STDMETHODIMP CConnect::QueryService(REFGUID guidService, REFIID riid, void** object)
{
	if(!object)
		return E_POINTER;

	if(!InlineIsEqualGUID(guidService, g_markerServGUID))
		return E_NOINTERFACE;

	if(!IsEqualIID(riid, __uuidof(IVsTextMarkerTypeProvider)))
		return E_NOINTERFACE;

	*object = (IVsTextMarkerTypeProvider*)this;
	((IVsTextMarkerTypeProvider*)this)->AddRef();
	return S_OK;
}

STDMETHODIMP CConnect::GetTextMarkerType(GUID* guidMarker, IVsPackageDefinedTextMarkerType** markerType)
{
	if(!markerType)
		return E_POINTER;

	if(!InlineIsEqualGUID(*guidMarker, g_markerTypeGUID))
		return E_NOINTERFACE;

	*markerType = (IVsPackageDefinedTextMarkerType*)this;
	(*markerType)->AddRef();
	return S_OK;
}

STDMETHODIMP CConnect::GetVisualStyle(DWORD* visualFlags)
{
	*visualFlags = MV_COLOR_ALWAYS;
	return S_OK;
}

STDMETHODIMP CConnect::GetDefaultColors(COLORINDEX* foreground, COLORINDEX* background)
{
	*foreground = CI_BLACK;
	*background = CI_CYAN;
	return S_OK;
}

STDMETHODIMP CConnect::GetDefaultLineStyle(COLORINDEX* lineColor, LINESTYLE* lineIndex)
{
	// This shouldn't be called because we haven't requested MV_LINE or MV_BORDER.
	*lineColor = CI_BLACK;
	*lineIndex = LI_SOLID;
	return S_OK;
}

STDMETHODIMP CConnect::GetDefaultFontFlags(DWORD* fontFlags)
{
	*fontFlags = FF_DEFAULT;
	return S_OK;
}

STDMETHODIMP CConnect::DrawGlyphWithColors(HDC /*hdc*/, RECT* /*rect*/, long /*markerType*/, IVsTextMarkerColorSet* /*markerColors*/, DWORD /*glyphDrawFlags*/, long /*lineHeight*/)
{
	// This shouldn't be called because we haven't requested MV_GLYPH.
	return S_OK;
}

STDMETHODIMP CConnect::GetBehaviorFlags(DWORD* flags)
{
	*flags = MB_DEFAULT;
	return S_OK;
}

STDMETHODIMP CConnect::GetPriorityIndex(long* priorityIndex)
{
	*priorityIndex = 100;
	return S_OK;
}
<commit_msg>COM paranoia.<commit_after>/*
*	Copyright 2009 Griffin Software
*
*	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 "MetalScrollPCH.h"
#include "Connect.h"
#include "MetalBar.h"
#include "ColorChip.h"

extern CAddInModule _AtlModule;

CComPtr<EnvDTE80::DTE2>				g_dte;
CComPtr<EnvDTE::AddIn>				g_addInInstance;
CComPtr<IVsTextManager>				g_textMgr;

// {454DB430-837D-4435-B7C1-F503B93EDA04} is the GUID for the marker service we register in MetalScrollPackage.rgs.
static const GUID g_markerServGUID = { 0x454DB430, 0x837D, 0x4435, { 0xB7, 0xC1, 0xF5, 0x03, 0xB9, 0x3E, 0xDA, 0x04 } };
// {FF447164-72F9-42ed-A767-C737B86BBD0B} is the GUID for the marker type.
static const GUID g_markerTypeGUID = { 0xFF447164, 0x72F9, 0x42ed, { 0xA7, 0x67, 0xC7, 0x37, 0xB8, 0x6B, 0xBD, 0x0B } };

CConnect::CConnect()
{
	m_textMgrEventsCookie = 0;
}

bool CConnect::GetTextManagerEventsPlug(IConnectionPoint** connPt)
{
	CComQIPtr<IConnectionPointContainer> connPoints = g_textMgr;
	if(!connPoints)
		return false;

	HRESULT hr = connPoints->FindConnectionPoint(__uuidof(IVsTextManagerEvents), connPt);
	return SUCCEEDED(hr) && *connPt;
}

bool CConnect::GetTextViewEventsPlug(IConnectionPoint** connPt, IVsTextView* view)
{
	CComQIPtr<IConnectionPointContainer> connPoints = view;
	if(!connPoints)
		return false;

	HRESULT hr = connPoints->FindConnectionPoint(__uuidof(IVsTextViewEvents), connPt);
	return SUCCEEDED(hr) && *connPt;
}

STDMETHODIMP CConnect::OnConnection(IDispatch* application, ext_ConnectMode /*connectMode*/, IDispatch* addInInst, SAFEARRAY** /*custom*/)
{
	HRESULT hr = application->QueryInterface(__uuidof(EnvDTE80::DTE2), (void**)&g_dte);
	if(FAILED(hr))
		return hr;

	hr = addInInst->QueryInterface(__uuidof(EnvDTE::AddIn), (void**)&g_addInInstance);
	if(FAILED(hr))
		return hr;

	CComQIPtr<IServiceProvider> sp = g_dte;
	if(!sp)
		return E_NOINTERFACE;

	hr = sp->QueryService(SID_SVsTextManager, IID_IVsTextManager, (void**)&g_textMgr);
	if(FAILED(hr) || !g_textMgr)
		return E_NOINTERFACE;

	CComPtr<IConnectionPoint> connPt;
	if(!GetTextManagerEventsPlug(&connPt))
		return E_NOINTERFACE;

	hr = connPt->Advise((IVsTextManagerEvents*)this, &m_textMgrEventsCookie);
	if(FAILED(hr))
		return hr;

	CComPtr<IProfferService> profferServ;
	hr = sp->QueryService(SID_SProfferService, IID_IProfferService, (void**)&profferServ);
	if(FAILED(hr) || !profferServ)
		return E_NOINTERFACE;

	DWORD cookie;
	hr = profferServ->ProfferService(g_markerServGUID, (IServiceProvider*)this, &cookie);
	if(FAILED(hr))
		return hr;

	MetalBar::ReadSettings();
	ColorChip::Register();

	return S_OK;
}

STDMETHODIMP CConnect::OnDisconnection(ext_DisconnectMode /*removeMode*/, SAFEARRAY** /*custom*/)
{
	CComPtr<IConnectionPoint> textMgrEventsPlug;
	if(m_textMgrEventsCookie && GetTextManagerEventsPlug(&textMgrEventsPlug))
	{
		textMgrEventsPlug->Unadvise(m_textMgrEventsCookie);
		m_textMgrEventsCookie = 0;
	}

	MetalBar::RemoveAllBars();
	ColorChip::Unregister();

	// Give up the global pointers.
	g_textMgr = 0;
	g_addInInstance = 0;
	g_dte = 0;
	return S_OK;
}

void STDMETHODCALLTYPE CConnect::OnRegisterView(IVsTextView* view)
{
	// Unfortunately, the window hasn't been created at this point yet, so we can't get the HWND
	// here. Register an even handler to catch SetFocus(), and get the HWND from there. We'll remove
	// the handler after the first SetFocus() as we don't care about getting more events once we
	// have the HWND.
	CComPtr<IConnectionPoint> connPt;
	if(!GetTextViewEventsPlug(&connPt, view))
		return;
	DWORD cookie;
	connPt->Advise((IVsTextViewEvents*)this, &cookie);
}

struct ScrollbarHandles
{
	HWND		vert;
	HWND		horiz;
};

static BOOL CALLBACK FindScrollbars(HWND hwnd, LPARAM param)
{
	ScrollbarHandles* handles = (ScrollbarHandles*)param;

	char className[64];
	GetClassNameA(hwnd, className, _countof(className));
	if(strcmp(className, "ScrollBar") == 0)
	{
		LONG styles = GetWindowLong(hwnd, GWL_STYLE);
		if(styles & SBS_VERT)
			handles->vert = hwnd;
		else
			handles->horiz = hwnd;
	}

	return TRUE;
}

static LRESULT FAR PASCAL ScrollBarProc(HWND hwnd, UINT message, WPARAM wparam, LPARAM lparam)
{
	MetalBar* bar = (MetalBar*)GetWindowLongPtr(hwnd, GWL_USERDATA);
	if(!bar)
		return 0;

	return bar->WndProc(hwnd, message, wparam, lparam);
}

void CConnect::HookScrollbar(IVsTextView* view)
{
	HWND editorHwnd = view->GetWindowHandle();
	if(!editorHwnd)
		return;

	ScrollbarHandles bars = { 0 };
	EnumChildWindows(GetParent(editorHwnd), FindScrollbars, (LPARAM)&bars);
	if(!bars.horiz || !bars.vert)
		return;

	if(GetWindowLongPtr(bars.vert, GWL_USERDATA) != 0)
		return;

	WNDPROC oldProc = (WNDPROC)SetWindowLongPtr(bars.vert, GWLP_WNDPROC, (::LONG_PTR)ScrollBarProc);
	MetalBar* bar = new MetalBar(bars.vert, editorHwnd, bars.horiz, oldProc, view);
	SetWindowLongPtr(bars.vert, GWL_USERDATA, (::LONG_PTR)bar);
}

void STDMETHODCALLTYPE CConnect::OnSetFocus(IVsTextView* view)
{
	HookScrollbar(view);

	// Remove ourselves from the event list. Since we can't store the cookie we got
	// when we registered the event handler, we'll have to scan the event list looking
	// for our pointer.
	CComPtr<IConnectionPoint> connPt;
	if(!GetTextViewEventsPlug(&connPt, view))
		return;

	CComPtr<IEnumConnections> enumerator;
	HRESULT hr = connPt->EnumConnections(&enumerator);
	if(FAILED(hr) || !enumerator)
		return;

	CONNECTDATA connData;
	ULONG numRet;
	bool found = false;
	while( SUCCEEDED(enumerator->Next(1, &connData, &numRet)) && (numRet > 0) )
	{
		if(connData.pUnk == (IVsTextViewEvents*)this)
		{
			connPt->Unadvise(connData.dwCookie);
			found = true;
		}
		connData.pUnk->Release();
	}
}

STDMETHODIMP CConnect::QueryService(REFGUID guidService, REFIID riid, void** object)
{
	if(!object)
		return E_POINTER;

	if(!InlineIsEqualGUID(guidService, g_markerServGUID))
		return E_NOINTERFACE;

	if(!IsEqualIID(riid, __uuidof(IVsTextMarkerTypeProvider)))
		return E_NOINTERFACE;

	*object = (IVsTextMarkerTypeProvider*)this;
	((IVsTextMarkerTypeProvider*)this)->AddRef();
	return S_OK;
}

STDMETHODIMP CConnect::GetTextMarkerType(GUID* guidMarker, IVsPackageDefinedTextMarkerType** markerType)
{
	if(!markerType)
		return E_POINTER;

	if(!InlineIsEqualGUID(*guidMarker, g_markerTypeGUID))
		return E_NOINTERFACE;

	*markerType = (IVsPackageDefinedTextMarkerType*)this;
	(*markerType)->AddRef();
	return S_OK;
}

STDMETHODIMP CConnect::GetVisualStyle(DWORD* visualFlags)
{
	if(!visualFlags)
		return E_POINTER;

	*visualFlags = MV_COLOR_ALWAYS;
	return S_OK;
}

STDMETHODIMP CConnect::GetDefaultColors(COLORINDEX* foreground, COLORINDEX* background)
{
	if(!foreground || !background)
		return E_POINTER;

	*foreground = CI_BLACK;
	*background = CI_CYAN;
	return S_OK;
}

STDMETHODIMP CConnect::GetDefaultLineStyle(COLORINDEX* lineColor, LINESTYLE* lineIndex)
{
	// This shouldn't be called because we haven't requested MV_LINE or MV_BORDER.
	if(!lineColor || !lineIndex)
		return E_POINTER;

	*lineColor = CI_BLACK;
	*lineIndex = LI_SOLID;
	return S_OK;
}

STDMETHODIMP CConnect::GetDefaultFontFlags(DWORD* fontFlags)
{
	if(!fontFlags)
		return E_POINTER;

	*fontFlags = FF_DEFAULT;
	return S_OK;
}

STDMETHODIMP CConnect::DrawGlyphWithColors(HDC /*hdc*/, RECT* /*rect*/, long /*markerType*/, IVsTextMarkerColorSet* /*markerColors*/, DWORD /*glyphDrawFlags*/, long /*lineHeight*/)
{
	// This shouldn't be called because we haven't requested MV_GLYPH.
	return S_OK;
}

STDMETHODIMP CConnect::GetBehaviorFlags(DWORD* flags)
{
	if(!flags)
		return E_POINTER;

	*flags = MB_DEFAULT;
	return S_OK;
}

STDMETHODIMP CConnect::GetPriorityIndex(long* priorityIndex)
{
	if(!priorityIndex)
		return E_POINTER;

	*priorityIndex = 100;
	return S_OK;
}
<|endoftext|>


<commit_before><commit_msg>Dead function prototypes gone.<commit_after><|endoftext|>


<commit_before>// This file is part of Parsito <http://github.com/ufal/parsito/>.
//
// Copyright 2015 Institute of Formal and Applied Linguistics, Faculty of
// Mathematics and Physics, Charles University in Prague, Czech Republic.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#include <algorithm>
#include <atomic>
#include <fstream>
#include <limits>
#include <random>
#include <thread>
#include <unordered_set>

#include "network/neural_network_trainer.h"
#include "parser_nn.h"
#include "parser_nn_trainer.h"
#include "utils/parse_double.h"
#include "utils/parse_int.h"
#include "utils/split.h"

namespace ufal {
namespace parsito {

void parser_nn_trainer::train(const string& transition_system_name, const string& transition_oracle_name,
                              const string& embeddings_description, const string& nodes_description, const network_parameters& parameters,
                              unsigned number_of_threads, const vector<tree>& train, const vector<tree>& heldout, binary_encoder& enc) {
  if (train.empty()) runtime_failure("No training data was given!");

  // Random generator with fixed seed for reproducibility
  mt19937 generator(42);

  // Check that all non-root nodes have heads and nonempty deprel
  for (auto&& tree : train)
    for (auto&& node : tree.nodes)
      if (node.id) {
        if (node.head < 0) runtime_failure("The node '" << node.form << "' with id " << node.id << " has no head set!");
        if (node.deprel.empty()) runtime_failure("The node '" << node.form << "' with id " << node.id << " has no deprel set!");
      }

  // Generate labels for transition system
  parser_nn parser;
  unordered_set<string> labels_set;

  for (auto&& tree : train)
    for (auto&& node : tree.nodes)
      if (node.id && !labels_set.count(node.deprel)) {
        labels_set.insert(node.deprel);
        parser.labels.push_back(node.deprel);
      }

  // Create transition system and transition oracle
  parser.system.reset(transition_system::create(transition_system_name, parser.labels));
  if (!parser.system) runtime_failure("Cannot create transition system '" << transition_system_name << "'!");

  unique_ptr<transition_oracle> oracle(parser.system->oracle(transition_oracle_name));
  if (!oracle) runtime_failure("Cannot create transition oracle '" << transition_oracle_name << "' for transition system '" << transition_system_name << "'!");

  // Create node_extractor
  string error;
  if (!parser.nodes.create(nodes_description, error)) runtime_failure(error);

  // Load value_extractors and embeddings
  vector<string> value_names;
  vector<string_piece> lines, tokens;
  split(embeddings_description, '\n', lines);
  for (auto&& line : lines) {
    // Ignore empty lines and comments
    if (!line.len || line.str[0] == '#') continue;

    split(line, ' ', tokens);
    if (!(tokens.size() >= 3 && tokens.size() <= 6))
      runtime_failure("Expected 3 to 6 columns on embedding description line '" << line << "'!");

    value_names.emplace_back(string(tokens[0].str, tokens[0].len));
    parser.values.emplace_back();
    if (!parser.values.back().create(tokens[0], error)) runtime_failure(error);

    int dimension = parse_int(tokens[1], "embedding dimension");
    int min_count = parse_int(tokens[2], "minimum frequency count");
    unsigned updatable_index = 0;
    unsigned embeddings_from_file = 0;
    string embeddings_from_file_comment;
    vector<pair<string, vector<float>>> weights;
    unordered_set<string> weights_set;

    // Compute words and counts present in the training data
    string word;
    unordered_map<string, int> word_counts;
    for (auto&& tree : train)
      for (auto&& node : tree.nodes)
        if (node.id) {
          parser.values.back().extract(node, word);
          word_counts[word]++;
        }

    // Load embedding if it was given
    if (tokens.size() >= 4) {
      int update_weights = tokens.size() >= 5 ? parse_int(tokens[4], "update weights") : 1;
      int max_embeddings = tokens.size() >= 6 ? parse_int(tokens[5], "maximum embeddings count") : numeric_limits<int>::max();
      ifstream in(string(tokens[3].str, tokens[3].len));
      if (!in.is_open()) runtime_failure("Cannot load '" << tokens[0] << "' embedding from file '" << tokens[3] << "'!");

      // Load first line containing dictionary size and dimensions
      string line;
      vector<string_piece> parts;
      if (!getline(in, line)) runtime_failure("Cannot read first line from embedding file '" << tokens[3] << "'!");
      split(line, ' ', parts);
      if (parts.size() != 2) runtime_failure("Expected two numbers on the first line of embedding file '" << tokens[3] << "'!");
      int file_dimension = parse_int(parts[1], "embedding file dimension");

      if (file_dimension < dimension) runtime_failure("The embedding file '" << tokens[3] << "' has lower dimension than required!");

      // Generate random projection when smaller dimension is required
      vector<vector<float>> projection;
      if (file_dimension > dimension) {
        embeddings_from_file_comment = "[dim" + to_string(file_dimension) + "->" + to_string(dimension) + "]";

        uniform_real_distribution<double> uniform(0, 1);
        projection.resize(dimension);
        for (auto&& row : projection) {
          row.resize(file_dimension);
          for (auto&& weight : row) weight = uniform(generator);

          double sum = 0;
          for (auto&& weight : row) sum += weight;
          for (auto&& weight : row) weight /= sum;
        }
      }

      // Load input embedding
      vector<double> input_weights(file_dimension);
      vector<float> projected_weights(dimension);
      while (getline(in, line) && int(weights.size()) < max_embeddings) {
        split(line, ' ', parts);
        if (!parts.empty() && !parts.back().len) parts.pop_back(); // Ignore space at the end of line
        if (int(parts.size()) != file_dimension + 1) runtime_failure("Wrong number of values on line '" << line << "' of embedding file '" << tokens[3]);
        for (int i = 0; i < file_dimension; i++)
          input_weights[i] = parse_double(parts[1 + i], "embedding weight");

        string word(parts[0].str, parts[0].len);

        // For update_weights == 2, ignore embeddings for unknown words
        if (update_weights == 2 && !word_counts.count(word))
          continue;

        for (int i = 0; i < dimension; i++)
          if (file_dimension == dimension) {
            projected_weights[i] = input_weights[i];
          } else {
            projected_weights[i] = 0;
            for (int j = 0; j < file_dimension; j++)
              projected_weights[i] += projection[i][j] * input_weights[j];
          }

        if (!weights_set.count(word)) {
          weights.emplace_back(word, projected_weights);
          weights_set.insert(word);
        }
      }
      embeddings_from_file = weights.size();
      updatable_index = update_weights ? 0 : embeddings_from_file;
    }

    // Add embedding for non-present word with min_count
    {
      vector<float> word_weights(dimension);
      uniform_real_distribution<float> uniform(-1, 1);
      for (auto&& word_count : word_counts)
        if (word_count.second >= min_count && !weights_set.count(word_count.first)) {
          for (auto&& word_weight : word_weights)
            word_weight = uniform(generator);

          weights.emplace_back(word_count.first, word_weights);
        }
    }

    // If there are unknown words in the training data, create initial embedding
    vector<float> unknown_weights(dimension);
    if (min_count > 1) {
      uniform_real_distribution<float> uniform(-1, 1);

      for (auto&& weight : unknown_weights)
        weight = uniform(generator);
    }

    // Add the embedding
    parser.embeddings.emplace_back();
    parser.embeddings.back().create(dimension, updatable_index, weights, unknown_weights);

    // Count the cover of this embedding
    string buffer;
    unsigned words_total = 0, words_covered = 0, words_covered_from_file = 0;
    for (auto&& tree : train)
      for (auto&& node : tree.nodes)
        if (node.id) {
          parser.values.back().extract(node, word);
          words_total++;
          int word_id = parser.embeddings.back().lookup_word(word, buffer);
          words_covered += word_id != parser.embeddings.back().unknown_word();
          words_covered_from_file += word_id != parser.embeddings.back().unknown_word() && unsigned(word_id) < embeddings_from_file;
        }

    cerr << "Initialized '" << tokens[0] << "' embedding with " << embeddings_from_file << embeddings_from_file_comment
         << "," << weights.size() << " words and " << fixed << setprecision(1) << 100. * words_covered_from_file / words_total
         << "%," << 100. * words_covered / words_total << "% coverage." << endl;
  }

  // Train the network
  unsigned total_dimension = 0, total_nodes = 0;
  for (auto&& embedding : parser.embeddings) total_dimension += embedding.dimension;
  for (auto&& tree : train) total_nodes += tree.nodes.size() - 1;
  auto scaled_parameters = parameters;
  scaled_parameters.l1_regularization /= train.size();
  scaled_parameters.l2_regularization /= total_nodes;
  neural_network_trainer network_trainer(parser.network, total_dimension * parser.nodes.node_count(), parser.system->transition_count(), scaled_parameters, generator);

  vector<int> permutation;
  for (size_t i = 0; i < train.size(); i++)
    permutation.push_back(permutation.size());

  for (int iteration = 1; network_trainer.next_iteration(); iteration++) {
    // Train on training data
    shuffle(permutation.begin(), permutation.end(), generator);

    atomic<unsigned> atomic_index(0);
    atomic<double> atomic_logprob(0);
    auto training = [&]() {
      tree t;
      configuration conf;
      string word, word_buffer;
      vector<vector<int>> nodes_embeddings;
      vector<int> extracted_nodes;
      vector<const vector<int>*> extracted_embeddings;
      neural_network_trainer::workspace workspace;
      double logprob = 0;

      for (unsigned current_index; (current_index = atomic_index++) < permutation.size();) {
        const tree& gold = train[permutation[current_index]];
        t = gold;
        t.unlink_all_nodes();
        conf.init(&t);

        // Compute embeddings
        if (t.nodes.size() > nodes_embeddings.size()) nodes_embeddings.resize(t.nodes.size());
        for (size_t i = 0; i < t.nodes.size(); i++) {
          nodes_embeddings[i].resize(parser.embeddings.size());
          for (size_t j = 0; j < parser.embeddings.size(); j++) {
            parser.values[j].extract(t.nodes[i], word);
            nodes_embeddings[i][j] = parser.embeddings[j].lookup_word(word, word_buffer);
          }
        }

        // Train the network
        while (!conf.final()) {
          // Extract nodes
          parser.nodes.extract(conf, extracted_nodes);
          extracted_embeddings.resize(extracted_nodes.size());
          for (size_t i = 0; i < extracted_nodes.size(); i++)
            extracted_embeddings[i] = extracted_nodes[i] >= 0 ? &nodes_embeddings[extracted_nodes[i]] : nullptr;

          // Propagate
          network_trainer.propagate(parser.embeddings, extracted_embeddings, workspace);

          // Find most probable applicable transition
          int network_best = -1;
          for (unsigned i = 0; i < workspace.outcomes.size(); i++)
            if (parser.system->applicable(conf, i) && (network_best < 0 || workspace.outcomes[i] > workspace.outcomes[network_best]))
              network_best = i;

          // Apply the oracle
          auto prediction = oracle->predict(conf, gold, network_best);

          // Update logprob
          if (workspace.outcomes[prediction.best])
            logprob += log(workspace.outcomes[prediction.best]);

          // Backpropagate the chosen outcome
          network_trainer.backpropagate(parser.embeddings, extracted_embeddings, prediction.best, workspace);

          // Follow the chosen outcome
          int child = parser.system->perform(conf, prediction.to_follow);

          // If a node was linked, recompute its embeddings as deprel has changed
          if (child >= 0)
            for (size_t i = 0; i < parser.embeddings.size(); i++) {
              parser.values[i].extract(t.nodes[child], word);
              nodes_embeddings[child][i] = parser.embeddings[i].lookup_word(word, word_buffer);
            }
        }
        network_trainer.finalize_sentence();
      }
      for (double old_atomic_logprob = atomic_logprob; atomic_logprob.compare_exchange_weak(old_atomic_logprob, old_atomic_logprob + logprob); ) {}
    };

    cerr << "Iteration " << iteration << ": ";
    if (number_of_threads > 1) {
      vector<thread> threads;
      for (unsigned i = 0; i < number_of_threads; i++) threads.emplace_back(training);
      for (; !threads.empty(); threads.pop_back()) threads.back().join();
    } else {
      training();
    }
    cerr << "training logprob " << scientific << setprecision(4) << atomic_logprob;

    // Evaluate heldout data if present
    if (!heldout.empty()) {
      network_trainer.finalize_dropout_weights();

      tree t;
      unsigned total = 0, correct_unlabelled = 0, correct_labelled = 0;
      for (auto&& gold : heldout) {
        t = gold;
        t.unlink_all_nodes();
        parser.parse(t);
        for (size_t i = 1; i < t.nodes.size(); i++) {
          total++;
          correct_unlabelled += t.nodes[i].head == gold.nodes[i].head;
          correct_labelled += t.nodes[i].head == gold.nodes[i].head && t.nodes[i].deprel == gold.nodes[i].deprel;
        }
      }

      network_trainer.finalize_dropout_weights(false);

      cerr << ", heldout UAS " << fixed << setprecision(2) << (100. * correct_unlabelled / total) << "%, LAS " << (100. * correct_labelled / total) << "%";
    }

    cerr << endl;
  }
  network_trainer.finalize_dropout_weights();

  // Encode transition system
  enc.add_2B(parser.labels.size());
  for (auto&& label : parser.labels)
    enc.add_str(label);
  enc.add_str(transition_system_name);

  // Encode nodes selector
  enc.add_str(nodes_description);

  // Encode value extractors and embeddings
  enc.add_2B(value_names.size());
  for (auto&& value_name : value_names)
    enc.add_str(value_name);
  for (auto&& embedding : parser.embeddings)
    embedding.save(enc);

  // Encode the network
  network_trainer.save_network(enc);
}

} // namespace parsito
} // namespace ufal
<commit_msg>Allow oracle to return inapplicable transition.<commit_after>// This file is part of Parsito <http://github.com/ufal/parsito/>.
//
// Copyright 2015 Institute of Formal and Applied Linguistics, Faculty of
// Mathematics and Physics, Charles University in Prague, Czech Republic.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#include <algorithm>
#include <atomic>
#include <fstream>
#include <limits>
#include <random>
#include <thread>
#include <unordered_set>

#include "network/neural_network_trainer.h"
#include "parser_nn.h"
#include "parser_nn_trainer.h"
#include "utils/parse_double.h"
#include "utils/parse_int.h"
#include "utils/split.h"

namespace ufal {
namespace parsito {

void parser_nn_trainer::train(const string& transition_system_name, const string& transition_oracle_name,
                              const string& embeddings_description, const string& nodes_description, const network_parameters& parameters,
                              unsigned number_of_threads, const vector<tree>& train, const vector<tree>& heldout, binary_encoder& enc) {
  if (train.empty()) runtime_failure("No training data was given!");

  // Random generator with fixed seed for reproducibility
  mt19937 generator(42);

  // Check that all non-root nodes have heads and nonempty deprel
  for (auto&& tree : train)
    for (auto&& node : tree.nodes)
      if (node.id) {
        if (node.head < 0) runtime_failure("The node '" << node.form << "' with id " << node.id << " has no head set!");
        if (node.deprel.empty()) runtime_failure("The node '" << node.form << "' with id " << node.id << " has no deprel set!");
      }

  // Generate labels for transition system
  parser_nn parser;
  unordered_set<string> labels_set;

  for (auto&& tree : train)
    for (auto&& node : tree.nodes)
      if (node.id && !labels_set.count(node.deprel)) {
        labels_set.insert(node.deprel);
        parser.labels.push_back(node.deprel);
      }

  // Create transition system and transition oracle
  parser.system.reset(transition_system::create(transition_system_name, parser.labels));
  if (!parser.system) runtime_failure("Cannot create transition system '" << transition_system_name << "'!");

  unique_ptr<transition_oracle> oracle(parser.system->oracle(transition_oracle_name));
  if (!oracle) runtime_failure("Cannot create transition oracle '" << transition_oracle_name << "' for transition system '" << transition_system_name << "'!");

  // Create node_extractor
  string error;
  if (!parser.nodes.create(nodes_description, error)) runtime_failure(error);

  // Load value_extractors and embeddings
  vector<string> value_names;
  vector<string_piece> lines, tokens;
  split(embeddings_description, '\n', lines);
  for (auto&& line : lines) {
    // Ignore empty lines and comments
    if (!line.len || line.str[0] == '#') continue;

    split(line, ' ', tokens);
    if (!(tokens.size() >= 3 && tokens.size() <= 6))
      runtime_failure("Expected 3 to 6 columns on embedding description line '" << line << "'!");

    value_names.emplace_back(string(tokens[0].str, tokens[0].len));
    parser.values.emplace_back();
    if (!parser.values.back().create(tokens[0], error)) runtime_failure(error);

    int dimension = parse_int(tokens[1], "embedding dimension");
    int min_count = parse_int(tokens[2], "minimum frequency count");
    unsigned updatable_index = 0;
    unsigned embeddings_from_file = 0;
    string embeddings_from_file_comment;
    vector<pair<string, vector<float>>> weights;
    unordered_set<string> weights_set;

    // Compute words and counts present in the training data
    string word;
    unordered_map<string, int> word_counts;
    for (auto&& tree : train)
      for (auto&& node : tree.nodes)
        if (node.id) {
          parser.values.back().extract(node, word);
          word_counts[word]++;
        }

    // Load embedding if it was given
    if (tokens.size() >= 4) {
      int update_weights = tokens.size() >= 5 ? parse_int(tokens[4], "update weights") : 1;
      int max_embeddings = tokens.size() >= 6 ? parse_int(tokens[5], "maximum embeddings count") : numeric_limits<int>::max();
      ifstream in(string(tokens[3].str, tokens[3].len));
      if (!in.is_open()) runtime_failure("Cannot load '" << tokens[0] << "' embedding from file '" << tokens[3] << "'!");

      // Load first line containing dictionary size and dimensions
      string line;
      vector<string_piece> parts;
      if (!getline(in, line)) runtime_failure("Cannot read first line from embedding file '" << tokens[3] << "'!");
      split(line, ' ', parts);
      if (parts.size() != 2) runtime_failure("Expected two numbers on the first line of embedding file '" << tokens[3] << "'!");
      int file_dimension = parse_int(parts[1], "embedding file dimension");

      if (file_dimension < dimension) runtime_failure("The embedding file '" << tokens[3] << "' has lower dimension than required!");

      // Generate random projection when smaller dimension is required
      vector<vector<float>> projection;
      if (file_dimension > dimension) {
        embeddings_from_file_comment = "[dim" + to_string(file_dimension) + "->" + to_string(dimension) + "]";

        uniform_real_distribution<double> uniform(0, 1);
        projection.resize(dimension);
        for (auto&& row : projection) {
          row.resize(file_dimension);
          for (auto&& weight : row) weight = uniform(generator);

          double sum = 0;
          for (auto&& weight : row) sum += weight;
          for (auto&& weight : row) weight /= sum;
        }
      }

      // Load input embedding
      vector<double> input_weights(file_dimension);
      vector<float> projected_weights(dimension);
      while (getline(in, line) && int(weights.size()) < max_embeddings) {
        split(line, ' ', parts);
        if (!parts.empty() && !parts.back().len) parts.pop_back(); // Ignore space at the end of line
        if (int(parts.size()) != file_dimension + 1) runtime_failure("Wrong number of values on line '" << line << "' of embedding file '" << tokens[3]);
        for (int i = 0; i < file_dimension; i++)
          input_weights[i] = parse_double(parts[1 + i], "embedding weight");

        string word(parts[0].str, parts[0].len);

        // For update_weights == 2, ignore embeddings for unknown words
        if (update_weights == 2 && !word_counts.count(word))
          continue;

        for (int i = 0; i < dimension; i++)
          if (file_dimension == dimension) {
            projected_weights[i] = input_weights[i];
          } else {
            projected_weights[i] = 0;
            for (int j = 0; j < file_dimension; j++)
              projected_weights[i] += projection[i][j] * input_weights[j];
          }

        if (!weights_set.count(word)) {
          weights.emplace_back(word, projected_weights);
          weights_set.insert(word);
        }
      }
      embeddings_from_file = weights.size();
      updatable_index = update_weights ? 0 : embeddings_from_file;
    }

    // Add embedding for non-present word with min_count
    {
      vector<float> word_weights(dimension);
      uniform_real_distribution<float> uniform(-1, 1);
      for (auto&& word_count : word_counts)
        if (word_count.second >= min_count && !weights_set.count(word_count.first)) {
          for (auto&& word_weight : word_weights)
            word_weight = uniform(generator);

          weights.emplace_back(word_count.first, word_weights);
        }
    }

    // If there are unknown words in the training data, create initial embedding
    vector<float> unknown_weights(dimension);
    if (min_count > 1) {
      uniform_real_distribution<float> uniform(-1, 1);

      for (auto&& weight : unknown_weights)
        weight = uniform(generator);
    }

    // Add the embedding
    parser.embeddings.emplace_back();
    parser.embeddings.back().create(dimension, updatable_index, weights, unknown_weights);

    // Count the cover of this embedding
    string buffer;
    unsigned words_total = 0, words_covered = 0, words_covered_from_file = 0;
    for (auto&& tree : train)
      for (auto&& node : tree.nodes)
        if (node.id) {
          parser.values.back().extract(node, word);
          words_total++;
          int word_id = parser.embeddings.back().lookup_word(word, buffer);
          words_covered += word_id != parser.embeddings.back().unknown_word();
          words_covered_from_file += word_id != parser.embeddings.back().unknown_word() && unsigned(word_id) < embeddings_from_file;
        }

    cerr << "Initialized '" << tokens[0] << "' embedding with " << embeddings_from_file << embeddings_from_file_comment
         << "," << weights.size() << " words and " << fixed << setprecision(1) << 100. * words_covered_from_file / words_total
         << "%," << 100. * words_covered / words_total << "% coverage." << endl;
  }

  // Train the network
  unsigned total_dimension = 0, total_nodes = 0;
  for (auto&& embedding : parser.embeddings) total_dimension += embedding.dimension;
  for (auto&& tree : train) total_nodes += tree.nodes.size() - 1;
  auto scaled_parameters = parameters;
  scaled_parameters.l1_regularization /= train.size();
  scaled_parameters.l2_regularization /= total_nodes;
  neural_network_trainer network_trainer(parser.network, total_dimension * parser.nodes.node_count(), parser.system->transition_count(), scaled_parameters, generator);

  vector<int> permutation;
  for (size_t i = 0; i < train.size(); i++)
    permutation.push_back(permutation.size());

  for (int iteration = 1; network_trainer.next_iteration(); iteration++) {
    // Train on training data
    shuffle(permutation.begin(), permutation.end(), generator);

    atomic<unsigned> atomic_index(0);
    atomic<double> atomic_logprob(0);
    auto training = [&]() {
      tree t;
      configuration conf;
      string word, word_buffer;
      vector<vector<int>> nodes_embeddings;
      vector<int> extracted_nodes;
      vector<const vector<int>*> extracted_embeddings;
      neural_network_trainer::workspace workspace;
      double logprob = 0;

      for (unsigned current_index; (current_index = atomic_index++) < permutation.size();) {
        const tree& gold = train[permutation[current_index]];
        t = gold;
        t.unlink_all_nodes();
        conf.init(&t);

        // Compute embeddings
        if (t.nodes.size() > nodes_embeddings.size()) nodes_embeddings.resize(t.nodes.size());
        for (size_t i = 0; i < t.nodes.size(); i++) {
          nodes_embeddings[i].resize(parser.embeddings.size());
          for (size_t j = 0; j < parser.embeddings.size(); j++) {
            parser.values[j].extract(t.nodes[i], word);
            nodes_embeddings[i][j] = parser.embeddings[j].lookup_word(word, word_buffer);
          }
        }

        // Train the network
        while (!conf.final()) {
          // Extract nodes
          parser.nodes.extract(conf, extracted_nodes);
          extracted_embeddings.resize(extracted_nodes.size());
          for (size_t i = 0; i < extracted_nodes.size(); i++)
            extracted_embeddings[i] = extracted_nodes[i] >= 0 ? &nodes_embeddings[extracted_nodes[i]] : nullptr;

          // Propagate
          network_trainer.propagate(parser.embeddings, extracted_embeddings, workspace);

          // Find most probable applicable transition
          int network_best = -1;
          for (unsigned i = 0; i < workspace.outcomes.size(); i++)
            if (parser.system->applicable(conf, i) && (network_best < 0 || workspace.outcomes[i] > workspace.outcomes[network_best]))
              network_best = i;

          // Apply the oracle
          auto prediction = oracle->predict(conf, gold, network_best);

          // If the best transition is applicable, train on it
          if (parser.system->applicable(conf, prediction.best)) {
            // Update logprob
            if (workspace.outcomes[prediction.best])
              logprob += log(workspace.outcomes[prediction.best]);

            // Backpropagate the chosen outcome
            network_trainer.backpropagate(parser.embeddings, extracted_embeddings, prediction.best, workspace);
          }

          // Emergency break if the to_follow transition is not applicable
          if (!parser.system->applicable(conf, prediction.to_follow))
              break;

          // Follow the chosen outcome
          int child = parser.system->perform(conf, prediction.to_follow);

          // If a node was linked, recompute its embeddings as deprel has changed
          if (child >= 0)
            for (size_t i = 0; i < parser.embeddings.size(); i++) {
              parser.values[i].extract(t.nodes[child], word);
              nodes_embeddings[child][i] = parser.embeddings[i].lookup_word(word, word_buffer);
            }
        }
        network_trainer.finalize_sentence();
      }
      for (double old_atomic_logprob = atomic_logprob; atomic_logprob.compare_exchange_weak(old_atomic_logprob, old_atomic_logprob + logprob); ) {}
    };

    cerr << "Iteration " << iteration << ": ";
    if (number_of_threads > 1) {
      vector<thread> threads;
      for (unsigned i = 0; i < number_of_threads; i++) threads.emplace_back(training);
      for (; !threads.empty(); threads.pop_back()) threads.back().join();
    } else {
      training();
    }
    cerr << "training logprob " << scientific << setprecision(4) << atomic_logprob;

    // Evaluate heldout data if present
    if (!heldout.empty()) {
      network_trainer.finalize_dropout_weights();

      tree t;
      unsigned total = 0, correct_unlabelled = 0, correct_labelled = 0;
      for (auto&& gold : heldout) {
        t = gold;
        t.unlink_all_nodes();
        parser.parse(t);
        for (size_t i = 1; i < t.nodes.size(); i++) {
          total++;
          correct_unlabelled += t.nodes[i].head == gold.nodes[i].head;
          correct_labelled += t.nodes[i].head == gold.nodes[i].head && t.nodes[i].deprel == gold.nodes[i].deprel;
        }
      }

      network_trainer.finalize_dropout_weights(false);

      cerr << ", heldout UAS " << fixed << setprecision(2) << (100. * correct_unlabelled / total) << "%, LAS " << (100. * correct_labelled / total) << "%";
    }

    cerr << endl;
  }
  network_trainer.finalize_dropout_weights();

  // Encode transition system
  enc.add_2B(parser.labels.size());
  for (auto&& label : parser.labels)
    enc.add_str(label);
  enc.add_str(transition_system_name);

  // Encode nodes selector
  enc.add_str(nodes_description);

  // Encode value extractors and embeddings
  enc.add_2B(value_names.size());
  for (auto&& value_name : value_names)
    enc.add_str(value_name);
  for (auto&& embedding : parser.embeddings)
    embedding.save(enc);

  // Encode the network
  network_trainer.save_network(enc);
}

} // namespace parsito
} // namespace ufal
<|endoftext|>


<commit_before>#include "prefixedoutstream.h"

#include <string>
#include <iostream>
#include <streambuf>
#include <string.h>

using namespace mlpack::io;

PrefixedOutStream& PrefixedOutStream::operator<< (bool& val) {
  BaseLogic<bool&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (short& val) {
  BaseLogic<short&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (unsigned short& val) {
 BaseLogic<unsigned short&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (int& val) {
  BaseLogic<int&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (unsigned int& val) {
  BaseLogic<unsigned int&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (long& val) {
  BaseLogic<long&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (unsigned long& val) {
  BaseLogic<unsigned long&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (float& val) {
  BaseLogic<float&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (double& val) {
  BaseLogic<double&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (long double& val) {
  BaseLogic<long double&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (void* val) {
  BaseLogic<void*>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (const char* str) {
  BaseLogic<const char*>(str);

  // No need to worry about platform-independent newline characters; because
  // stdout and stderr are text-mode streams, they do the relevant conversions
  // automatically.
  if (strstr(str, "\n") != NULL) {
    carriageReturned = true;
    if (fatal) // Terminate application.
      exit(1);
  }

  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (std::string& str) {

  BaseLogic<std::string&>(str);

  // No need to worry about platform-independent newline characters; because
  // stdout and stderr are text-mode streams, they do the relevant conversions
  // automatically.
  if (str.find("\n") != std::string::npos) {
    carriageReturned = true;
    if (fatal) // Terminate application.
      exit(1);
  }

  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (std::streambuf* sb) {
  destination << sb;
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<<
    (std::ostream& (*pf) (std::ostream&)) {
  // We don't want to prefix on what will show up as empty lines.
  destination << pf;
  carriageReturned = true;
  if (fatal) // Terminate application.
    exit(1);

  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (std::ios& (*pf) (std::ios&)) {

  BaseLogic<std::ios& (*) (std::ios&)>(pf);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<<
    (std::ios_base& (*pf) (std::ios_base&)) {

  BaseLogic<std::ios_base& (*) (std::ios_base&)>(pf);
  return *this;
}

<commit_msg>Oops, this isn't included by default.<commit_after>#include "prefixedoutstream.h"

#include <string>
#include <iostream>
#include <streambuf>
#include <string.h>
#include <stdlib.h>

using namespace mlpack::io;

PrefixedOutStream& PrefixedOutStream::operator<< (bool& val) {
  BaseLogic<bool&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (short& val) {
  BaseLogic<short&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (unsigned short& val) {
 BaseLogic<unsigned short&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (int& val) {
  BaseLogic<int&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (unsigned int& val) {
  BaseLogic<unsigned int&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (long& val) {
  BaseLogic<long&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (unsigned long& val) {
  BaseLogic<unsigned long&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (float& val) {
  BaseLogic<float&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (double& val) {
  BaseLogic<double&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (long double& val) {
  BaseLogic<long double&>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (void* val) {
  BaseLogic<void*>(val);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (const char* str) {
  BaseLogic<const char*>(str);

  // No need to worry about platform-independent newline characters; because
  // stdout and stderr are text-mode streams, they do the relevant conversions
  // automatically.
  if (strstr(str, "\n") != NULL) {
    carriageReturned = true;
    if (fatal) // Terminate application.
      exit(1);
  }

  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (std::string& str) {

  BaseLogic<std::string&>(str);

  // No need to worry about platform-independent newline characters; because
  // stdout and stderr are text-mode streams, they do the relevant conversions
  // automatically.
  if (str.find("\n") != std::string::npos) {
    carriageReturned = true;
    if (fatal) // Terminate application.
      exit(1);
  }

  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (std::streambuf* sb) {
  destination << sb;
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<<
    (std::ostream& (*pf) (std::ostream&)) {
  // We don't want to prefix on what will show up as empty lines.
  destination << pf;
  carriageReturned = true;
  if (fatal) // Terminate application.
    exit(1);

  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<< (std::ios& (*pf) (std::ios&)) {

  BaseLogic<std::ios& (*) (std::ios&)>(pf);
  return *this;
}

PrefixedOutStream& PrefixedOutStream::operator<<
    (std::ios_base& (*pf) (std::ios_base&)) {

  BaseLogic<std::ios_base& (*) (std::ios_base&)>(pf);
  return *this;
}

<|endoftext|>


<commit_before>#include "continuous_fmm.h"

int main(int argc, char *argv[]) {

  // Initialize FastExec (parameter handling stuff).
  fx_init(argc, argv, NULL);

  ////////// READING PARAMETERS AND LOADING DATA /////////////////////

  // FASTexec organizes parameters and results into submodules.  Think
  // of this as creating a new folder named "cfmm_module" under the
  // root directory (NULL) for the continuous FMM object to work
  // inside.  Here, we initialize it with all parameters defined
  // "--cfmm/...=...".
  struct datanode* cfmm_module = fx_submodule(fx_root, "cfmm");

  // The reference data file is a required parameter.
  const char* references_file_name = fx_param_str_req(fx_root, "data");

  // The query data file defaults to the references.
  const char* queries_file_name =
    fx_param_str(fx_root, "query", references_file_name);

  // Query and reference datasets, reference weight dataset, and the
  // reference bandwidth dataset.
  Matrix references;
  Matrix reference_weights;
  Matrix reference_bandwidths;
  Matrix queries;

  // Flag for telling whether references are equal to queries
  bool queries_equal_references =
    !strcmp(queries_file_name, references_file_name);

  // data::Load inits a matrix with the contents of a .csv or .arff.
  data::Load(references_file_name, &references);
  la::Scale(100.0, &references);
  if(queries_equal_references) {
    queries.Alias(references);
  }
  else {
    data::Load(queries_file_name, &queries);
  }

  // If the reference weight file name is specified, then read in,
  // otherwise, initialize to uniform weights.
  if(fx_param_exists(fx_root, "dwgts")) {
    data::Load(fx_param_str(fx_root, "dwgts", NULL), &reference_weights);
  }
  else {
    reference_weights.Init(1, references.n_cols());
    reference_weights.SetAll(1);
  }

  // Read in the bandwidth for each point.
  if(fx_param_exists(fx_root, "bandwidths")) {
    data::Load(fx_param_str(fx_root, "bandwidths", NULL),
	       &reference_bandwidths);
  }
  else {
    reference_bandwidths.Init(1, references.n_cols());
    reference_bandwidths.SetAll(1);
  }

  // Declare the fast multipole method object and initialize it.
  ContinuousFmm cfmm_algorithm;
  cfmm_algorithm.Init(queries, references, reference_weights,
		      reference_bandwidths, queries_equal_references,
		      cfmm_module);

  // Start computation.
  cfmm_algorithm.Compute();

  // Compute naively.
  Vector naive_results;
  cfmm_algorithm.NaiveCompute(&naive_results);

  fx_done(NULL);
  return 0;
}
<commit_msg>Final commit of CFMM.<commit_after>#include "continuous_fmm.h"

int main(int argc, char *argv[]) {

  // Initialize FastExec (parameter handling stuff).
  fx_init(argc, argv, NULL);

  ////////// READING PARAMETERS AND LOADING DATA /////////////////////

  // FASTexec organizes parameters and results into submodules.  Think
  // of this as creating a new folder named "cfmm_module" under the
  // root directory (NULL) for the continuous FMM object to work
  // inside.  Here, we initialize it with all parameters defined
  // "--cfmm/...=...".
  struct datanode* cfmm_module = fx_submodule(fx_root, "cfmm");

  // The reference data file is a required parameter.
  const char* references_file_name = fx_param_str_req(fx_root, "data");

  // The query data file defaults to the references.
  const char* queries_file_name =
    fx_param_str(fx_root, "query", references_file_name);

  // Query and reference datasets, reference weight dataset, and the
  // reference bandwidth dataset.
  Matrix references;
  Matrix reference_weights;
  Matrix reference_bandwidths;
  Matrix queries;

  // Flag for telling whether references are equal to queries
  bool queries_equal_references =
    !strcmp(queries_file_name, references_file_name);

  // data::Load inits a matrix with the contents of a .csv or .arff.
  data::Load(references_file_name, &references);
  if(queries_equal_references) {
    queries.Alias(references);
  }
  else {
    data::Load(queries_file_name, &queries);
  }

  // If the reference weight file name is specified, then read in,
  // otherwise, initialize to uniform weights.
  if(fx_param_exists(fx_root, "dwgts")) {
    data::Load(fx_param_str(fx_root, "dwgts", NULL), &reference_weights);
  }
  else {
    reference_weights.Init(1, references.n_cols());
    reference_weights.SetAll(1);
  }

  // Read in the bandwidth for each point.
  if(fx_param_exists(fx_root, "bandwidths")) {
    data::Load(fx_param_str(fx_root, "bandwidths", NULL),
	       &reference_bandwidths);
  }
  else {
    reference_bandwidths.Init(1, references.n_cols());
    reference_bandwidths.SetAll(1);
  }

  // Declare the fast multipole method object and initialize it.
  ContinuousFmm cfmm_algorithm;
  cfmm_algorithm.Init(queries, references, reference_weights,
		      reference_bandwidths, queries_equal_references,
		      cfmm_module);

  // Start computation.
  cfmm_algorithm.Compute();

  // Compute naively.
  Vector naive_results;
  cfmm_algorithm.NaiveCompute(&naive_results);

  fx_done(NULL);
  return 0;
}
<|endoftext|>


<commit_before>#include <cvd/config.h>
#include <cvd/image_convert.h>
#include <cvd/utility.h>
#include <arm_neon.h>

namespace CVD {

	namespace Internal {
		void convert_rgb_to_grey_NEON(const SubImage<Rgb<byte> >& from, SubImage<byte>& to, const unsigned r, const unsigned g, const unsigned b) {
			const uint8x8_t red = vdup_n_u8(r);
			const uint8x8_t green = vdup_n_u8(g);
			const uint8x8_t blue = vdup_n_u8(b);

			for( int y = 0; y < from.size().y; ++y){
				const Rgb<byte> * in = from[y];
				byte * out = to[y];
				for( int x = 0; x < from.size().x; x+=8, in += 8, out += 8){
					uint8x8x3_t in_data  = vld3_u8((const uint8_t *)in);

					uint16x8_t sum = vmull_u8(in_data.val[0], red);
					sum = vmlal_u8(sum, in_data.val[1], green);
					sum = vmlal_u8(sum, in_data.val[2], blue);

					uint8x8_t final_sum = vshrn_n_u16(sum, 8); // divide by 256
					vst1_u8(out, final_sum);
				}
			}
		}
	}

	void ConvertImage<Rgb<byte>, byte, Pixel::CIE<Rgb<byte>, byte>, 1>::convert(const SubImage<Rgb<byte> >& from, SubImage<byte>& to)
	{

		if((from.size().x % 8) == 0){
            // red (77) + green (150) + blue (29) = 256
            Internal::convert_rgb_to_grey_NEON( from, to, 77, 150, 29);
		} else {
			const int multiple_of_8_width = (from.size().x / 8) * 8;
			const ImageRef end_fast(multiple_of_8_width, from.size().y);

            // red (77) + green (150) + blue (29) = 256
			Internal::convert_rgb_to_grey_NEON( from.sub_image(ImageRef_zero, end_fast), to.sub_image(ImageRef_zero, end_fast).ref(), 77, 150, 29);

			for(int y = 0; y < from.size().y; ++y){
				const Rgb<byte> * in = from[y]+multiple_of_8_width;
				byte * out = to[y]+multiple_of_8_width;
				for(int x = multiple_of_8_width; x < from.size().x; ++x, ++in, ++out){
					Pixel::CIE<Rgb<byte>,byte>::convert(*in, *out);
				}
			}
		}
	}
}
<commit_msg>cleanup<commit_after>#include <cvd/image_convert.h>
#include <arm_neon.h>

namespace CVD {

	namespace Internal {
		void convert_rgb_to_grey_NEON(const SubImage<Rgb<byte> >& from, SubImage<byte>& to, const unsigned r, const unsigned g, const unsigned b) {
			const uint8x8_t red = vdup_n_u8(r);
			const uint8x8_t green = vdup_n_u8(g);
			const uint8x8_t blue = vdup_n_u8(b);

			for( int y = 0; y < from.size().y; ++y){
				const Rgb<byte> * in = from[y];
				byte * out = to[y];
				for( int x = 0; x < from.size().x; x+=8, in += 8, out += 8){
					uint8x8x3_t in_data  = vld3_u8((const uint8_t *)in);

					uint16x8_t sum = vmull_u8(in_data.val[0], red);
					sum = vmlal_u8(sum, in_data.val[1], green);
					sum = vmlal_u8(sum, in_data.val[2], blue);

					uint8x8_t final_sum = vshrn_n_u16(sum, 8); // divide by 256
					vst1_u8(out, final_sum);
				}
			}
		}
	}

	void ConvertImage<Rgb<byte>, byte, Pixel::CIE<Rgb<byte>, byte>, 1>::convert(const SubImage<Rgb<byte> >& from, SubImage<byte>& to)
	{

		if((from.size().x % 8) == 0){
            // red (77) + green (150) + blue (29) = 256
            Internal::convert_rgb_to_grey_NEON( from, to, 77, 150, 29);
		} else {
			const int multiple_of_8_width = (from.size().x / 8) * 8;
			const ImageRef end_fast(multiple_of_8_width, from.size().y);

            // red (77) + green (150) + blue (29) = 256
			Internal::convert_rgb_to_grey_NEON( from.sub_image(ImageRef_zero, end_fast), to.sub_image(ImageRef_zero, end_fast).ref(), 77, 150, 29);

			for(int y = 0; y < from.size().y; ++y){
				const Rgb<byte> * in = from[y]+multiple_of_8_width;
				byte * out = to[y]+multiple_of_8_width;
				for(int x = multiple_of_8_width; x < from.size().x; ++x, ++in, ++out){
					Pixel::CIE<Rgb<byte>,byte>::convert(*in, *out);
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
 * #%L
 * %%
 * Copyright (C) 2011 - 2017 BMW Car IT GmbH
 * %%
 * 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.
 * #L%
 */

#include "joynr/MulticastReceiverDirectory.h"

namespace joynr
{

INIT_LOGGER(MulticastReceiverDirectory);

MulticastReceiverDirectory::~MulticastReceiverDirectory()
{
    JOYNR_LOG_TRACE(logger, "destructor: number of entries = {}", multicastReceivers.size());
    multicastReceivers.clear();
}

void MulticastReceiverDirectory::registerMulticastReceiver(const std::string& multicastId,
                                                           const std::string& receiverId)
{
    JOYNR_LOG_TRACE(logger,
                    "register multicast receiver: multicastId={}, receiverId={}",
                    multicastId,
                    receiverId);
    std::lock_guard<std::recursive_mutex> lock(mutex);
    MulticastMatcher matcher(multicastId);
    multicastReceivers[matcher].insert(receiverId);
}

bool MulticastReceiverDirectory::unregisterMulticastReceiver(const std::string& multicastId,
                                                             const std::string& receiverId)
{
    JOYNR_LOG_TRACE(logger,
                    "unregister multicast receiver: multicastId={}, receiverId={}",
                    multicastId,
                    receiverId);
    std::lock_guard<std::recursive_mutex> lock(mutex);
    if (contains(multicastId)) {
        MulticastMatcher matcher(multicastId);
        std::unordered_set<std::string>& receivers = multicastReceivers[matcher];
        receivers.erase(receiverId);
        JOYNR_LOG_TRACE(logger,
                        "removed multicast receiver: multicastId={}, receiverId={}",
                        multicastId,
                        receiverId);
        if (receivers.empty()) {
            JOYNR_LOG_TRACE(logger, "removed last multicast receiver: multicastId={}", multicastId);
            multicastReceivers.erase(matcher);
        }
        return true;
    }
    return false;
}

std::unordered_set<std::string> MulticastReceiverDirectory::getReceivers(
        const std::string& multicastId)
{
    JOYNR_LOG_TRACE(logger, "get multicast receivers: multicastId={}", multicastId);
    std::lock_guard<std::recursive_mutex> lock(mutex);

    std::unordered_set<std::string> foundReceivers;
    for (const std::pair<MulticastMatcher, std::unordered_set<std::string>>& pair :
         multicastReceivers) {
        if (pair.first.doesMatch(multicastId)) {
            foundReceivers.insert(pair.second.cbegin(), pair.second.cend());
        }
    }
    return foundReceivers;
}

std::vector<std::string> MulticastReceiverDirectory::getMulticastIds() const
{
    std::vector<std::string> multicastIds;

    for (const auto& multicastReceiver : multicastReceivers) {
        multicastIds.push_back(multicastReceiver.first.multicastId);
    }

    return multicastIds;
}

bool MulticastReceiverDirectory::contains(const std::string& multicastId)
{
    std::lock_guard<std::recursive_mutex> lock(mutex);
    MulticastMatcher matcher(multicastId);
    return multicastReceivers.find(matcher) != multicastReceivers.cend();
}

bool MulticastReceiverDirectory::contains(const std::string& multicastId,
                                          const std::string& receiverId)
{
    std::lock_guard<std::recursive_mutex> lock(mutex);
    auto receivers = getReceivers(multicastId);
    return receivers.find(receiverId) != receivers.cend();
}

} // namespace joynr
<commit_msg>[C++] avoid implicit cast<commit_after>/*
 * #%L
 * %%
 * Copyright (C) 2011 - 2017 BMW Car IT GmbH
 * %%
 * 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.
 * #L%
 */

#include "joynr/MulticastReceiverDirectory.h"

namespace joynr
{

INIT_LOGGER(MulticastReceiverDirectory);

MulticastReceiverDirectory::~MulticastReceiverDirectory()
{
    JOYNR_LOG_TRACE(logger, "destructor: number of entries = {}", multicastReceivers.size());
    multicastReceivers.clear();
}

void MulticastReceiverDirectory::registerMulticastReceiver(const std::string& multicastId,
                                                           const std::string& receiverId)
{
    JOYNR_LOG_TRACE(logger,
                    "register multicast receiver: multicastId={}, receiverId={}",
                    multicastId,
                    receiverId);
    std::lock_guard<std::recursive_mutex> lock(mutex);
    MulticastMatcher matcher(multicastId);
    multicastReceivers[matcher].insert(receiverId);
}

bool MulticastReceiverDirectory::unregisterMulticastReceiver(const std::string& multicastId,
                                                             const std::string& receiverId)
{
    JOYNR_LOG_TRACE(logger,
                    "unregister multicast receiver: multicastId={}, receiverId={}",
                    multicastId,
                    receiverId);
    std::lock_guard<std::recursive_mutex> lock(mutex);
    if (contains(multicastId)) {
        MulticastMatcher matcher(multicastId);
        std::unordered_set<std::string>& receivers = multicastReceivers[matcher];
        receivers.erase(receiverId);
        JOYNR_LOG_TRACE(logger,
                        "removed multicast receiver: multicastId={}, receiverId={}",
                        multicastId,
                        receiverId);
        if (receivers.empty()) {
            JOYNR_LOG_TRACE(logger, "removed last multicast receiver: multicastId={}", multicastId);
            multicastReceivers.erase(matcher);
        }
        return true;
    }
    return false;
}

std::unordered_set<std::string> MulticastReceiverDirectory::getReceivers(
        const std::string& multicastId)
{
    JOYNR_LOG_TRACE(logger, "get multicast receivers: multicastId={}", multicastId);
    std::lock_guard<std::recursive_mutex> lock(mutex);

    std::unordered_set<std::string> foundReceivers;
    for (const auto& entry : multicastReceivers) {
        if (entry.first.doesMatch(multicastId)) {
            foundReceivers.insert(entry.second.cbegin(), entry.second.cend());
        }
    }
    return foundReceivers;
}

std::vector<std::string> MulticastReceiverDirectory::getMulticastIds() const
{
    std::vector<std::string> multicastIds;

    for (const auto& multicastReceiver : multicastReceivers) {
        multicastIds.push_back(multicastReceiver.first.multicastId);
    }

    return multicastIds;
}

bool MulticastReceiverDirectory::contains(const std::string& multicastId)
{
    std::lock_guard<std::recursive_mutex> lock(mutex);
    MulticastMatcher matcher(multicastId);
    return multicastReceivers.find(matcher) != multicastReceivers.cend();
}

bool MulticastReceiverDirectory::contains(const std::string& multicastId,
                                          const std::string& receiverId)
{
    std::lock_guard<std::recursive_mutex> lock(mutex);
    auto receivers = getReceivers(multicastId);
    return receivers.find(receiverId) != receivers.cend();
}

} // namespace joynr
<|endoftext|>


<commit_before>//===-- SIMachineFunctionInfo.cpp - SI Machine Function Info -------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
/// \file
//===----------------------------------------------------------------------===//

#include "SIMachineFunctionInfo.h"
#include "AMDGPUSubtarget.h"
#include "SIInstrInfo.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/LLVMContext.h"

#define MAX_LANES 64

using namespace llvm;

// Pin the vtable to this file.
void SIMachineFunctionInfo::anchor() {}

SIMachineFunctionInfo::SIMachineFunctionInfo(const MachineFunction &MF)
  : AMDGPUMachineFunction(MF),
    TIDReg(AMDGPU::NoRegister),
    ScratchRSrcReg(AMDGPU::NoRegister),
    ScratchWaveOffsetReg(AMDGPU::NoRegister),
    PrivateSegmentBufferUserSGPR(AMDGPU::NoRegister),
    DispatchPtrUserSGPR(AMDGPU::NoRegister),
    QueuePtrUserSGPR(AMDGPU::NoRegister),
    KernargSegmentPtrUserSGPR(AMDGPU::NoRegister),
    DispatchIDUserSGPR(AMDGPU::NoRegister),
    FlatScratchInitUserSGPR(AMDGPU::NoRegister),
    PrivateSegmentSizeUserSGPR(AMDGPU::NoRegister),
    GridWorkGroupCountXUserSGPR(AMDGPU::NoRegister),
    GridWorkGroupCountYUserSGPR(AMDGPU::NoRegister),
    GridWorkGroupCountZUserSGPR(AMDGPU::NoRegister),
    WorkGroupIDXSystemSGPR(AMDGPU::NoRegister),
    WorkGroupIDYSystemSGPR(AMDGPU::NoRegister),
    WorkGroupIDZSystemSGPR(AMDGPU::NoRegister),
    WorkGroupInfoSystemSGPR(AMDGPU::NoRegister),
    PrivateSegmentWaveByteOffsetSystemSGPR(AMDGPU::NoRegister),
    PSInputAddr(0),
    ReturnsVoid(true),
    MaximumWorkGroupSize(0),
    DebuggerReservedVGPRCount(0),
    LDSWaveSpillSize(0),
    PSInputEna(0),
    NumUserSGPRs(0),
    NumSystemSGPRs(0),
    HasSpilledSGPRs(false),
    HasSpilledVGPRs(false),
    HasNonSpillStackObjects(false),
    HasFlatInstructions(false),
    PrivateSegmentBuffer(false),
    DispatchPtr(false),
    QueuePtr(false),
    DispatchID(false),
    KernargSegmentPtr(false),
    FlatScratchInit(false),
    GridWorkgroupCountX(false),
    GridWorkgroupCountY(false),
    GridWorkgroupCountZ(false),
    WorkGroupIDX(false),
    WorkGroupIDY(false),
    WorkGroupIDZ(false),
    WorkGroupInfo(false),
    PrivateSegmentWaveByteOffset(false),
    WorkItemIDX(false),
    WorkItemIDY(false),
    WorkItemIDZ(false) {
  const AMDGPUSubtarget &ST = MF.getSubtarget<AMDGPUSubtarget>();
  const Function *F = MF.getFunction();

  PSInputAddr = AMDGPU::getInitialPSInputAddr(*F);

  const MachineFrameInfo *FrameInfo = MF.getFrameInfo();

  if (!AMDGPU::isShader(F->getCallingConv())) {
    KernargSegmentPtr = true;
    WorkGroupIDX = true;
    WorkItemIDX = true;
  }

  if (F->hasFnAttribute("amdgpu-work-group-id-y"))
    WorkGroupIDY = true;

  if (F->hasFnAttribute("amdgpu-work-group-id-z"))
    WorkGroupIDZ = true;

  if (F->hasFnAttribute("amdgpu-work-item-id-y"))
    WorkItemIDY = true;

  if (F->hasFnAttribute("amdgpu-work-item-id-z"))
    WorkItemIDZ = true;

  // X, XY, and XYZ are the only supported combinations, so make sure Y is
  // enabled if Z is.
  if (WorkItemIDZ)
    WorkItemIDY = true;

  bool MaySpill = ST.isVGPRSpillingEnabled(*F);
  bool HasStackObjects = FrameInfo->hasStackObjects();

  if (HasStackObjects || MaySpill)
    PrivateSegmentWaveByteOffset = true;

  if (ST.isAmdHsaOS()) {
    if (HasStackObjects || MaySpill)
      PrivateSegmentBuffer = true;

    if (F->hasFnAttribute("amdgpu-dispatch-ptr"))
      DispatchPtr = true;

    if (F->hasFnAttribute("amdgpu-queue-ptr"))
      QueuePtr = true;
  }

  // We don't need to worry about accessing spills with flat instructions.
  // TODO: On VI where we must use flat for global, we should be able to omit
  // this if it is never used for generic access.
  if (HasStackObjects && ST.getGeneration() >= AMDGPUSubtarget::SEA_ISLANDS &&
      ST.isAmdHsaOS())
    FlatScratchInit = true;

  if (AMDGPU::isCompute(F->getCallingConv()))
    MaximumWorkGroupSize = AMDGPU::getMaximumWorkGroupSize(*F);
  else
    MaximumWorkGroupSize = ST.getWavefrontSize();

  if (ST.debuggerReserveRegs())
    DebuggerReservedVGPRCount = 4;
}

unsigned SIMachineFunctionInfo::addPrivateSegmentBuffer(
  const SIRegisterInfo &TRI) {
  PrivateSegmentBufferUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_128RegClass);
  NumUserSGPRs += 4;
  return PrivateSegmentBufferUserSGPR;
}

unsigned SIMachineFunctionInfo::addDispatchPtr(const SIRegisterInfo &TRI) {
  DispatchPtrUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return DispatchPtrUserSGPR;
}

unsigned SIMachineFunctionInfo::addQueuePtr(const SIRegisterInfo &TRI) {
  QueuePtrUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return QueuePtrUserSGPR;
}

unsigned SIMachineFunctionInfo::addKernargSegmentPtr(const SIRegisterInfo &TRI) {
  KernargSegmentPtrUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return KernargSegmentPtrUserSGPR;
}

unsigned SIMachineFunctionInfo::addFlatScratchInit(const SIRegisterInfo &TRI) {
  FlatScratchInitUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return FlatScratchInitUserSGPR;
}

SIMachineFunctionInfo::SpilledReg SIMachineFunctionInfo::getSpilledReg(
                                                       MachineFunction *MF,
                                                       unsigned FrameIndex,
                                                       unsigned SubIdx) {
  MachineFrameInfo *FrameInfo = MF->getFrameInfo();
  const SIRegisterInfo *TRI = static_cast<const SIRegisterInfo *>(
      MF->getSubtarget<AMDGPUSubtarget>().getRegisterInfo());
  MachineRegisterInfo &MRI = MF->getRegInfo();
  int64_t Offset = FrameInfo->getObjectOffset(FrameIndex);
  Offset += SubIdx * 4;

  unsigned LaneVGPRIdx = Offset / (64 * 4);
  unsigned Lane = (Offset / 4) % 64;

  struct SpilledReg Spill;
  Spill.Lane = Lane;

  if (!LaneVGPRs.count(LaneVGPRIdx)) {
    unsigned LaneVGPR = TRI->findUnusedRegister(MRI, &AMDGPU::VGPR_32RegClass);

    if (LaneVGPR == AMDGPU::NoRegister)
      // We have no VGPRs left for spilling SGPRs.
      return Spill;

    LaneVGPRs[LaneVGPRIdx] = LaneVGPR;

    // Add this register as live-in to all blocks to avoid machine verifer
    // complaining about use of an undefined physical register.
    for (MachineFunction::iterator BI = MF->begin(), BE = MF->end();
         BI != BE; ++BI) {
      BI->addLiveIn(LaneVGPR);
    }
  }

  Spill.VGPR = LaneVGPRs[LaneVGPRIdx];
  return Spill;
}

unsigned SIMachineFunctionInfo::getMaximumWorkGroupSize(
                                              const MachineFunction &MF) const {
  return MaximumWorkGroupSize;
}
<commit_msg>AMDGPU: Add option to disable spilling SGPRs to VGPRs.<commit_after>//===-- SIMachineFunctionInfo.cpp -------- SI Machine Function Info -------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "SIMachineFunctionInfo.h"
#include "AMDGPUSubtarget.h"
#include "SIInstrInfo.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/LLVMContext.h"

#define MAX_LANES 64

using namespace llvm;

static cl::opt<bool> EnableSpillSGPRToVGPR(
  "amdgpu-spill-sgpr-to-vgpr",
  cl::desc("Enable spilling VGPRs to SGPRs"),
  cl::ReallyHidden,
  cl::init(true));

// Pin the vtable to this file.
void SIMachineFunctionInfo::anchor() {}

SIMachineFunctionInfo::SIMachineFunctionInfo(const MachineFunction &MF)
  : AMDGPUMachineFunction(MF),
    TIDReg(AMDGPU::NoRegister),
    ScratchRSrcReg(AMDGPU::NoRegister),
    ScratchWaveOffsetReg(AMDGPU::NoRegister),
    PrivateSegmentBufferUserSGPR(AMDGPU::NoRegister),
    DispatchPtrUserSGPR(AMDGPU::NoRegister),
    QueuePtrUserSGPR(AMDGPU::NoRegister),
    KernargSegmentPtrUserSGPR(AMDGPU::NoRegister),
    DispatchIDUserSGPR(AMDGPU::NoRegister),
    FlatScratchInitUserSGPR(AMDGPU::NoRegister),
    PrivateSegmentSizeUserSGPR(AMDGPU::NoRegister),
    GridWorkGroupCountXUserSGPR(AMDGPU::NoRegister),
    GridWorkGroupCountYUserSGPR(AMDGPU::NoRegister),
    GridWorkGroupCountZUserSGPR(AMDGPU::NoRegister),
    WorkGroupIDXSystemSGPR(AMDGPU::NoRegister),
    WorkGroupIDYSystemSGPR(AMDGPU::NoRegister),
    WorkGroupIDZSystemSGPR(AMDGPU::NoRegister),
    WorkGroupInfoSystemSGPR(AMDGPU::NoRegister),
    PrivateSegmentWaveByteOffsetSystemSGPR(AMDGPU::NoRegister),
    PSInputAddr(0),
    ReturnsVoid(true),
    MaximumWorkGroupSize(0),
    DebuggerReservedVGPRCount(0),
    LDSWaveSpillSize(0),
    PSInputEna(0),
    NumUserSGPRs(0),
    NumSystemSGPRs(0),
    HasSpilledSGPRs(false),
    HasSpilledVGPRs(false),
    HasNonSpillStackObjects(false),
    HasFlatInstructions(false),
    PrivateSegmentBuffer(false),
    DispatchPtr(false),
    QueuePtr(false),
    DispatchID(false),
    KernargSegmentPtr(false),
    FlatScratchInit(false),
    GridWorkgroupCountX(false),
    GridWorkgroupCountY(false),
    GridWorkgroupCountZ(false),
    WorkGroupIDX(false),
    WorkGroupIDY(false),
    WorkGroupIDZ(false),
    WorkGroupInfo(false),
    PrivateSegmentWaveByteOffset(false),
    WorkItemIDX(false),
    WorkItemIDY(false),
    WorkItemIDZ(false) {
  const AMDGPUSubtarget &ST = MF.getSubtarget<AMDGPUSubtarget>();
  const Function *F = MF.getFunction();

  PSInputAddr = AMDGPU::getInitialPSInputAddr(*F);

  const MachineFrameInfo *FrameInfo = MF.getFrameInfo();

  if (!AMDGPU::isShader(F->getCallingConv())) {
    KernargSegmentPtr = true;
    WorkGroupIDX = true;
    WorkItemIDX = true;
  }

  if (F->hasFnAttribute("amdgpu-work-group-id-y"))
    WorkGroupIDY = true;

  if (F->hasFnAttribute("amdgpu-work-group-id-z"))
    WorkGroupIDZ = true;

  if (F->hasFnAttribute("amdgpu-work-item-id-y"))
    WorkItemIDY = true;

  if (F->hasFnAttribute("amdgpu-work-item-id-z"))
    WorkItemIDZ = true;

  // X, XY, and XYZ are the only supported combinations, so make sure Y is
  // enabled if Z is.
  if (WorkItemIDZ)
    WorkItemIDY = true;

  bool MaySpill = ST.isVGPRSpillingEnabled(*F);
  bool HasStackObjects = FrameInfo->hasStackObjects();

  if (HasStackObjects || MaySpill)
    PrivateSegmentWaveByteOffset = true;

  if (ST.isAmdHsaOS()) {
    if (HasStackObjects || MaySpill)
      PrivateSegmentBuffer = true;

    if (F->hasFnAttribute("amdgpu-dispatch-ptr"))
      DispatchPtr = true;

    if (F->hasFnAttribute("amdgpu-queue-ptr"))
      QueuePtr = true;
  }

  // We don't need to worry about accessing spills with flat instructions.
  // TODO: On VI where we must use flat for global, we should be able to omit
  // this if it is never used for generic access.
  if (HasStackObjects && ST.getGeneration() >= AMDGPUSubtarget::SEA_ISLANDS &&
      ST.isAmdHsaOS())
    FlatScratchInit = true;

  if (AMDGPU::isCompute(F->getCallingConv()))
    MaximumWorkGroupSize = AMDGPU::getMaximumWorkGroupSize(*F);
  else
    MaximumWorkGroupSize = ST.getWavefrontSize();

  if (ST.debuggerReserveRegs())
    DebuggerReservedVGPRCount = 4;
}

unsigned SIMachineFunctionInfo::addPrivateSegmentBuffer(
  const SIRegisterInfo &TRI) {
  PrivateSegmentBufferUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_128RegClass);
  NumUserSGPRs += 4;
  return PrivateSegmentBufferUserSGPR;
}

unsigned SIMachineFunctionInfo::addDispatchPtr(const SIRegisterInfo &TRI) {
  DispatchPtrUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return DispatchPtrUserSGPR;
}

unsigned SIMachineFunctionInfo::addQueuePtr(const SIRegisterInfo &TRI) {
  QueuePtrUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return QueuePtrUserSGPR;
}

unsigned SIMachineFunctionInfo::addKernargSegmentPtr(const SIRegisterInfo &TRI) {
  KernargSegmentPtrUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return KernargSegmentPtrUserSGPR;
}

unsigned SIMachineFunctionInfo::addFlatScratchInit(const SIRegisterInfo &TRI) {
  FlatScratchInitUserSGPR = TRI.getMatchingSuperReg(
    getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass);
  NumUserSGPRs += 2;
  return FlatScratchInitUserSGPR;
}

SIMachineFunctionInfo::SpilledReg SIMachineFunctionInfo::getSpilledReg(
                                                       MachineFunction *MF,
                                                       unsigned FrameIndex,
                                                       unsigned SubIdx) {
  if (!EnableSpillSGPRToVGPR)
    return SpilledReg();

  MachineFrameInfo *FrameInfo = MF->getFrameInfo();
  const SIRegisterInfo *TRI = static_cast<const SIRegisterInfo *>(
      MF->getSubtarget<AMDGPUSubtarget>().getRegisterInfo());
  MachineRegisterInfo &MRI = MF->getRegInfo();
  int64_t Offset = FrameInfo->getObjectOffset(FrameIndex);
  Offset += SubIdx * 4;

  unsigned LaneVGPRIdx = Offset / (64 * 4);
  unsigned Lane = (Offset / 4) % 64;

  struct SpilledReg Spill;
  Spill.Lane = Lane;

  if (!LaneVGPRs.count(LaneVGPRIdx)) {
    unsigned LaneVGPR = TRI->findUnusedRegister(MRI, &AMDGPU::VGPR_32RegClass);

    if (LaneVGPR == AMDGPU::NoRegister)
      // We have no VGPRs left for spilling SGPRs.
      return Spill;

    LaneVGPRs[LaneVGPRIdx] = LaneVGPR;

    // Add this register as live-in to all blocks to avoid machine verifer
    // complaining about use of an undefined physical register.
    for (MachineFunction::iterator BI = MF->begin(), BE = MF->end();
         BI != BE; ++BI) {
      BI->addLiveIn(LaneVGPR);
    }
  }

  Spill.VGPR = LaneVGPRs[LaneVGPRIdx];
  return Spill;
}

unsigned SIMachineFunctionInfo::getMaximumWorkGroupSize(
                                              const MachineFunction &MF) const {
  return MaximumWorkGroupSize;
}
<|endoftext|>