<commit_before>// Copyright 2014 Alessio Sclocco <a.sclocco@vu.nl>
//
// 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 <iostream>
#include <string>
#include <vector>
#include <exception>
#include <fstream>
#include <iomanip>
#include <limits>

#include <ArgumentList.hpp>
#include <Observation.hpp>
#include <InitializeOpenCL.hpp>
#include <Kernel.hpp>
#include <Shifts.hpp>
#include <Dedispersion.hpp>
#include <utils.hpp>
#include <Timer.hpp>
#include <Stats.hpp>

// Define the data types
typedef float inputDataType;
std::string inputTypeName("float");
std::string intermediateTypeName("float");
typedef float outputDataType;
std::string outputTypeName("float");

void initializeDeviceMemory(cl::Context & clContext, cl::CommandQueue * clQueue, std::vector< float > * shifts, cl::Buffer * shifts_d, cl::Buffer * dispersedData_d, const unsigned int dispersedData_size, cl::Buffer * dedispersedData_d, const unsigned int dedispersedData_size);

int main(int argc, char * argv[]) {
  bool reInit = false;
  uint8_t inputBits = 0;
  unsigned int splitSeconds = 0;
	unsigned int nrIterations = 0;
	unsigned int clPlatformID = 0;
	unsigned int clDeviceID = 0;
	unsigned int minThreads = 0;
  unsigned int maxThreads = 0;
	unsigned int maxRows = 0;
	unsigned int maxColumns = 0;
  unsigned int threadUnit = 0;
  unsigned int threadIncrement = 0;
  unsigned int maxItems = 0;
  unsigned int maxUnroll = 0;
  unsigned int maxLoopBodySize = 0;
  AstroData::Observation observation;
  PulsarSearch::DedispersionConf conf;
  cl::Event event;

	try {
    isa::utils::ArgumentList args(argc, argv);

		nrIterations = args.getSwitchArgument< unsigned int >("-iterations");
		clPlatformID = args.getSwitchArgument< unsigned int >("-opencl_platform");
		clDeviceID = args.getSwitchArgument< unsigned int >("-opencl_device");
    conf.setSplitSeconds(args.getSwitch("-split_seconds"));
    conf.setLocalMem(args.getSwitch("-local"));
		observation.setPadding(args.getSwitchArgument< unsigned int >("-padding"));
    inputBits = args.getSwitchArgument< unsigned int >("-input_bits");
    threadUnit = args.getSwitchArgument< unsigned int >("-thread_unit");
		minThreads = args.getSwitchArgument< unsigned int >("-min_threads");
		maxThreads = args.getSwitchArgument< unsigned int >("-max_threads");
		maxRows = args.getSwitchArgument< unsigned int >("-max_rows");
		maxColumns = args.getSwitchArgument< unsigned int >("-max_columns");
    threadIncrement = args.getSwitchArgument< unsigned int >("-thread_increment");
		maxItems = args.getSwitchArgument< unsigned int >("-max_items");
    maxUnroll = args.getSwitchArgument< unsigned int >("-max_unroll");
    maxLoopBodySize = args.getSwitchArgument< unsigned int >("-max_loopsize");
    observation.setFrequencyRange(args.getSwitchArgument< unsigned int >("-channels"), args.getSwitchArgument< float >("-min_freq"), args.getSwitchArgument< float >("-channel_bandwidth"));
		observation.setNrSamplesPerSecond(args.getSwitchArgument< unsigned int >("-samples"));
    observation.setDMRange(args.getSwitchArgument< unsigned int >("-dms"), args.getSwitchArgument< float >("-dm_first"), args.getSwitchArgument< float >("-dm_step"));
	} catch ( isa::utils::EmptyCommandLine & err ) {
		std::cerr << argv[0] << " -iterations ... -opencl_platform ... -opencl_device ... [-split_seconds] [-local] -padding ... -thread_unit ... -min_threads ... -max_threads ... -max_items ... -max_unroll ... -max_loopsize ... -max_columns ... -max_rows ... -thread_increment ... -min_freq ... -channel_bandwidth ... -samples ... -channels ... -dms ... -dm_first ... -dm_step ..." << std::endl;
		return 1;
	} catch ( std::exception & err ) {
		std::cerr << err.what() << std::endl;
		return 1;
	}

  // Allocate host memory
  std::vector< float > * shifts = PulsarSearch::getShifts(observation);
  if ( conf.getSplitSeconds() ) {
    if ( (observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep())))) % observation.getNrSamplesPerSecond() == 0 ) {
      splitSeconds = (observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep())))) / observation.getNrSamplesPerSecond();
    } else {
      splitSeconds = ((observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep())))) / observation.getNrSamplesPerSecond()) + 1;
    }
  } else {
    observation.setNrSamplesPerDispersedChannel(observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep()))));
  }

	// Initialize OpenCL
	cl::Context clContext;
	std::vector< cl::Platform > * clPlatforms = new std::vector< cl::Platform >();
	std::vector< cl::Device > * clDevices = new std::vector< cl::Device >();
	std::vector< std::vector< cl::CommandQueue > > * clQueues = new std::vector< std::vector < cl::CommandQueue > >();
  isa::OpenCL::initializeOpenCL(clPlatformID, 1, clPlatforms, &clContext, clDevices, clQueues);

	// Allocate device memory
  cl::Buffer shifts_d;
  cl::Buffer dispersedData_d;
  cl::Buffer dedispersedData_d;

  try {
    if ( conf.getSplitSeconds() ) {
      initializeDeviceMemory(clContext, &(clQueues->at(clDeviceID)[0]), shifts, &shifts_d, &dispersedData_d, splitSeconds * observation.getNrChannels() * observation.getNrSamplesPerPaddedSecond(), &dedispersedData_d, observation.getNrDMs() * observation.getNrSamplesPerPaddedSecond());
    } else {
      initializeDeviceMemory(clContext, &(clQueues->at(clDeviceID)[0]), shifts, &shifts_d, &dispersedData_d, observation.getNrChannels() * observation.getNrSamplesPerDispersedChannel(), &dedispersedData_d, observation.getNrDMs() * observation.getNrSamplesPerPaddedSecond());
    }
  } catch ( cl::Error & err ) {
    std::cerr << err.what() << std::endl;
    return -1;
  }

	// Find the parameters
	std::vector< unsigned int > samplesPerBlock;
	for ( unsigned int samples = minThreads; samples <= maxColumns; samples += threadIncrement ) {
		if ( (observation.getNrSamplesPerSecond() % samples) == 0 || (observation.getNrSamplesPerPaddedSecond() % samples) == 0 ) {
			samplesPerBlock.push_back(samples);
		}
	}
	std::vector< unsigned int > DMsPerBlock;
	for ( unsigned int DMs = 1; DMs <= maxRows; DMs++ ) {
		if ( (observation.getNrDMs() % DMs) == 0 ) {
			DMsPerBlock.push_back(DMs);
		}
	}

	std::cout << std::fixed << std::endl;
	std::cout << "# nrDMs nrChannels nrSamples splitSeconds local unroll samplesPerBlock DMsPerBlock samplesPerThread DMsPerThread GFLOP/s time stdDeviation COV" << std::endl << std::endl;

	for ( std::vector< unsigned int >::iterator samples = samplesPerBlock.begin(); samples != samplesPerBlock.end(); ++samples ) {
    conf.setNrSamplesPerBlock(*samples);

		for ( std::vector< unsigned int >::iterator DMs = DMsPerBlock.begin(); DMs != DMsPerBlock.end(); ++DMs ) {
      conf.setNrDMsPerBlock(*DMs);
			if ( conf.getNrSamplesPerBlock() * conf.getNrDMsPerBlock() > maxThreads ) {
				break;
			} else if ( (conf.getNrSamplesPerBlock() * conf.getNrDMsPerBlock()) % threadUnit != 0 ) {
        continue;
      }

			for ( unsigned int samplesPerThread = 1; samplesPerThread <= maxItems; samplesPerThread++ ) {
        conf.setNrSamplesPerThread(samplesPerThread);
				if ( (observation.getNrSamplesPerSecond() % (conf.getNrSamplesPerBlock() * conf.getNrSamplesPerThread())) != 0 && (observation.getNrSamplesPerPaddedSecond() % (conf.getNrSamplesPerBlock() * conf.getNrSamplesPerThread())) != 0 ) {
					continue;
				}

				for ( unsigned int DMsPerThread = 1; DMsPerThread <= maxItems; DMsPerThread++ ) {
          conf.setNrDMsPerThread(DMsPerThread);
					if ( (observation.getNrDMs() % (conf.getNrDMsPerBlock() * conf.getNrDMsPerThread())) != 0 ) {
						continue;
					} else if ( (conf.getNrSamplesPerThread() * conf.getNrDMsPerThread()) + conf.getNrDMsPerThread() > maxItems ) {
						break;
					}

          for ( unsigned int unroll = 1; unroll <= maxUnroll; unroll++ ) {
            conf.setUnroll(unroll);
            if ( (observation.getNrChannels() - 1) % conf.getUnroll() != 0 ) {
              continue;
            } else if ( (conf.getNrSamplesPerThread() * conf.getNrDMsPerThread() * conf.getUnroll()) > maxLoopBodySize ) {
              break;
            }
            // Generate kernel
            double gflops = isa::utils::giga(static_cast< long long unsigned int >(observation.getNrDMs()) * observation.getNrChannels() * observation.getNrSamplesPerSecond());
            isa::utils::Timer timer;
            cl::Kernel * kernel;
            std::string * code = PulsarSearch::getDedispersionOpenCL(conf, inputBits, inputTypeName, intermediateTypeName, outputTypeName, observation, *shifts);

            if ( reInit ) {
              delete clQueues;
              clQueues = new std::vector< std::vector < cl::CommandQueue > >();
              isa::OpenCL::initializeOpenCL(clPlatformID, 1, clPlatforms, &clContext, clDevices, clQueues);
              try {
                initializeDeviceMemory(clContext, &(clQueues->at(clDeviceID)[0]), shifts, &shifts_d, &dispersedData_d, observation.getNrChannels() * observation.getNrSamplesPerDispersedChannel(), &dedispersedData_d, observation.getNrDMs() * observation.getNrSamplesPerPaddedSecond());
              } catch ( cl::Error & err ) {
                return -1;
              }
              reInit = false;
            }
            try {
              kernel = isa::OpenCL::compile("dedispersion", *code, "-cl-mad-enable -Werror", clContext, clDevices->at(clDeviceID));
            } catch ( isa::OpenCL::OpenCLError & err ) {
              std::cerr << err.what() << std::endl;
              delete code;
              break;
            }
            delete code;

            unsigned int nrThreads = 0;
            if ( observation.getNrSamplesPerSecond() % (conf.getNrSamplesPerBlock() * conf.getNrSamplesPerThread()) == 0 ) {
              nrThreads = observation.getNrSamplesPerSecond() / conf.getNrSamplesPerThread();
            } else {
              nrThreads = observation.getNrSamplesPerPaddedSecond() / conf.getNrSamplesPerThread();
            }
            cl::NDRange global(nrThreads, observation.getNrDMs() / conf.getNrDMsPerThread());
            cl::NDRange local(conf.getNrSamplesPerBlock(), conf.getNrDMsPerBlock());

            kernel->setArg(0, dispersedData_d);
            kernel->setArg(1, dedispersedData_d);
            kernel->setArg(2, shifts_d);

            try {
              // Warm-up run
              clQueues->at(clDeviceID)[0].finish();
              clQueues->at(clDeviceID)[0].enqueueNDRangeKernel(*kernel, cl::NullRange, global, local, 0, &event);
              event.wait();
              // Tuning runs
              for ( unsigned int iteration = 0; iteration < nrIterations; iteration++ ) {
                timer.start();
                clQueues->at(clDeviceID)[0].enqueueNDRangeKernel(*kernel, cl::NullRange, global, local, 0, &event);
                event.wait();
                timer.stop();
              }
            } catch ( cl::Error & err ) {
              std::cerr << "OpenCL error kernel execution (";
              std::cerr << conf.print() << "): ";
              std::cerr << isa::utils::toString(err.err()) << "." << std::endl;
              delete kernel;
              if ( err.err() == -4 || err.err() == -61 ) {
                return -1;
              }
              reInit = true;
              break;
            }
            delete kernel;

            std::cout << observation.getNrDMs() << " " << observation.getNrChannels() << " " << observation.getNrSamplesPerSecond() << " ";
            std::cout << conf.print() << " ";
            std::cout << std::setprecision(3);
            std::cout << gflops / timer.getAverageTime() << " ";
            std::cout << std::setprecision(6);
            std::cout << timer.getAverageTime() << " " << timer.getStandardDeviation() << " ";
            std::cout << timer.getCoefficientOfVariation() <<  std::endl;
          }
				}
			}
		}
	}

	std::cout << std::endl;

	return 0;
}

