udo_regression.cpp

#include <array>
#include <atomic>
#include <memory>
#include <string_view>
#ifdef UDO_STANDALONE
#include <atomic>
#include <cerrno>
#include <charconv>
#include <chrono>
#include <iostream>
#include <map>
#include <string>
#include <thread>
#include <udo/UDOStandalone.hpp>
#include <fcntl.h>
#include <sched.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#endif
//---------------------------------------------------------------------------
#include <udo/UDOperator.hpp>
//---------------------------------------------------------------------------
using namespace std;
using namespace std::literals::string_view_literals;
//---------------------------------------------------------------------------
/// A tuple this UDO takes as an input
struct Input {
   // The value for x
   double x;
   // The measurement of y that will be fitted
   double y;
};
//---------------------------------------------------------------------------
/// An tuple generated by this UDO
struct Output {
   // The value for the parameter a
   double a;
   // The value for the parameter b
   double b;
   // The value for the parameter c
   double c;
};
//---------------------------------------------------------------------------
/// The linear regression operator. It solves the following problem:
/// y_i = a + bx_i + cx_i^2
/// Determine a, b and c for the given values for x and y while minimizing the
/// sum of the squared errors:
/// Sum_i (a + bx_i + cx_i^2 - y_i)^2
///
/// This can be solved as follows:
///
/// Setting the partial derivatives of the squared errors to 0 results in:
///
/// /                           \   /   \   /          \
/// | Sum 1    Sum x    Sum x^2 |   | a |   | Sum y    |
/// | Sum x    Sum x^2  Sum x^3 | * | b | = | Sum xy   |
/// | Sum x^2  Sum x^3  Sum x^4 |   | c |   | Sum x^2y |
/// \                           /   \   /   \          /
///
/// To calculate a, b, and c we calculate the inverse of the first matrix and
/// multiply it from the left:
///
/// /   \   /                           \-1  /          \
/// | a |   | Sum 1    Sum x    Sum x^2 |    | Sum y    |
/// | b | = | Sum x    Sum x^2  Sum x^3 |  * | Sum xy   |
/// | c |   | Sum x^2  Sum x^3  Sum x^4 |    | Sum x^2y |
/// \   /   \                           /    \          /
///
/// This results in a closed form solution for a, b, and c:
///
/// a = 1 / det(A) * (
///        Sum y (Sum x^2 Sum x^4 - (Sum x^3)^2) +
///        Sum xy (Sum x^2 Sum x^3 - Sum x Sum x^4) +
///        Sum x^2y (Sum x Sum x^3 - (Sum x^2)^2)
///     )
///
/// b = 1 / det(A) * (
///        Sum y (Sum x^2 Sum x^3 - Sum x Sum x^4) +
///        Sum xy (Sum 1 Sum x^4 - (Sum x^2)^2) +
///        Sum x^2y (Sum x Sum x^2 - Sum 1 Sum x^3)
///     )
///
/// c = 1 / det(A) * (
///        Sum y (Sum x Sum x^3 - (Sum x^2)^2) +
///        Sum xy (Sum x Sum x^2 - Sum 1 Sum x^3) +
///        Sum x^2y (Sum 1 Sum x^2 - (Sum x)^2)
///     )
///
/// with det(A) =
///      Sum 1 Sum x^2 Sum X^4
///      + 2 Sum x Sum x^2 Sum x^3
///      - (Sum x^2)^3
///      - Sum 1 (Sum x^3)^2
///      - (Sum x)^2 Sum x^4
///
/// Since everything we calculate are sums, we can trivially parallelize this by
/// letting each thread calculate the partial sums of the values it receives and
/// then sum up all partial sums once at the end. With the partial sums we then
/// determine det(A) and finally a, b, and c.
class LinearRegression : public udo::UDOperator<Input, Output> {
   private:
   /// The partial sums of a thread
   struct alignas(64) PartialSums {
      // The value for Sum 1
      double sum1 = 0.0;
      // The value for Sum x
      double sumx = 0.0;
      // The value for Sum x^2
      double sumx2 = 0.0;
      // The value for Sum x^3
      double sumx3 = 0.0;
      // The value for Sum x^4
      double sumx4 = 0.0;
      // The value for Sum y
      double sumy = 0.0;
      // The value for Sum xy
      double sumxy = 0.0;
      // The value for Sum x^2y
      double sumx2y = 0.0;
   };

   /// The local state of a thread in the regression
   struct RegressionLocalState {
      /// The partial sums
      PartialSums partialSums;
      /// The pointer to the next local state
      RegressionLocalState* next = nullptr;
   };