void initializeDeviceMemory(cl::Context & clContext, cl::CommandQueue * clQueue, std::vector< float > * shifts, cl::Buffer * shifts_d, cl::Buffer * dispersedData_d, const unsigned int dispersedData_size, cl::Buffer * dedispersedData_d, const unsigned int dedispersedData_size) {
  try {
    *shifts_d = cl::Buffer(clContext, CL_MEM_READ_ONLY, shifts->size() * sizeof(float), 0, 0);
    *dispersedData_d = cl::Buffer(clContext, CL_MEM_READ_ONLY, dispersedData_size * sizeof(inputDataType), 0, 0);
    *dedispersedData_d = cl::Buffer(clContext, CL_MEM_READ_WRITE, dedispersedData_size * sizeof(outputDataType), 0, 0);
    clQueue->enqueueWriteBuffer(*shifts_d, CL_FALSE, 0, shifts->size() * sizeof(float), reinterpret_cast< void * >(shifts->data()));
    clQueue->finish();
  } catch ( cl::Error & err ) {
    std::cerr << "OpenCL error: " << isa::utils::toString(err.err()) << "." << std::endl;
    throw;
  }
}

<commit_msg>Added -input_bits to the command line help message.<commit_after>// Copyright 2014 Alessio Sclocco <a.sclocco@vu.nl>
//
// 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 <iostream>
#include <string>
#include <vector>
#include <exception>
#include <fstream>
#include <iomanip>
#include <limits>

#include <ArgumentList.hpp>
#include <Observation.hpp>
#include <InitializeOpenCL.hpp>
#include <Kernel.hpp>
#include <Shifts.hpp>
#include <Dedispersion.hpp>
#include <utils.hpp>
#include <Timer.hpp>
#include <Stats.hpp>

// Define the data types
typedef float inputDataType;
std::string inputTypeName("float");
std::string intermediateTypeName("float");
typedef float outputDataType;
std::string outputTypeName("float");

void initializeDeviceMemory(cl::Context & clContext, cl::CommandQueue * clQueue, std::vector< float > * shifts, cl::Buffer * shifts_d, cl::Buffer * dispersedData_d, const unsigned int dispersedData_size, cl::Buffer * dedispersedData_d, const unsigned int dedispersedData_size);

int main(int argc, char * argv[]) {
  bool reInit = false;
  uint8_t inputBits = 0;
  unsigned int splitSeconds = 0;
	unsigned int nrIterations = 0;
	unsigned int clPlatformID = 0;
	unsigned int clDeviceID = 0;
	unsigned int minThreads = 0;
  unsigned int maxThreads = 0;
	unsigned int maxRows = 0;
	unsigned int maxColumns = 0;
  unsigned int threadUnit = 0;
  unsigned int threadIncrement = 0;
  unsigned int maxItems = 0;
  unsigned int maxUnroll = 0;
  unsigned int maxLoopBodySize = 0;
  AstroData::Observation observation;
  PulsarSearch::DedispersionConf conf;
  cl::Event event;

	try {
    isa::utils::ArgumentList args(argc, argv);

		nrIterations = args.getSwitchArgument< unsigned int >("-iterations");
		clPlatformID = args.getSwitchArgument< unsigned int >("-opencl_platform");
		clDeviceID = args.getSwitchArgument< unsigned int >("-opencl_device");
    conf.setSplitSeconds(args.getSwitch("-split_seconds"));
    conf.setLocalMem(args.getSwitch("-local"));
		observation.setPadding(args.getSwitchArgument< unsigned int >("-padding"));
    inputBits = args.getSwitchArgument< unsigned int >("-input_bits");
    threadUnit = args.getSwitchArgument< unsigned int >("-thread_unit");
		minThreads = args.getSwitchArgument< unsigned int >("-min_threads");
		maxThreads = args.getSwitchArgument< unsigned int >("-max_threads");
		maxRows = args.getSwitchArgument< unsigned int >("-max_rows");
		maxColumns = args.getSwitchArgument< unsigned int >("-max_columns");
    threadIncrement = args.getSwitchArgument< unsigned int >("-thread_increment");
		maxItems = args.getSwitchArgument< unsigned int >("-max_items");
    maxUnroll = args.getSwitchArgument< unsigned int >("-max_unroll");
    maxLoopBodySize = args.getSwitchArgument< unsigned int >("-max_loopsize");
    observation.setFrequencyRange(args.getSwitchArgument< unsigned int >("-channels"), args.getSwitchArgument< float >("-min_freq"), args.getSwitchArgument< float >("-channel_bandwidth"));
		observation.setNrSamplesPerSecond(args.getSwitchArgument< unsigned int >("-samples"));
    observation.setDMRange(args.getSwitchArgument< unsigned int >("-dms"), args.getSwitchArgument< float >("-dm_first"), args.getSwitchArgument< float >("-dm_step"));
	} catch ( isa::utils::EmptyCommandLine & err ) {
		std::cerr << argv[0] << " -iterations ... -opencl_platform ... -opencl_device ... [-split_seconds] [-local] -input_bits ... -padding ... -thread_unit ... -min_threads ... -max_threads ... -max_items ... -max_unroll ... -max_loopsize ... -max_columns ... -max_rows ... -thread_increment ... -min_freq ... -channel_bandwidth ... -samples ... -channels ... -dms ... -dm_first ... -dm_step ..." << std::endl;
		return 1;
	} catch ( std::exception & err ) {
		std::cerr << err.what() << std::endl;
		return 1;
	}

  // Allocate host memory
  std::vector< float > * shifts = PulsarSearch::getShifts(observation);
  if ( conf.getSplitSeconds() ) {
    if ( (observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep())))) % observation.getNrSamplesPerSecond() == 0 ) {
      splitSeconds = (observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep())))) / observation.getNrSamplesPerSecond();
    } else {
      splitSeconds = ((observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep())))) / observation.getNrSamplesPerSecond()) + 1;
    }
  } else {
    observation.setNrSamplesPerDispersedChannel(observation.getNrSamplesPerSecond() + static_cast< unsigned int >(shifts->at(0) * (observation.getFirstDM() + ((observation.getNrDMs() - 1) * observation.getDMStep()))));
  }

	// Initialize OpenCL
	cl::Context clContext;
	std::vector< cl::Platform > * clPlatforms = new std::vector< cl::Platform >();
	std::vector< cl::Device > * clDevices = new std::vector< cl::Device >();
	std::vector< std::vector< cl::CommandQueue > > * clQueues = new std::vector< std::vector < cl::CommandQueue > >();
  isa::OpenCL::initializeOpenCL(clPlatformID, 1, clPlatforms, &clContext, clDevices, clQueues);

	// Allocate device memory
  cl::Buffer shifts_d;
  cl::Buffer dispersedData_d;
  cl::Buffer dedispersedData_d;

  try {
    if ( conf.getSplitSeconds() ) {
      initializeDeviceMemory(clContext, &(clQueues->at(clDeviceID)[0]), shifts, &shifts_d, &dispersedData_d, splitSeconds * observation.getNrChannels() * observation.getNrSamplesPerPaddedSecond(), &dedispersedData_d, observation.getNrDMs() * observation.getNrSamplesPerPaddedSecond());
    } else {
      initializeDeviceMemory(clContext, &(clQueues->at(clDeviceID)[0]), shifts, &shifts_d, &dispersedData_d, observation.getNrChannels() * observation.getNrSamplesPerDispersedChannel(), &dedispersedData_d, observation.getNrDMs() * observation.getNrSamplesPerPaddedSecond());
    }
  } catch ( cl::Error & err ) {
    std::cerr << err.what() << std::endl;
    return -1;
  }

	// Find the parameters
	std::vector< unsigned int > samplesPerBlock;
	for ( unsigned int samples = minThreads; samples <= maxColumns; samples += threadIncrement ) {
		if ( (observation.getNrSamplesPerSecond() % samples) == 0 || (observation.getNrSamplesPerPaddedSecond() % samples) == 0 ) {
			samplesPerBlock.push_back(samples);
		}
	}
	std::vector< unsigned int > DMsPerBlock;
	for ( unsigned int DMs = 1; DMs <= maxRows; DMs++ ) {
		if ( (observation.getNrDMs() % DMs) == 0 ) {
			DMsPerBlock.push_back(DMs);
		}
	}

	std::cout << std::fixed << std::endl;
	std::cout << "# nrDMs nrChannels nrSamples splitSeconds local unroll samplesPerBlock DMsPerBlock samplesPerThread DMsPerThread GFLOP/s time stdDeviation COV" << std::endl << std::endl;

	for ( std::vector< unsigned int >::iterator samples = samplesPerBlock.begin(); samples != samplesPerBlock.end(); ++samples ) {
    conf.setNrSamplesPerBlock(*samples);

		for ( std::vector< unsigned int >::iterator DMs = DMsPerBlock.begin(); DMs != DMsPerBlock.end(); ++DMs ) {
      conf.setNrDMsPerBlock(*DMs);
			if ( conf.getNrSamplesPerBlock() * conf.getNrDMsPerBlock() > maxThreads ) {
				break;
			} else if ( (conf.getNrSamplesPerBlock() * conf.getNrDMsPerBlock()) % threadUnit != 0 ) {
        continue;
      }

			for ( unsigned int samplesPerThread = 1; samplesPerThread <= maxItems; samplesPerThread++ ) {
        conf.setNrSamplesPerThread(samplesPerThread);
				if ( (observation.getNrSamplesPerSecond() % (conf.getNrSamplesPerBlock() * conf.getNrSamplesPerThread())) != 0 && (observation.getNrSamplesPerPaddedSecond() % (conf.getNrSamplesPerBlock() * conf.getNrSamplesPerThread())) != 0 ) {
					continue;
				}

				for ( unsigned int DMsPerThread = 1; DMsPerThread <= maxItems; DMsPerThread++ ) {
          conf.setNrDMsPerThread(DMsPerThread);
					if ( (observation.getNrDMs() % (conf.getNrDMsPerBlock() * conf.getNrDMsPerThread())) != 0 ) {
						continue;
					} else if ( (conf.getNrSamplesPerThread() * conf.getNrDMsPerThread()) + conf.getNrDMsPerThread() > maxItems ) {
						break;
					}

          for ( unsigned int unroll = 1; unroll <= maxUnroll; unroll++ ) {
            conf.setUnroll(unroll);
            if ( (observation.getNrChannels() - 1) % conf.getUnroll() != 0 ) {
              continue;
            } else if ( (conf.getNrSamplesPerThread() * conf.getNrDMsPerThread() * conf.getUnroll()) > maxLoopBodySize ) {
              break;
            }
            // Generate kernel
            double gflops = isa::utils::giga(static_cast< long long unsigned int >(observation.getNrDMs()) * observation.getNrChannels() * observation.getNrSamplesPerSecond());
            isa::utils::Timer timer;
            cl::Kernel * kernel;
            std::string * code = PulsarSearch::getDedispersionOpenCL(conf, inputBits, inputTypeName, intermediateTypeName, outputTypeName, observation, *shifts);

            if ( reInit ) {
              delete clQueues;
              clQueues = new std::vector< std::vector < cl::CommandQueue > >();
              isa::OpenCL::initializeOpenCL(clPlatformID, 1, clPlatforms, &clContext, clDevices, clQueues);
              try {
                initializeDeviceMemory(clContext, &(clQueues->at(clDeviceID)[0]), shifts, &shifts_d, &dispersedData_d, observation.getNrChannels() * observation.getNrSamplesPerDispersedChannel(), &dedispersedData_d, observation.getNrDMs() * observation.getNrSamplesPerPaddedSecond());
              } catch ( cl::Error & err ) {
                return -1;
              }
              reInit = false;
            }
            try {
              kernel = isa::OpenCL::compile("dedispersion", *code, "-cl-mad-enable -Werror", clContext, clDevices->at(clDeviceID));
            } catch ( isa::OpenCL::OpenCLError & err ) {
              std::cerr << err.what() << std::endl;
              delete code;
              break;
            }
            delete code;

            unsigned int nrThreads = 0;
            if ( observation.getNrSamplesPerSecond() % (conf.getNrSamplesPerBlock() * conf.getNrSamplesPerThread()) == 0 ) {
              nrThreads = observation.getNrSamplesPerSecond() / conf.getNrSamplesPerThread();
            } else {
              nrThreads = observation.getNrSamplesPerPaddedSecond() / conf.getNrSamplesPerThread();
            }
            cl::NDRange global(nrThreads, observation.getNrDMs() / conf.getNrDMsPerThread());
            cl::NDRange local(conf.getNrSamplesPerBlock(), conf.getNrDMsPerBlock());

            kernel->setArg(0, dispersedData_d);
            kernel->setArg(1, dedispersedData_d);
            kernel->setArg(2, shifts_d);

            try {
              // Warm-up run
              clQueues->at(clDeviceID)[0].finish();
              clQueues->at(clDeviceID)[0].enqueueNDRangeKernel(*kernel, cl::NullRange, global, local, 0, &event);
              event.wait();
              // Tuning runs
              for ( unsigned int iteration = 0; iteration < nrIterations; iteration++ ) {
                timer.start();
                clQueues->at(clDeviceID)[0].enqueueNDRangeKernel(*kernel, cl::NullRange, global, local, 0, &event);
                event.wait();
                timer.stop();
              }
            } catch ( cl::Error & err ) {
              std::cerr << "OpenCL error kernel execution (";
              std::cerr << conf.print() << "): ";
              std::cerr << isa::utils::toString(err.err()) << "." << std::endl;
              delete kernel;
              if ( err.err() == -4 || err.err() == -61 ) {
                return -1;
              }
              reInit = true;
              break;
            }
            delete kernel;

            std::cout << observation.getNrDMs() << " " << observation.getNrChannels() << " " << observation.getNrSamplesPerSecond() << " ";
            std::cout << conf.print() << " ";
            std::cout << std::setprecision(3);
            std::cout << gflops / timer.getAverageTime() << " ";
            std::cout << std::setprecision(6);
            std::cout << timer.getAverageTime() << " " << timer.getStandardDeviation() << " ";
            std::cout << timer.getCoefficientOfVariation() <<  std::endl;
          }
				}
			}
		}
	}

	std::cout << std::endl;

	return 0;
}

void initializeDeviceMemory(cl::Context & clContext, cl::CommandQueue * clQueue, std::vector< float > * shifts, cl::Buffer * shifts_d, cl::Buffer * dispersedData_d, const unsigned int dispersedData_size, cl::Buffer * dedispersedData_d, const unsigned int dedispersedData_size) {
  try {
    *shifts_d = cl::Buffer(clContext, CL_MEM_READ_ONLY, shifts->size() * sizeof(float), 0, 0);
    *dispersedData_d = cl::Buffer(clContext, CL_MEM_READ_ONLY, dispersedData_size * sizeof(inputDataType), 0, 0);
    *dedispersedData_d = cl::Buffer(clContext, CL_MEM_READ_WRITE, dedispersedData_size * sizeof(outputDataType), 0, 0);
    clQueue->enqueueWriteBuffer(*shifts_d, CL_FALSE, 0, shifts->size() * sizeof(float), reinterpret_cast< void * >(shifts->data()));
    clQueue->finish();
  } catch ( cl::Error & err ) {
    std::cerr << "OpenCL error: " << isa::utils::toString(err.err()) << "." << std::endl;
    throw;
  }
}

<|endoftext|>


<commit_before>#include <iostream>
#include <immintrin.h>
#include <new>
#include <cassert>
#include <sstream>
#include "simd_util.h"
#include "pauli_string.h"
#include <cstring>
#include <cstdlib>
#include <bit>

PauliStringPtr::PauliStringPtr(
            size_t init_size,
            BitPtr init_sign,
            uint64_t *init_x,
            uint64_t *init_z,
            size_t init_stride256) :
        size(init_size),
        bit_ptr_sign(init_sign),
        _x(init_x),
        _z(init_z),
        stride256(init_stride256) {
}

PauliStringVal::PauliStringVal(size_t init_size) :
        val_sign(false),
        x_data(init_size),
        z_data(init_size) {
}

PauliStringVal::PauliStringVal(const PauliStringPtr &other) :
        val_sign(other.bit_ptr_sign.get()),
        x_data(other.size, other._x),
        z_data(other.size, other._z) {
}

PauliStringPtr PauliStringVal::ptr() const {
    return PauliStringPtr(*this);
}

std::string PauliStringVal::str() const {
    return ptr().str();
}

void PauliStringPtr::swap_with(PauliStringPtr &other) {
    assert(size == other.size);
    bit_ptr_sign.swap(other.bit_ptr_sign);
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        std::swap(*x256, *ox256);
        std::swap(*z256, *oz256);
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }
}

void PauliStringPtr::overwrite_with(const PauliStringPtr &other) {
    assert(size == other.size);
    bit_ptr_sign.set(other.bit_ptr_sign.get());
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        *x256 = *ox256;
        *z256 = *oz256;
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }
}

PauliStringVal& PauliStringVal::operator=(const PauliStringPtr &other) noexcept {
    (*this).~PauliStringVal();
    new(this) PauliStringVal(other);
    return *this;
}

PauliStringPtr::PauliStringPtr(const PauliStringVal &other) :
        size(other.x_data.num_bits),
        bit_ptr_sign(BitPtr((void *)&other.val_sign, 0)),
        _x(other.x_data.u64),
        _z(other.z_data.u64),
        stride256(1) {
}

uint8_t PauliStringPtr::log_i_scalar_byproduct(const PauliStringPtr &other) const {
    assert(size == other.size);
    union {__m256i u256; uint64_t u64[4]; } cnt1 {};
    union {__m256i u256; uint64_t u64[4]; } cnt2 {};

    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        auto y0 = *x256 & *z256;
        auto x0 = y0 ^ *x256;
        auto z0 = y0 ^ *z256;
        auto y1 = *ox256 & *oz256;
        auto x1 = y1 ^ *ox256;
        auto z1 = y1 ^ *oz256;
        auto f = (x0 & y1) | (y0 & z1) | (z0 & x1);
        auto b = (x0 & z1) | (y0 & x1) | (z0 & y1);
        cnt1.u256 ^= b;
        cnt2.u256 ^= cnt1.u256 & (f ^ b);
        cnt1.u256 ^= f;
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }
    size_t s = 0;
    for (size_t k = 0; k < 4; k++) {
        s += (uint8_t) std::popcount(cnt1.u64[k]);
        s ^= (uint8_t) std::popcount(cnt2.u64[k]) << 1;
    }
    return s & 3;
}

std::string PauliStringPtr::str() const {
    std::stringstream ss;
    ss << *this;
    return ss.str();
}

bool PauliStringPtr::operator==(const PauliStringPtr &other) const {
    if (size != other.size || bit_ptr_sign.get() != other.bit_ptr_sign.get()) {
        return false;
    }
    __m256i acc {};
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        acc |= *x256 ^ *ox256;
        acc |= *z256 ^ *oz256;
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }

    auto acc64 = (uint64_t *)&acc;
    for (size_t k = 0; k < 4; k++) {
        if (acc64[k]) {
            return false;
        }
    }

    return true;
}

bool PauliStringPtr::operator!=(const PauliStringPtr &other) const {
    return !(*this == other);
}

std::ostream &operator<<(std::ostream &out, const PauliStringVal &ps) {
    return out << ps.ptr();
}

size_t PauliStringPtr::num_words256() const {
    return ceil256(size) >> 8;
}

std::ostream &operator<<(std::ostream &out, const PauliStringPtr &ps) {
    out << (ps.bit_ptr_sign.get() ? '-' : '+');
    auto x256 = (__m256i *)ps._x;
    auto z256 = (__m256i *)ps._z;
    auto end = &x256[ps.num_words256() * ps.stride256];
    size_t remaining = ps.size;
    while (x256 != end) {
        auto xs = m256i_to_bits(*x256);
        auto zs = m256i_to_bits(*z256);
        for (int j = 0; j < 256 && remaining; j++) {
            out << "_XZY"[xs[j] + 2 * zs[j]];
            remaining--;
        }
        x256 += ps.stride256;
        z256 += ps.stride256;
    }
    return out;
}

PauliStringVal PauliStringVal::from_pattern(bool sign, size_t size, const std::function<char(size_t)> &func) {
    PauliStringVal result(size);
    result.val_sign = sign;
    for (size_t i = 0; i < size; i++) {
        char c = func(i);
        bool x;
        bool z;
        if (c == 'X') {
            x = true;
            z = false;
        } else if (c == 'Y') {
            x = true;
            z = true;
        } else if (c == 'Z') {
            x = false;
            z = true;
        } else if (c == '_' || c == 'I') {
            x = false;
            z = false;
        } else {
            throw std::runtime_error("Unrecognized pauli character. " + std::to_string(c));
        }
        result.x_data.u64[i / 64] ^= (uint64_t)x << (i & 63);
        result.z_data.u64[i / 64] ^= (uint64_t)z << (i & 63);
    }
    return result;
}