   /// The list of local states
   atomic<RegressionLocalState*> localStateList = nullptr;
   /// The mutex flag to return the result
   atomic_flag resultMutex = false;

   public:
   /// Consume an input tuple
   void consume(LocalState& rawLocalState, const Input& input) {
      auto*& localState = reinterpret_cast<RegressionLocalState*&>(rawLocalState.data);
      if (!localState) {
         auto newLocalState = make_unique<RegressionLocalState>();
         newLocalState->next = localStateList.load();
         while (!localStateList.compare_exchange_weak(newLocalState->next, newLocalState.get()))
            ;

         localState = newLocalState.get();
         // This will be deallocated in postProduce()
         newLocalState.release();
      }

      double x = input.x;
      double y = input.y;

      auto x2 = x * x;
      auto x3 = x2 * x;
      auto x4 = x2 * x2;
      auto xy = x * y;
      auto x2y = x2 * y;

      auto& sums = localState->partialSums;
      sums.sum1 += 1;
      sums.sumx += x;
      sums.sumx2 += x2;
      sums.sumx3 += x3;
      sums.sumx4 += x4;
      sums.sumy += y;
      sums.sumxy += xy;
      sums.sumx2y += x2y;
   }

   /// Produce the output
   bool postProduce(LocalState& /*localState*/) {
      if (resultMutex.test_and_set())
         return true;

      // Sum up all partial sums from the local states
      PartialSums sums;
      for (auto* localState = localStateList.load(); localState;) {
         unique_ptr<RegressionLocalState> localStatePtr(localState);

         auto& lsums = localStatePtr->partialSums;
         sums.sum1 += lsums.sum1;
         sums.sumx += lsums.sumx;
         sums.sumx2 += lsums.sumx2;
         sums.sumx3 += lsums.sumx3;
         sums.sumx4 += lsums.sumx4;
         sums.sumy += lsums.sumy;
         sums.sumxy += lsums.sumxy;
         sums.sumx2y += lsums.sumx2y;

         localState = localStatePtr->next;
      }

      // clang-format off
      double detInv = 1 / (
         sums.sum1 * sums.sumx2 * sums.sumx4
         + 2 * sums.sumx * sums.sumx2 * sums.sumx3
         - sums.sumx2 * sums.sumx2 * sums.sumx2
         - sums.sum1 * sums.sumx3 * sums.sumx3
         - sums.sumx * sums.sumx * sums.sumx4
      );
      double a = detInv * (
         sums.sumy * (sums.sumx2 * sums.sumx4 - sums.sumx3 * sums.sumx3)
         + sums.sumxy * (sums.sumx2 * sums.sumx3 - sums.sumx * sums.sumx4)
         + sums.sumx2y * (sums.sumx * sums.sumx3 - sums.sumx2 * sums.sumx2)
      );
      double b = detInv * (
         sums.sumy * (sums.sumx2 * sums.sumx3 - sums.sumx * sums.sumx4)
         + sums.sumxy * (sums.sum1 * sums.sumx4 - sums.sumx2 * sums.sumx2)
         + sums.sumx2y * (sums.sumx * sums.sumx2 - sums.sum1 * sums.sumx3)
      );
      double c = detInv * (
         sums.sumy * (sums.sumx * sums.sumx3 - sums.sumx2 * sums.sumx2)
         + sums.sumxy * (sums.sumx * sums.sumx2 - sums.sum1 * sums.sumx3)
         + sums.sumx2y * (sums.sum1 * sums.sumx2 - sums.sumx * sums.sumx)
      );
      // clang-format on

      produceOutputTuple({a, b, c});

      return true;
   }
};
//---------------------------------------------------------------------------
#ifdef UDO_STANDALONE
//---------------------------------------------------------------------------
static size_t getNumThreads()
/// Get the number of available threads
{
   ::cpu_set_t cpuSet = {};
   if (::sched_getaffinity(0, sizeof(cpuSet), &cpuSet) != 0)
      return ~0ull;

   size_t threadCount = CPU_COUNT(&cpuSet);
   return threadCount;
}
//---------------------------------------------------------------------------
int main(int argc, const char** argv) {
   bool argError = false;
   bool benchmark = false;
   string_view inputFileName;

   const char** argIt = argv;
   ++argIt;
   const char** argEnd = argv + argc;
   for (; argIt != argEnd; ++argIt) {
      string_view arg(*argIt);
      if (arg.empty())
         continue;
      if (arg == "--benchmark") {
         benchmark = true;
      } else {
         if (inputFileName.empty()) {
            inputFileName = arg;
         } else {
            argError = true;
            break;
         }
      }
   }

   if (!argError && inputFileName.empty())
      argError = true;

   if (argError) {
      cerr << "Usage: " << argv[0] << " [--benchmark] <input file>" << endl;
      return 2;
   }

   int inputFileFd = ::open(inputFileName.data(), O_RDONLY | O_CLOEXEC);
   if (inputFileFd < 0) {
      cerr << "Failed opening " << inputFileName << ": " << strerror(errno) << endl;
      return 1;
   }

   struct ::stat fileStat{};
   if (::fstat(inputFileFd, &fileStat) < 0) {
      cerr << "stat(" << inputFileName << ") failed: " << strerror(errno) << endl;
      return 1;
   }

   void* inputFilePtr = ::mmap(nullptr, fileStat.st_size, PROT_READ, MAP_PRIVATE, inputFileFd, 0);
   if (inputFilePtr == MAP_FAILED) {
      cerr << "mmap(" << inputFileName << ") failed: " << strerror(errno) << endl;
      return 1;
   }
   ::close(inputFileFd);

   string_view inputFileData(static_cast<const char*>(inputFilePtr), fileStat.st_size);

   // Discard the header line
   inputFileData.remove_prefix(inputFileData.find('\n') + 1);

   static constexpr size_t sizePerThread = 4096 * 4;

   size_t numThreads = getNumThreads();

   vector<vector<Input>> threadInputs(numThreads);
   vector<thread> threads;
   atomic<size_t> currentOffset = 0;

   for (size_t threadId = 0; threadId < numThreads; ++threadId) {
      threads.emplace_back([&, threadId, inputFileData] {
         auto& inputs = threadInputs[threadId];

         while (true) {
            size_t localOffset = currentOffset.load();
            if (localOffset >= inputFileData.size())
               break;

            size_t offsetEnd = localOffset + sizePerThread;

            if (offsetEnd >= inputFileData.size()) {
               offsetEnd = inputFileData.size();
            } else {
               // Go forward until the next newline
               size_t newlineOffset = inputFileData.find('\n', offsetEnd);
               if (newlineOffset == string_view::npos)
                  offsetEnd = inputFileData.size();
               else
                  offsetEnd = newlineOffset + 1;
            }

            if (!currentOffset.compare_exchange_weak(localOffset, offsetEnd))
               continue;

            string_view inputStr = inputFileData.substr(localOffset, offsetEnd - localOffset);

            char strBuffer[64];
            while (!inputStr.empty()) {
               Input in;

               size_t commaPos = inputStr.find(',');
               memcpy(strBuffer, inputStr.data(), commaPos - 1);
               strBuffer[commaPos] = '\0';
               inputStr.remove_prefix(commaPos + 1);
               in.x = strtod(strBuffer, nullptr);

               size_t nlPos = inputStr.find('\n');
               memcpy(strBuffer, inputStr.data(), nlPos - 1);
               strBuffer[nlPos] = '\0';
               inputStr.remove_prefix(nlPos + 1);
               in.y = strtod(strBuffer, nullptr);

               inputs.push_back(in);
            }
         }
      });
   }

   for (auto& thread : threads)
      thread.join();

   ::munmap(inputFilePtr, fileStat.st_size);

   vector<Input> inputs;
   for (auto& threadInput : threadInputs)
      inputs.insert(inputs.end(), threadInput.begin(), threadInput.end());

   vector<Output> outputs(3);

   if (benchmark) {
      for (unsigned i = 0; i < 11; ++i) {
         udo::UDOStandalone<LinearRegression> standalone(numThreads, 10000);
         LinearRegression regression;

         auto start = chrono::steady_clock::now();
         standalone.run(regression, inputs, outputs);
         auto end = chrono::steady_clock::now();
         auto duration_ms = chrono::duration_cast<chrono::nanoseconds>(end - start).count();
         // Don't measure the first run
         if (i > 0)
            cout << duration_ms << '\n';
      }
   } else {
      udo::UDOStandalone<LinearRegression> standalone(numThreads, 10000);
      LinearRegression regression;
      standalone.run(regression, inputs, outputs);

      auto& params = standalone.getOutput()[0];
      cout << "a = " << params.a << '\n';
      cout << "b = " << params.b << '\n';
      cout << "c = " << params.c << '\n';
      cout << "-> y = " << params.a << " + " << params.b << "x" << " + " << params.c << "x^2\n";
   }

   return 0;
}
//---------------------------------------------------------------------------
#endif