PauliStringVal PauliStringVal::from_str(const char *text) {
    auto sign = text[0] == '-';
    if (text[0] == '+' || text[0] == '-') {
        text++;
    }
    return PauliStringVal::from_pattern(sign, strlen(text), [&](size_t i){ return text[i]; });
}

PauliStringVal PauliStringVal::identity(size_t size) {
    return PauliStringVal(size);
}

PauliStringPtr& PauliStringPtr::operator*=(const PauliStringPtr& rhs) {
    uint8_t log_i = inplace_right_mul_with_scalar_output(rhs);
    assert((log_i & 1) == 0);
    bit_ptr_sign.toggle_if(log_i & 2);
    return *this;
}

uint8_t PauliStringPtr::inplace_right_mul_with_scalar_output(const PauliStringPtr& rhs) {
    uint8_t result = log_i_scalar_byproduct(rhs);
    bit_ptr_sign.toggle_if(rhs.bit_ptr_sign.get());
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)rhs._x;
    auto oz256 = (__m256i *)rhs._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        *x256 ^= *ox256;
        *z256 ^= *oz256;
        x256 += stride256;
        z256 += stride256;
        ox256 += rhs.stride256;
        oz256 += rhs.stride256;
    }
    return result;
}

bool PauliStringPtr::get_x_bit(size_t k) const {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    return ((_x[i0] >> i1) & 1) != 0;
}

bool PauliStringPtr::get_z_bit(size_t k) const {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    return ((_z[i0] >> i1) & 1) != 0;
}

void PauliStringPtr::toggle_x_bit(size_t k) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _x[i0] ^= 1ull << i1;
}

void PauliStringPtr::toggle_z_bit(size_t k) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _z[i0] ^= 1ull << i1;
}

void PauliStringPtr::set_x_bit(size_t k, bool val) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _x[i0] &= ~(1ull << i1);
    _x[i0] ^= (uint64_t)val << i1;
}

void PauliStringPtr::set_z_bit(size_t k, bool val) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _z[i0] &= ~(1ull << i1);
    _z[i0] ^= (uint64_t)val << i1;
}

void PauliStringPtr::gather_into(PauliStringPtr &out, const std::vector<size_t> &in_indices) const {
    assert(in_indices.size() == out.size);
    for (size_t k_out = 0; k_out < out.size; k_out++) {
        size_t k_in = in_indices[k_out];
        out.set_x_bit(k_out, get_x_bit(k_in));
        out.set_z_bit(k_out, get_z_bit(k_in));
    }
}

void PauliStringPtr::scatter_into(PauliStringPtr &out, const std::vector<size_t> &out_indices) const {
    assert(size == out_indices.size());
    for (size_t k_in = 0; k_in < size; k_in++) {
        size_t k_out = out_indices[k_in];
        out.set_x_bit(k_out, get_x_bit(k_in));
        out.set_z_bit(k_out, get_z_bit(k_in));
    }
    out.bit_ptr_sign.toggle_if(bit_ptr_sign.get());
}

bool PauliStringVal::operator==(const PauliStringPtr &other) const {
    return ptr() == other;
}
bool PauliStringVal::operator!=(const PauliStringPtr &other) const {
    return ptr() != other;
}
<commit_msg>Make Pauli product multiplication another 50% faster - Found a way to compute the forward/back bits in fewer operations - Runs at 90 GigaPaulis per second (!) (!!)<commit_after>#include <iostream>
#include <immintrin.h>
#include <new>
#include <cassert>
#include <sstream>
#include "simd_util.h"
#include "pauli_string.h"
#include <cstring>
#include <cstdlib>
#include <bit>

PauliStringPtr::PauliStringPtr(
            size_t init_size,
            BitPtr init_sign,
            uint64_t *init_x,
            uint64_t *init_z,
            size_t init_stride256) :
        size(init_size),
        bit_ptr_sign(init_sign),
        _x(init_x),
        _z(init_z),
        stride256(init_stride256) {
}

PauliStringVal::PauliStringVal(size_t init_size) :
        val_sign(false),
        x_data(init_size),
        z_data(init_size) {
}

PauliStringVal::PauliStringVal(const PauliStringPtr &other) :
        val_sign(other.bit_ptr_sign.get()),
        x_data(other.size, other._x),
        z_data(other.size, other._z) {
}

PauliStringPtr PauliStringVal::ptr() const {
    return PauliStringPtr(*this);
}

std::string PauliStringVal::str() const {
    return ptr().str();
}

void PauliStringPtr::swap_with(PauliStringPtr &other) {
    assert(size == other.size);
    bit_ptr_sign.swap(other.bit_ptr_sign);
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        std::swap(*x256, *ox256);
        std::swap(*z256, *oz256);
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }
}

void PauliStringPtr::overwrite_with(const PauliStringPtr &other) {
    assert(size == other.size);
    bit_ptr_sign.set(other.bit_ptr_sign.get());
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        *x256 = *ox256;
        *z256 = *oz256;
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }
}

PauliStringVal& PauliStringVal::operator=(const PauliStringPtr &other) noexcept {
    (*this).~PauliStringVal();
    new(this) PauliStringVal(other);
    return *this;
}

PauliStringPtr::PauliStringPtr(const PauliStringVal &other) :
        size(other.x_data.num_bits),
        bit_ptr_sign(BitPtr((void *)&other.val_sign, 0)),
        _x(other.x_data.u64),
        _z(other.z_data.u64),
        stride256(1) {
}

uint8_t PauliStringPtr::log_i_scalar_byproduct(const PauliStringPtr &other) const {
    assert(size == other.size);
    union {__m256i u256; uint64_t u64[4]; } cnt1 {};
    union {__m256i u256; uint64_t u64[4]; } cnt2 {};
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        auto x1 = *x256;
        auto z2 = *oz256;
        auto z1 = *z256;
        auto x2 = *ox256;
        auto t1 = x1 & z2;
        auto t2 = x2 & z1;
        auto f = t1 ^ t2;
        auto b = ((x1 ^ z2) & t2) ^ _mm256_andnot_si256(z1, _mm256_andnot_si256(x2, t1));
        cnt1.u256 ^= b;
        cnt2.u256 ^= cnt1.u256 & f;
        cnt1.u256 ^= f ^ b;
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }
    size_t s = 0;
    for (size_t k = 0; k < 4; k++) {
        s += (uint8_t) std::popcount(cnt1.u64[k]);
        s ^= (uint8_t) std::popcount(cnt2.u64[k]) << 1;
    }
    return s & 3;
}

std::string PauliStringPtr::str() const {
    std::stringstream ss;
    ss << *this;
    return ss.str();
}

bool PauliStringPtr::operator==(const PauliStringPtr &other) const {
    if (size != other.size || bit_ptr_sign.get() != other.bit_ptr_sign.get()) {
        return false;
    }
    __m256i acc {};
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)other._x;
    auto oz256 = (__m256i *)other._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        acc |= *x256 ^ *ox256;
        acc |= *z256 ^ *oz256;
        x256 += stride256;
        z256 += stride256;
        ox256 += other.stride256;
        oz256 += other.stride256;
    }

    auto acc64 = (uint64_t *)&acc;
    for (size_t k = 0; k < 4; k++) {
        if (acc64[k]) {
            return false;
        }
    }

    return true;
}

bool PauliStringPtr::operator!=(const PauliStringPtr &other) const {
    return !(*this == other);
}

std::ostream &operator<<(std::ostream &out, const PauliStringVal &ps) {
    return out << ps.ptr();
}

size_t PauliStringPtr::num_words256() const {
    return ceil256(size) >> 8;
}

std::ostream &operator<<(std::ostream &out, const PauliStringPtr &ps) {
    out << (ps.bit_ptr_sign.get() ? '-' : '+');
    auto x256 = (__m256i *)ps._x;
    auto z256 = (__m256i *)ps._z;
    auto end = &x256[ps.num_words256() * ps.stride256];
    size_t remaining = ps.size;
    while (x256 != end) {
        auto xs = m256i_to_bits(*x256);
        auto zs = m256i_to_bits(*z256);
        for (int j = 0; j < 256 && remaining; j++) {
            out << "_XZY"[xs[j] + 2 * zs[j]];
            remaining--;
        }
        x256 += ps.stride256;
        z256 += ps.stride256;
    }
    return out;
}

PauliStringVal PauliStringVal::from_pattern(bool sign, size_t size, const std::function<char(size_t)> &func) {
    PauliStringVal result(size);
    result.val_sign = sign;
    for (size_t i = 0; i < size; i++) {
        char c = func(i);
        bool x;
        bool z;
        if (c == 'X') {
            x = true;
            z = false;
        } else if (c == 'Y') {
            x = true;
            z = true;
        } else if (c == 'Z') {
            x = false;
            z = true;
        } else if (c == '_' || c == 'I') {
            x = false;
            z = false;
        } else {
            throw std::runtime_error("Unrecognized pauli character. " + std::to_string(c));
        }
        result.x_data.u64[i / 64] ^= (uint64_t)x << (i & 63);
        result.z_data.u64[i / 64] ^= (uint64_t)z << (i & 63);
    }
    return result;
}

PauliStringVal PauliStringVal::from_str(const char *text) {
    auto sign = text[0] == '-';
    if (text[0] == '+' || text[0] == '-') {
        text++;
    }
    return PauliStringVal::from_pattern(sign, strlen(text), [&](size_t i){ return text[i]; });
}

PauliStringVal PauliStringVal::identity(size_t size) {
    return PauliStringVal(size);
}

PauliStringPtr& PauliStringPtr::operator*=(const PauliStringPtr& rhs) {
    uint8_t log_i = inplace_right_mul_with_scalar_output(rhs);
    assert((log_i & 1) == 0);
    bit_ptr_sign.toggle_if(log_i & 2);
    return *this;
}

uint8_t PauliStringPtr::inplace_right_mul_with_scalar_output(const PauliStringPtr& rhs) {
    uint8_t result = log_i_scalar_byproduct(rhs);
    bit_ptr_sign.toggle_if(rhs.bit_ptr_sign.get());
    auto x256 = (__m256i *)_x;
    auto z256 = (__m256i *)_z;
    auto ox256 = (__m256i *)rhs._x;
    auto oz256 = (__m256i *)rhs._z;
    auto end = &x256[num_words256() * stride256];
    while (x256 != end) {
        *x256 ^= *ox256;
        *z256 ^= *oz256;
        x256 += stride256;
        z256 += stride256;
        ox256 += rhs.stride256;
        oz256 += rhs.stride256;
    }
    return result;
}

bool PauliStringPtr::get_x_bit(size_t k) const {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    return ((_x[i0] >> i1) & 1) != 0;
}

bool PauliStringPtr::get_z_bit(size_t k) const {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    return ((_z[i0] >> i1) & 1) != 0;
}

void PauliStringPtr::toggle_x_bit(size_t k) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _x[i0] ^= 1ull << i1;
}

void PauliStringPtr::toggle_z_bit(size_t k) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _z[i0] ^= 1ull << i1;
}

void PauliStringPtr::set_x_bit(size_t k, bool val) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _x[i0] &= ~(1ull << i1);
    _x[i0] ^= (uint64_t)val << i1;
}

void PauliStringPtr::set_z_bit(size_t k, bool val) {
    k = (k & ~0xFF)*stride256 | (k & 0xFF);
    size_t i0 = k >> 6;
    size_t i1 = k & 63;
    _z[i0] &= ~(1ull << i1);
    _z[i0] ^= (uint64_t)val << i1;
}

void PauliStringPtr::gather_into(PauliStringPtr &out, const std::vector<size_t> &in_indices) const {
    assert(in_indices.size() == out.size);
    for (size_t k_out = 0; k_out < out.size; k_out++) {
        size_t k_in = in_indices[k_out];
        out.set_x_bit(k_out, get_x_bit(k_in));
        out.set_z_bit(k_out, get_z_bit(k_in));
    }
}

void PauliStringPtr::scatter_into(PauliStringPtr &out, const std::vector<size_t> &out_indices) const {
    assert(size == out_indices.size());
    for (size_t k_in = 0; k_in < size; k_in++) {
        size_t k_out = out_indices[k_in];
        out.set_x_bit(k_out, get_x_bit(k_in));
        out.set_z_bit(k_out, get_z_bit(k_in));
    }
    out.bit_ptr_sign.toggle_if(bit_ptr_sign.get());
}

bool PauliStringVal::operator==(const PauliStringPtr &other) const {
    return ptr() == other;
}
bool PauliStringVal::operator!=(const PauliStringPtr &other) const {
    return ptr() != other;
}
<|endoftext|>


<commit_before>#include "../../../shared/generated/cpp/IntentBase.h"
#include "common/RhoStd.h"
#include "common/AutoPointer.h"
#include "common/RhodesApp.h"
#include "common/RhoConf.h"
#include "logging/RhoLog.h"
#include "SystemImpl.h"

namespace rho {

    using namespace apiGenerator;
    using namespace common;

    class CIntentSingletonImpl: public CIntentSingletonBase
    {
    public:

        CIntentSingletonImpl(): CIntentSingletonBase(){}

        virtual void send( const rho::Hashtable<rho::String, rho::String>& params, rho::apiGenerator::CMethodResult& oResult) {
			// TODO: support all other appropriate params
			String appName = params.get("appName");
			String data = params.get("data");
			if (appName.length() <= 0) {
				oResult.setError("\"appName\" can't be empty");
				return;
			}
			if (data.length() <= 0) {
				oResult.setError("\"data\" can't be empty");
				return;
			}
			CSystemFactory::getSystemSingletonS()->sendApplicationMessage(appName, data, oResult);
        }

        virtual void startListening(rho::apiGenerator::CMethodResult& oResult) {
			CSystemFactory::getSystemSingletonS()->startApplicationMessageNotifications(oResult);
        }

        virtual void stopListening(rho::apiGenerator::CMethodResult& oResult) {
			CSystemFactory::getSystemSingletonS()->stopApplicationMessageNotifications(oResult);
        }

    };

    class CIntentImpl : public CIntentBase
    {
    public:
        virtual ~CIntentImpl() {}

        //methods

    };

    ////////////////////////////////////////////////////////////////////////

    class CIntentFactory: public CIntentFactoryBase    {
    public:
        CIntentFactory(){}

        IIntentSingleton* createModuleSingleton()
        {
            return new CIntentSingletonImpl();
        }

        virtual IIntent* createModuleByID(const rho::String& strID){ return new CIntentImpl(); };

    };

}

extern "C" void Init_Intent()
{
    rho::CIntentFactory::setInstance( new rho::CIntentFactory() );
    rho::Init_Intent_API();

}<commit_msg>wm: Rho.Intent.send(...) accepts empty 'appName' and 'data' parameters<commit_after>#include "../../../shared/generated/cpp/IntentBase.h"
#include "common/RhoStd.h"
#include "common/AutoPointer.h"
#include "common/RhodesApp.h"
#include "common/RhoConf.h"
#include "logging/RhoLog.h"
#include "SystemImpl.h"

namespace rho {

    using namespace apiGenerator;
    using namespace common;

    class CIntentSingletonImpl: public CIntentSingletonBase
    {
    public:

        CIntentSingletonImpl(): CIntentSingletonBase(){}

        virtual void send( const rho::Hashtable<rho::String, rho::String>& params, rho::apiGenerator::CMethodResult& oResult) {
			// TODO: support all other appropriate params
			CSystemFactory::getSystemSingletonS()->sendApplicationMessage(params.get("appName"), params.get("data"), oResult);
        }

        virtual void startListening(rho::apiGenerator::CMethodResult& oResult) {
			CSystemFactory::getSystemSingletonS()->startApplicationMessageNotifications(oResult);
        }

        virtual void stopListening(rho::apiGenerator::CMethodResult& oResult) {
			CSystemFactory::getSystemSingletonS()->stopApplicationMessageNotifications(oResult);
        }

    };

    class CIntentImpl : public CIntentBase
    {
    public:
        virtual ~CIntentImpl() {}

        //methods

    };

    ////////////////////////////////////////////////////////////////////////

    class CIntentFactory: public CIntentFactoryBase    {
    public:
        CIntentFactory(){}

        IIntentSingleton* createModuleSingleton()
        {
            return new CIntentSingletonImpl();
        }

        virtual IIntent* createModuleByID(const rho::String& strID){ return new CIntentImpl(); };

    };

}

extern "C" void Init_Intent()
{
    rho::CIntentFactory::setInstance( new rho::CIntentFactory() );
    rho::Init_Intent_API();

}<|endoftext|>


<commit_before>#include "config.hpp"
#include "core/FPSCamera.h"
#include "core/helpers.hpp"
#include "core/Log.h"
#include "core/LogView.h"
#include "core/Misc.h"
#include "core/node.hpp"
#include "core/ShaderProgramManager.hpp"
#include "core/WindowManager.hpp"

#include <imgui.h>
#include <external/imgui_impl_glfw_gl3.h>

#include <stack>

#include <cstdlib>

int main()
{
	//
	// Set up the logging system
	//
	Log::Init();
	Log::View::Init();

	//
	// Set up the camera
	//
	InputHandler input_handler;
	FPSCameraf camera(0.5f * glm::half_pi<float>(),
	                  static_cast<float>(config::resolution_x) / static_cast<float>(config::resolution_y),
	                  0.01f, 1000.0f);
	camera.mWorld.SetTranslate(glm::vec3(0.0f, 0.0f, 6.0f));
	camera.mMouseSensitivity = 0.003f;
	camera.mMovementSpeed = 0.25f * 12.0f;

	//
	// Set up the windowing system and create the window
	//
	WindowManager window_manager;
	WindowManager::WindowDatum window_datum{ input_handler, camera, config::resolution_x, config::resolution_y, 0, 0, 0, 0};
	GLFWwindow* window = window_manager.CreateWindow("EDAF80: Assignment 1", window_datum, config::msaa_rate);
	if (window == nullptr) {
		LogError("Failed to get a window: exiting.");

		Log::View::Destroy();
		Log::Destroy();

		return EXIT_FAILURE;
	}

	//
	// Load the sphere geometry
	//
	std::vector<bonobo::mesh_data> const objects = bonobo::loadObjects("sphere.obj");
	if (objects.empty()) {
		LogError("Failed to load the sphere geometry: exiting.");

		Log::View::Destroy();
		Log::Destroy();

		return EXIT_FAILURE;
	}
	bonobo::mesh_data const& sphere = objects.front();

	//
	// Create the shader program
	//
	ShaderProgramManager program_manager;
	GLuint shader = 0u;
	program_manager.CreateAndRegisterProgram({ { ShaderType::vertex, "EDAF80/default.vert" },
	                                           { ShaderType::fragment, "EDAF80/default.frag" } },
	                                         shader);
	if (shader == 0u) {
		LogError("Failed to generate the shader program: exiting.");

		Log::View::Destroy();
		Log::Destroy();

		return EXIT_FAILURE;
	}

	//
	// Set up the sun node and other related attributes
	//
	Node sun_node;
	sun_node.set_geometry(sphere);
	GLuint const sun_texture = bonobo::loadTexture2D("sunmap.png");
	sun_node.add_texture("diffuse", sun_texture, GL_TEXTURE_2D);
	float const sun_spin_speed = glm::two_pi<float>() / 6.0f; // Full rotation in six seconds

	Node solar_system_node;
	solar_system_node.add_child(&sun_node);

	//
	// TODO: Create nodes for the remaining of the solar system
	//

	glViewport(0, 0, config::resolution_x, config::resolution_y);
	glClearDepthf(1.0f);
	glClearColor(0.2f, 0.2f, 0.2f, 1.0f);
	glEnable(GL_DEPTH_TEST);

	size_t fpsSamples = 0;
	double lastTime = GetTimeSeconds();
	double fpsNextTick = lastTime + 1.0;

	bool show_logs = true;
	bool show_gui = true;

	while (!glfwWindowShouldClose(window)) {
		//
		// Compute timings information
		//
		double const nowTime = GetTimeSeconds();
		double const delta_time = nowTime - lastTime;
		lastTime = nowTime;
		if (nowTime > fpsNextTick) {
			fpsNextTick += 1.0;
			fpsSamples = 0;
		}
		++fpsSamples;

		//
		// Process inputs
		//
		glfwPollEvents();

		ImGuiIO const& io = ImGui::GetIO();
		input_handler.SetUICapture(io.WantCaptureMouse, io.WantCaptureKeyboard);
		input_handler.Advance();
		camera.Update(delta_time, input_handler);

		if (input_handler.GetKeycodeState(GLFW_KEY_F3) & JUST_RELEASED)
			show_logs = !show_logs;
		if (input_handler.GetKeycodeState(GLFW_KEY_F2) & JUST_RELEASED)
			show_gui = !show_gui;

		//
		// Start a new frame for Dear ImGui
		//
		ImGui_ImplGlfwGL3_NewFrame();

		//
		// Clear the screen
		//
		glClear(GL_DEPTH_BUFFER_BIT | GL_COLOR_BUFFER_BIT);

		//
		// Update the transforms
		//
		sun_node.rotate_y(sun_spin_speed * delta_time);

		//
		// Traverse the scene graph and render all nodes
		//
		std::stack<Node const*> node_stack({ &solar_system_node });
		std::stack<glm::mat4> matrix_stack({ glm::mat4(1.0f) });
		// TODO: Replace this explicit rendering of the Sun with a
		// traversal of the scene graph and rendering of all its nodes.
		sun_node.render(camera.GetWorldToClipMatrix(), sun_node.get_transform(), shader, [](GLuint /*program*/){});

		//
		// Display Dear ImGui windows
		//
		if (show_logs)
			Log::View::Render();
		if (show_gui)
			ImGui::Render();

		//
		// Queue the computed frame for display on screen
		//
		glfwSwapBuffers(window);
	}

	glDeleteTextures(1, &sun_texture);

	Log::View::Destroy();
	Log::Destroy();

	return EXIT_SUCCESS;
}
<commit_msg>Use the correct variable name for the diffuse texture<commit_after>#include "config.hpp"
#include "core/FPSCamera.h"
#include "core/helpers.hpp"
#include "core/Log.h"
#include "core/LogView.h"
#include "core/Misc.h"
#include "core/node.hpp"
#include "core/ShaderProgramManager.hpp"
#include "core/WindowManager.hpp"

#include <imgui.h>
#include <external/imgui_impl_glfw_gl3.h>

#include <stack>

#include <cstdlib>

int main()
{
	//
	// Set up the logging system
	//
	Log::Init();
	Log::View::Init();

	//
	// Set up the camera
	//
	InputHandler input_handler;
	FPSCameraf camera(0.5f * glm::half_pi<float>(),
	                  static_cast<float>(config::resolution_x) / static_cast<float>(config::resolution_y),
	                  0.01f, 1000.0f);
	camera.mWorld.SetTranslate(glm::vec3(0.0f, 0.0f, 6.0f));
	camera.mMouseSensitivity = 0.003f;
	camera.mMovementSpeed = 0.25f * 12.0f;

	//
	// Set up the windowing system and create the window
	//
	WindowManager window_manager;
	WindowManager::WindowDatum window_datum{ input_handler, camera, config::resolution_x, config::resolution_y, 0, 0, 0, 0};
	GLFWwindow* window = window_manager.CreateWindow("EDAF80: Assignment 1", window_datum, config::msaa_rate);
	if (window == nullptr) {
		LogError("Failed to get a window: exiting.");

		Log::View::Destroy();
		Log::Destroy();

		return EXIT_FAILURE;
	}

	//
	// Load the sphere geometry
	//
	std::vector<bonobo::mesh_data> const objects = bonobo::loadObjects("sphere.obj");
	if (objects.empty()) {
		LogError("Failed to load the sphere geometry: exiting.");

		Log::View::Destroy();
		Log::Destroy();

		return EXIT_FAILURE;
	}
	bonobo::mesh_data const& sphere = objects.front();

	//
	// Create the shader program
	//
	ShaderProgramManager program_manager;
	GLuint shader = 0u;
	program_manager.CreateAndRegisterProgram({ { ShaderType::vertex, "EDAF80/default.vert" },
	                                           { ShaderType::fragment, "EDAF80/default.frag" } },
	                                         shader);
	if (shader == 0u) {
		LogError("Failed to generate the shader program: exiting.");

		Log::View::Destroy();
		Log::Destroy();

		return EXIT_FAILURE;
	}

	//
	// Set up the sun node and other related attributes
	//
	Node sun_node;
	sun_node.set_geometry(sphere);
	GLuint const sun_texture = bonobo::loadTexture2D("sunmap.png");
	sun_node.add_texture("diffuse_texture", sun_texture, GL_TEXTURE_2D);
	float const sun_spin_speed = glm::two_pi<float>() / 6.0f; // Full rotation in six seconds

	Node solar_system_node;
	solar_system_node.add_child(&sun_node);

	//
	// TODO: Create nodes for the remaining of the solar system
	//

	glViewport(0, 0, config::resolution_x, config::resolution_y);
	glClearDepthf(1.0f);
	glClearColor(0.2f, 0.2f, 0.2f, 1.0f);
	glEnable(GL_DEPTH_TEST);

	size_t fpsSamples = 0;
	double lastTime = GetTimeSeconds();
	double fpsNextTick = lastTime + 1.0;

	bool show_logs = true;
	bool show_gui = true;

	while (!glfwWindowShouldClose(window)) {
		//
		// Compute timings information
		//
		double const nowTime = GetTimeSeconds();
		double const delta_time = nowTime - lastTime;
		lastTime = nowTime;
		if (nowTime > fpsNextTick) {
			fpsNextTick += 1.0;
			fpsSamples = 0;
		}
		++fpsSamples;

		//
		// Process inputs
		//
		glfwPollEvents();

		ImGuiIO const& io = ImGui::GetIO();
		input_handler.SetUICapture(io.WantCaptureMouse, io.WantCaptureKeyboard);
		input_handler.Advance();
		camera.Update(delta_time, input_handler);

		if (input_handler.GetKeycodeState(GLFW_KEY_F3) & JUST_RELEASED)
			show_logs = !show_logs;
		if (input_handler.GetKeycodeState(GLFW_KEY_F2) & JUST_RELEASED)
			show_gui = !show_gui;

		//
		// Start a new frame for Dear ImGui
		//
		ImGui_ImplGlfwGL3_NewFrame();

		//
		// Clear the screen
		//
		glClear(GL_DEPTH_BUFFER_BIT | GL_COLOR_BUFFER_BIT);

		//
		// Update the transforms
		//
		sun_node.rotate_y(sun_spin_speed * delta_time);

		//
		// Traverse the scene graph and render all nodes
		//
		std::stack<Node const*> node_stack({ &solar_system_node });
		std::stack<glm::mat4> matrix_stack({ glm::mat4(1.0f) });
		// TODO: Replace this explicit rendering of the Sun with a
		// traversal of the scene graph and rendering of all its nodes.
		sun_node.render(camera.GetWorldToClipMatrix(), sun_node.get_transform(), shader, [](GLuint /*program*/){});

		//
		// Display Dear ImGui windows
		//
		if (show_logs)
			Log::View::Render();
		if (show_gui)
			ImGui::Render();

		//
		// Queue the computed frame for display on screen
		//
		glfwSwapBuffers(window);
	}

	glDeleteTextures(1, &sun_texture);

	Log::View::Destroy();
	Log::Destroy();

	return EXIT_SUCCESS;
}
<|endoftext|>


<commit_before>#pragma once
#include <sstream>
#include <imgui.h>

namespace GUI {
    class LogView {
        public:
        /// Constructor.
        LogView();

        /// Destructor.
        ~LogView();

        /// Shows the log.
        void Show();

        /// If the log is visible or not.
        /**
         * @return Whether the log is visible or not.
         */
        bool IsVisible();

        private:
        std::stringstream defaultStringstream;
        std::stringstream infoStringstream;
        std::stringstream warningStringstream;
        std::stringstream errorStringstream;

        ImGuiTextBuffer textBuffer;
        ImGuiTextFilter textFilter;
        ImVector<int> lineOffsets;

        bool visible = true;

        static const int splitterSize = 2;
        int logHeight = 250;
        bool logResize = false;

        int sceneWidth = 250;
        bool sceneResize = false;

        int editorWidth = 250;
        bool editorResize = false;
    };
} // namespace GUI
<commit_msg>Update LogView.hpp<commit_after>#pragma once
#include <sstream>
#include <imgui.h>

namespace GUI {
    class LogView {
        public:
        /// Constructor.
        LogView();

        /// Destructor.
        ~LogView();

        /// Shows the log.
        void Show();

        /// If the log is visible or not.
        /**
         * @return Whether the log is visible or not.
         */
        const bool IsVisible();

        private:
        std::stringstream defaultStringstream;
        std::stringstream infoStringstream;
        std::stringstream warningStringstream;
        std::stringstream errorStringstream;

        ImGuiTextBuffer textBuffer;
        ImGuiTextFilter textFilter;
        ImVector<int> lineOffsets;

        bool visible = true;

        static const int splitterSize = 2;
        int logHeight = 250;
        bool logResize = false;

        int sceneWidth = 250;
        bool sceneResize = false;

        int editorWidth = 250;
        bool editorResize = false;
    };
} // namespace GUI
<|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/compiler/mlir/tensorflow/analysis/resource_dataflow.h"

#include <algorithm>
#include <vector>

#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"  // from @llvm-project
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"  // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/SymbolTable.h"  // from @llvm-project
#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

#define DEBUG_TYPE "resource-dataflow"

namespace mlir {
namespace TF {

ResourceConstructingOps::ResourceConstructingOps(Operation *op) {
  if (op) ops.insert(op);
}

ResourceConstructingOps ResourceConstructingOps::getPessimisticValueState(
    MLIRContext *context) {
  return ResourceConstructingOps();
}
ResourceConstructingOps ResourceConstructingOps::getPessimisticValueState(
    Value value) {
  if (auto barg = value.dyn_cast<BlockArgument>()) {
    if (func::FuncOp func =
            dyn_cast<func::FuncOp>(barg.getOwner()->getParentOp())) {
      SymbolTable symbol_table(func->getParentOfType<ModuleOp>());
      auto global_tensor = tf_saved_model::LookupBoundInputOfType<
          tf_saved_model::GlobalTensorOp>(func, barg.getArgNumber(),
                                          symbol_table);
      return ResourceConstructingOps(global_tensor);
    }
  } else if (auto vh = dyn_cast<TF::VarHandleOp>(value.getDefiningOp())) {
    return ResourceConstructingOps(vh);
  }
  return ResourceConstructingOps();
}

ResourceConstructingOps ResourceConstructingOps::join(
    const ResourceConstructingOps &lhs, const ResourceConstructingOps &rhs) {
  // Take union of both sets of possible GlobalTensorOp values that can be
  // referenced here.
  ResourceConstructingOps ret;
  ret.ops.insert(lhs.ops.begin(), lhs.ops.end());
  ret.ops.insert(rhs.ops.begin(), rhs.ops.end());
  return ret;
}

void ResourceConstructingOps::print(raw_ostream &os) const {
  llvm::interleaveComma(ops, os << "["), os << "]";
}

void ResourceDataflowAnalysis::visitOperation(Operation *op,
                                              ArrayRef<const StateT *> operands,
                                              ArrayRef<StateT *> results) {
  LLVM_DEBUG(llvm::dbgs() << "ResAn: Visiting operation: " << *op << "\n");
  setAllToEntryStates(results);
}

}  // namespace TF
}  // namespace mlir
<commit_msg>Make resource dataflow aware of TF_CastOp and TF_WhileRegionOp.<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/compiler/mlir/tensorflow/analysis/resource_dataflow.h"

#include <algorithm>
#include <vector>

#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"  // from @llvm-project
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"  // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/SymbolTable.h"  // from @llvm-project
#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

#define DEBUG_TYPE "resource-dataflow"

namespace mlir {
namespace TF {

ResourceConstructingOps::ResourceConstructingOps(Operation *op) {
  if (op) ops.insert(op);
}

ResourceConstructingOps ResourceConstructingOps::getPessimisticValueState(
    MLIRContext *context) {
  return ResourceConstructingOps();
}
ResourceConstructingOps ResourceConstructingOps::getPessimisticValueState(
    Value value) {
  if (auto barg = value.dyn_cast<BlockArgument>()) {
    if (func::FuncOp func =
            dyn_cast<func::FuncOp>(barg.getOwner()->getParentOp())) {
      SymbolTable symbol_table(func->getParentOfType<ModuleOp>());
      auto global_tensor = tf_saved_model::LookupBoundInputOfType<
          tf_saved_model::GlobalTensorOp>(func, barg.getArgNumber(),
                                          symbol_table);
      return ResourceConstructingOps(global_tensor);
    }
  } else if (auto vh = dyn_cast<TF::VarHandleOp>(value.getDefiningOp())) {
    return ResourceConstructingOps(vh);
  }
  return ResourceConstructingOps();
}

ResourceConstructingOps ResourceConstructingOps::join(
    const ResourceConstructingOps &lhs, const ResourceConstructingOps &rhs) {
  // Take union of both sets of possible GlobalTensorOp values that can be
  // referenced here.
  ResourceConstructingOps ret;
  ret.ops.insert(lhs.ops.begin(), lhs.ops.end());
  ret.ops.insert(rhs.ops.begin(), rhs.ops.end());
  return ret;
}

void ResourceConstructingOps::print(raw_ostream &os) const {
  llvm::interleaveComma(ops, os << "["), os << "]";
}

void ResourceDataflowAnalysis::visitOperation(Operation *op,
                                              ArrayRef<const StateT *> operands,
                                              ArrayRef<StateT *> results) {
  LLVM_DEBUG(llvm::dbgs() << "ResAn: Visiting operation: " << *op << "\n");
  if (auto cast = dyn_cast<TF::CastOp>(op)) {
    join(results[0], *operands[0]);
  } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
    for (auto &region : while_op->getRegions()) {
      for (auto [arg, value] :
           llvm::zip(region.getArguments(), while_op->getOperands())) {
        join(getLatticeElement(arg), *getLatticeElement(value));
      }
    }
  } else {
    setAllToEntryStates(results);
  }
}

}  // namespace TF
}  // namespace mlir
<|endoftext|>


<commit_before>/* Copyright 2019 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/kernels/data/experimental/lmdb_dataset_op.h"

#include "tensorflow/core/kernels/data/dataset_test_base.h"

#include "tensorflow/core/platform/posix/posix_file_system.h"

namespace tensorflow {
namespace data {
namespace experimental {
namespace {

constexpr char kNodeName[] = "lmdb_dataset";
constexpr char kIteratorPrefix[] = "Iterator";
constexpr char kDataFileName[] = "data.mdb";
constexpr char kDataFileLoc[] = "tensorflow/core/lib/lmdb/testdata";

class LMDBDatasetParams : public DatasetParams {
 public:
  LMDBDatasetParams(std::vector<std::string> filenames,
                    DataTypeVector output_dtypes,
                    std::vector<PartialTensorShape> output_shapes,
                    string node_name)
      : DatasetParams(std::move(output_dtypes), std::move(output_shapes),
                      std::move(kNodeName)),
        filenames(CreateTensor<std::string>(
            TensorShape({static_cast<int>(filenames.size())}), filenames)) {}

  Status MakeInputs(gtl::InlinedVector<TensorValue, 4>* inputs) override {
    *inputs = {TensorValue(&filenames)};
    return Status::OK();
  }

 private:
  // Names of binary database files to read, boxed up inside a Tensor of
  // strings
  Tensor filenames;
};

class LMDBDatasetOpTest : public DatasetOpsTestBaseV2<LMDBDatasetParams> {
 public:
  Status Initialize(LMDBDatasetParams* dataset_params) override {
    // Step 1: Set up enough of a TF runtime to be able to invoke a kernel.
    TF_RETURN_IF_ERROR(InitThreadPool(thread_num_));
    TF_RETURN_IF_ERROR(InitFunctionLibraryRuntime({}, cpu_num_));

    // Step 2: Box up the inputs to the kernel inside TensorValue objects inside
    // a vector.
    gtl::InlinedVector<TensorValue, 4> inputs;
    TF_RETURN_IF_ERROR(dataset_params->MakeInputs(&inputs));

    // Step 3: Create a dataset kernel to test, passing in attributes of the
    // kernel.
    TF_RETURN_IF_ERROR(MakeDatasetOpKernel(*dataset_params, &dataset_kernel_));

    // Step 4: Create a context in which the kernel will operate. This is where
    // the kernel gets initialized with its inputs
    TF_RETURN_IF_ERROR(
        CreateDatasetContext(dataset_kernel_.get(), &inputs, &dataset_ctx_));

    // Step 5: Unbox the DatasetBase object inside the variant tensor backing
    // the kernel.
    TF_RETURN_IF_ERROR(
        CreateDataset(dataset_kernel_.get(), dataset_ctx_.get(), &dataset_));

    // Step 6: Create an iterator in case the test needs to read the output of
    // the dataset.
    TF_RETURN_IF_ERROR(
        CreateIteratorContext(dataset_ctx_.get(), &iterator_ctx_));
    TF_RETURN_IF_ERROR(dataset_->MakeIterator(iterator_ctx_.get(),
                                              kIteratorPrefix, &iterator_));

    return Status::OK();
  }

  // Create instance of the kernel for `LMDBDataset`.
  // Does not initialize or validate the filenames list, because filenames are
  // inputs, not attributes.
  Status MakeDatasetOpKernel(const LMDBDatasetParams& dataset_params,
                             std::unique_ptr<OpKernel>* op_kernel) override {
    NodeDef node_def = test::function::NDef(
        kNodeName, name_utils::OpName(LMDBDatasetOp::kDatasetType),
        // Inputs
        {LMDBDatasetOp::kFileNames},
        // Attributes
        {{LMDBDatasetOp::kOutputTypes, dataset_params.output_dtypes},
         {LMDBDatasetOp::kOutputShapes, dataset_params.output_shapes}});
    TF_RETURN_IF_ERROR(CreateOpKernel(node_def, op_kernel));
    return Status::OK();
  }
};

// Copy our test data file to the current test program's temporary
// directory, and return the full path to the copied file.
// This copying is necessary because LMDB creates lock files adjacent
// to files that it reads.
string MaybeCopyDataFile() {
  string src_loc = io::JoinPath(kDataFileLoc, kDataFileName);
  string dest_loc = io::JoinPath(testing::TmpDir(), kDataFileName);

  PosixFileSystem fs;
  // PosixFileSystem::FileExists currently returns Status::OK() if the file
  // exists and errors::NotFound() if the file doesn't exist. There's no
  // indication in the code or docs about whether other error codes may be
  // added in the future, so we code defensively here.
  Status exists_status = fs.FileExists(dest_loc);
  if (exists_status.code() == error::NOT_FOUND) {
    TF_EXPECT_OK(fs.CopyFile(src_loc, dest_loc));
  } else {
    TF_EXPECT_OK(exists_status);
  }

  return dest_loc;
}

LMDBDatasetParams SingleValidInput() {
  return {/*filenames=*/{MaybeCopyDataFile()},
          /*output_dtypes=*/{DT_STRING, DT_STRING},
          /*output_shapes=*/{PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams TwoValidInputs() {
  return {/*filenames*/ {MaybeCopyDataFile(), MaybeCopyDataFile()},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams EmptyInput() {
  return {/*filenames*/ {},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams InvalidPathAtStart() {
  return {/*filenames*/ {"This is not a valid filename", MaybeCopyDataFile()},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams InvalidPathInMiddle() {
  return {/*filenames*/ {MaybeCopyDataFile(), "This is not a valid filename",
                         MaybeCopyDataFile()},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

std::vector<Tensor> StringsToTensors(const std::vector<string> strs) {
  std::vector<Tensor> ret;
  for (string str : strs) {
    ret.push_back(CreateTensor<string>(TensorShape({}), {str}));
  }
  return ret;
}

// The tensors we expect to see each time we read through the input
// data file.
const std::vector<Tensor> kFileOutput = StringsToTensors({
    // Each call to GetNext() produces two scalar string tensors, but the test
    // harness expects to receive a flat vector
    "0", "a",  //
    "1", "b",  //
    "2", "c",  //
    "3", "d",  //
    "4", "e",  //
    "5", "f",  //
    "6", "g",  //
    "7", "h",  //
    "8", "i",  //
    "9", "j",  //
});

std::vector<GetNextTestCase<LMDBDatasetParams>> GetNextTestCases() {
  // STL vectors don't have a "concatenate two vectors into a new vector"
  // operation, so...
  std::vector<Tensor> output_twice;
  output_twice.insert(output_twice.end(), kFileOutput.cbegin(),
                      kFileOutput.cend());
  output_twice.insert(output_twice.end(), kFileOutput.cbegin(),
                      kFileOutput.cend());

  return {
      {/*dataset_params=*/SingleValidInput(), /*expected_outputs=*/kFileOutput},
      {/*dataset_params=*/TwoValidInputs(), /*expected_outputs=*/output_twice},
      {/*dataset_params=*/EmptyInput(), /*expected_outputs=*/{}}};
}

ITERATOR_GET_NEXT_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                         GetNextTestCases());

TEST_F(LMDBDatasetOpTest, InvalidPathAtStart) {
  auto dataset_params = InvalidPathAtStart();
  TF_ASSERT_OK(Initialize(&dataset_params));

  // Errors about invalid files are only raised when attempting to read data.
  bool end_of_sequence = false;
  std::vector<Tensor> out_tensors;
  std::vector<Tensor> next;

  Status get_next_status =
      iterator_->GetNext(iterator_ctx_.get(), &next, &end_of_sequence);

  EXPECT_EQ(get_next_status.code(), error::INVALID_ARGUMENT);
}

TEST_F(LMDBDatasetOpTest, InvalidPathInMiddle) {
  auto dataset_params = InvalidPathInMiddle();
  TF_ASSERT_OK(Initialize(&dataset_params));

  bool end_of_sequence = false;
  std::vector<Tensor> out_tensors;
  std::vector<Tensor> next;

  // First 10 rows should be ok
  for (int i = 0; i < 10; ++i) {
    TF_ASSERT_OK(
        iterator_->GetNext(iterator_ctx_.get(), &next, &end_of_sequence));
    EXPECT_FALSE(end_of_sequence);
  }

  // Next read operation should raise an error
  Status get_next_status =
      iterator_->GetNext(iterator_ctx_.get(), &next, &end_of_sequence);
  EXPECT_EQ(get_next_status.code(), error::INVALID_ARGUMENT);
}

std::vector<DatasetNodeNameTestCase<LMDBDatasetParams>>
DatasetNodeNameTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_node_name=*/kNodeName}};
}

DATASET_NODE_NAME_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                         DatasetNodeNameTestCases());

std::vector<DatasetTypeStringTestCase<LMDBDatasetParams>>
DatasetTypeStringTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_dataset_type_string=*/name_utils::OpName(
               LMDBDatasetOp::kDatasetType)}};
}

DATASET_TYPE_STRING_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                           DatasetTypeStringTestCases());

std::vector<DatasetOutputDtypesTestCase<LMDBDatasetParams>>
DatasetOutputDtypesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_dtypes=*/{DT_STRING, DT_STRING}}};
}

DATASET_OUTPUT_DTYPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                             DatasetOutputDtypesTestCases());

std::vector<DatasetOutputShapesTestCase<LMDBDatasetParams>>
DatasetOutputShapesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_shapes=*/{PartialTensorShape({}),
                                       PartialTensorShape({})}}};
}

DATASET_OUTPUT_SHAPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                             DatasetOutputShapesTestCases());

std::vector<CardinalityTestCase<LMDBDatasetParams>> CardinalityTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_cardinality=*/kUnknownCardinality}};
}

DATASET_CARDINALITY_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                           CardinalityTestCases());

std::vector<IteratorOutputDtypesTestCase<LMDBDatasetParams>>
IteratorOutputDtypesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_dtypes=*/{DT_STRING, DT_STRING}}};
}

ITERATOR_OUTPUT_DTYPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                              IteratorOutputDtypesTestCases());

std::vector<IteratorOutputShapesTestCase<LMDBDatasetParams>>
IteratorOutputShapesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_shapes=*/{PartialTensorShape({}),
                                       PartialTensorShape({})}}};
}

ITERATOR_OUTPUT_SHAPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                              IteratorOutputShapesTestCases());

std::vector<IteratorPrefixTestCase<LMDBDatasetParams>>
IteratorOutputPrefixTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_iterator_prefix=*/name_utils::IteratorPrefix(
               LMDBDatasetOp::kDatasetType, kIteratorPrefix)}};
}

ITERATOR_PREFIX_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                       IteratorOutputPrefixTestCases());

// No test of save and restore; save/restore functionality is not implemented
// for this dataset.

}  // namespace
}  // namespace experimental
}  // namespace data
}  // namespace tensorflow
<commit_msg>Switch to using CreateTensors function<commit_after>/* Copyright 2019 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/kernels/data/experimental/lmdb_dataset_op.h"

#include "tensorflow/core/kernels/data/dataset_test_base.h"

#include "tensorflow/core/platform/posix/posix_file_system.h"

namespace tensorflow {
namespace data {
namespace experimental {
namespace {

constexpr char kNodeName[] = "lmdb_dataset";
constexpr char kIteratorPrefix[] = "Iterator";
constexpr char kDataFileName[] = "data.mdb";
constexpr char kDataFileLoc[] = "tensorflow/core/lib/lmdb/testdata";

class LMDBDatasetParams : public DatasetParams {
 public:
  LMDBDatasetParams(std::vector<std::string> filenames,
                    DataTypeVector output_dtypes,
                    std::vector<PartialTensorShape> output_shapes,
                    string node_name)
      : DatasetParams(std::move(output_dtypes), std::move(output_shapes),
                      std::move(kNodeName)),
        filenames(CreateTensor<std::string>(
            TensorShape({static_cast<int>(filenames.size())}), filenames)) {}

  Status MakeInputs(gtl::InlinedVector<TensorValue, 4>* inputs) override {
    *inputs = {TensorValue(&filenames)};
    return Status::OK();
  }

 private:
  // Names of binary database files to read, boxed up inside a Tensor of
  // strings
  Tensor filenames;
};

class LMDBDatasetOpTest : public DatasetOpsTestBaseV2<LMDBDatasetParams> {
 public:
  Status Initialize(LMDBDatasetParams* dataset_params) override {
    // Step 1: Set up enough of a TF runtime to be able to invoke a kernel.
    TF_RETURN_IF_ERROR(InitThreadPool(thread_num_));
    TF_RETURN_IF_ERROR(InitFunctionLibraryRuntime({}, cpu_num_));

    // Step 2: Box up the inputs to the kernel inside TensorValue objects inside
    // a vector.
    gtl::InlinedVector<TensorValue, 4> inputs;
    TF_RETURN_IF_ERROR(dataset_params->MakeInputs(&inputs));

    // Step 3: Create a dataset kernel to test, passing in attributes of the
    // kernel.
    TF_RETURN_IF_ERROR(MakeDatasetOpKernel(*dataset_params, &dataset_kernel_));

    // Step 4: Create a context in which the kernel will operate. This is where
    // the kernel gets initialized with its inputs
    TF_RETURN_IF_ERROR(
        CreateDatasetContext(dataset_kernel_.get(), &inputs, &dataset_ctx_));

    // Step 5: Unbox the DatasetBase object inside the variant tensor backing
    // the kernel.
    TF_RETURN_IF_ERROR(
        CreateDataset(dataset_kernel_.get(), dataset_ctx_.get(), &dataset_));

    // Step 6: Create an iterator in case the test needs to read the output of
    // the dataset.
    TF_RETURN_IF_ERROR(
        CreateIteratorContext(dataset_ctx_.get(), &iterator_ctx_));
    TF_RETURN_IF_ERROR(dataset_->MakeIterator(iterator_ctx_.get(),
                                              kIteratorPrefix, &iterator_));

    return Status::OK();
  }

  // Create instance of the kernel for `LMDBDataset`.
  // Does not initialize or validate the filenames list, because filenames are
  // inputs, not attributes.
  Status MakeDatasetOpKernel(const LMDBDatasetParams& dataset_params,
                             std::unique_ptr<OpKernel>* op_kernel) override {
    NodeDef node_def = test::function::NDef(
        kNodeName, name_utils::OpName(LMDBDatasetOp::kDatasetType),
        // Inputs
        {LMDBDatasetOp::kFileNames},
        // Attributes
        {{LMDBDatasetOp::kOutputTypes, dataset_params.output_dtypes},
         {LMDBDatasetOp::kOutputShapes, dataset_params.output_shapes}});
    TF_RETURN_IF_ERROR(CreateOpKernel(node_def, op_kernel));
    return Status::OK();
  }
};

// Copy our test data file to the current test program's temporary
// directory, and return the full path to the copied file.
// This copying is necessary because LMDB creates lock files adjacent
// to files that it reads.
string MaybeCopyDataFile() {
  string src_loc = io::JoinPath(kDataFileLoc, kDataFileName);
  string dest_loc = io::JoinPath(testing::TmpDir(), kDataFileName);

  PosixFileSystem fs;
  // PosixFileSystem::FileExists currently returns Status::OK() if the file
  // exists and errors::NotFound() if the file doesn't exist. There's no
  // indication in the code or docs about whether other error codes may be
  // added in the future, so we code defensively here.
  Status exists_status = fs.FileExists(dest_loc);
  if (exists_status.code() == error::NOT_FOUND) {
    TF_EXPECT_OK(fs.CopyFile(src_loc, dest_loc));
  } else {
    TF_EXPECT_OK(exists_status);
  }

  return dest_loc;
}

LMDBDatasetParams SingleValidInput() {
  return {/*filenames=*/{MaybeCopyDataFile()},
          /*output_dtypes=*/{DT_STRING, DT_STRING},
          /*output_shapes=*/{PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams TwoValidInputs() {
  return {/*filenames*/ {MaybeCopyDataFile(), MaybeCopyDataFile()},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams EmptyInput() {
  return {/*filenames*/ {},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams InvalidPathAtStart() {
  return {/*filenames*/ {"This is not a valid filename", MaybeCopyDataFile()},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

LMDBDatasetParams InvalidPathInMiddle() {
  return {/*filenames*/ {MaybeCopyDataFile(), "This is not a valid filename",
                         MaybeCopyDataFile()},
          /*output_dtypes*/ {DT_STRING, DT_STRING},
          /*output_shapes*/ {PartialTensorShape({}), PartialTensorShape({})},
          /*node_name=*/kNodeName};
}

// The tensors we expect to see each time we read through the input data file.
const std::vector<Tensor> kFileOutput = CreateTensors<string>(
    TensorShape({}),
    {
        // Each call to GetNext() produces two scalar string tensors, but the
        // test harness expects to receive a flat vector
        {"0"}, {"a"},  //
        {"1"}, {"b"},  //
        {"2"}, {"c"},  //
        {"3"}, {"d"},  //
        {"4"}, {"e"},  //
        {"5"}, {"f"},  //
        {"6"}, {"g"},  //
        {"7"}, {"h"},  //
        {"8"}, {"i"},  //
        {"9"}, {"j"},  //
    });

std::vector<GetNextTestCase<LMDBDatasetParams>> GetNextTestCases() {
  // STL vectors don't have a "concatenate two vectors into a new vector"
  // operation, so...
  std::vector<Tensor> output_twice;
  output_twice.insert(output_twice.end(), kFileOutput.cbegin(),
                      kFileOutput.cend());
  output_twice.insert(output_twice.end(), kFileOutput.cbegin(),
                      kFileOutput.cend());

  return {
      {/*dataset_params=*/SingleValidInput(), /*expected_outputs=*/kFileOutput},
      {/*dataset_params=*/TwoValidInputs(), /*expected_outputs=*/output_twice},
      {/*dataset_params=*/EmptyInput(), /*expected_outputs=*/{}}};
}

ITERATOR_GET_NEXT_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                         GetNextTestCases());

TEST_F(LMDBDatasetOpTest, InvalidPathAtStart) {
  auto dataset_params = InvalidPathAtStart();
  TF_ASSERT_OK(Initialize(&dataset_params));

  // Errors about invalid files are only raised when attempting to read data.
  bool end_of_sequence = false;
  std::vector<Tensor> out_tensors;
  std::vector<Tensor> next;

  Status get_next_status =
      iterator_->GetNext(iterator_ctx_.get(), &next, &end_of_sequence);

  EXPECT_EQ(get_next_status.code(), error::INVALID_ARGUMENT);
}

TEST_F(LMDBDatasetOpTest, InvalidPathInMiddle) {
  auto dataset_params = InvalidPathInMiddle();
  TF_ASSERT_OK(Initialize(&dataset_params));

  bool end_of_sequence = false;
  std::vector<Tensor> out_tensors;
  std::vector<Tensor> next;

  // First 10 rows should be ok
  for (int i = 0; i < 10; ++i) {
    TF_ASSERT_OK(
        iterator_->GetNext(iterator_ctx_.get(), &next, &end_of_sequence));
    EXPECT_FALSE(end_of_sequence);
  }

  // Next read operation should raise an error
  Status get_next_status =
      iterator_->GetNext(iterator_ctx_.get(), &next, &end_of_sequence);
  EXPECT_EQ(get_next_status.code(), error::INVALID_ARGUMENT);
}

std::vector<DatasetNodeNameTestCase<LMDBDatasetParams>>
DatasetNodeNameTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_node_name=*/kNodeName}};
}

DATASET_NODE_NAME_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                         DatasetNodeNameTestCases());

std::vector<DatasetTypeStringTestCase<LMDBDatasetParams>>
DatasetTypeStringTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_dataset_type_string=*/name_utils::OpName(
               LMDBDatasetOp::kDatasetType)}};
}

DATASET_TYPE_STRING_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                           DatasetTypeStringTestCases());

std::vector<DatasetOutputDtypesTestCase<LMDBDatasetParams>>
DatasetOutputDtypesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_dtypes=*/{DT_STRING, DT_STRING}}};
}

DATASET_OUTPUT_DTYPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                             DatasetOutputDtypesTestCases());

std::vector<DatasetOutputShapesTestCase<LMDBDatasetParams>>
DatasetOutputShapesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_shapes=*/{PartialTensorShape({}),
                                       PartialTensorShape({})}}};
}

DATASET_OUTPUT_SHAPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                             DatasetOutputShapesTestCases());

std::vector<CardinalityTestCase<LMDBDatasetParams>> CardinalityTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_cardinality=*/kUnknownCardinality}};
}

DATASET_CARDINALITY_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                           CardinalityTestCases());

std::vector<IteratorOutputDtypesTestCase<LMDBDatasetParams>>
IteratorOutputDtypesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_dtypes=*/{DT_STRING, DT_STRING}}};
}

ITERATOR_OUTPUT_DTYPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                              IteratorOutputDtypesTestCases());

std::vector<IteratorOutputShapesTestCase<LMDBDatasetParams>>
IteratorOutputShapesTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_output_shapes=*/{PartialTensorShape({}),
                                       PartialTensorShape({})}}};
}

ITERATOR_OUTPUT_SHAPES_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                              IteratorOutputShapesTestCases());

std::vector<IteratorPrefixTestCase<LMDBDatasetParams>>
IteratorOutputPrefixTestCases() {
  return {{/*dataset_params=*/SingleValidInput(),
           /*expected_iterator_prefix=*/name_utils::IteratorPrefix(
               LMDBDatasetOp::kDatasetType, kIteratorPrefix)}};
}

ITERATOR_PREFIX_TEST_P(LMDBDatasetOpTest, LMDBDatasetParams,
                       IteratorOutputPrefixTestCases());

// No test of save and restore; save/restore functionality is not implemented
// for this dataset.

}  // namespace
}  // namespace experimental
}  // namespace data
}  // namespace tensorflow
<|endoftext|>