FourierTransform.cpp

#include "FourierTransform.h"
#include "GPGPUQuadRenderer.h"
#include "RenderContext.h"
 
#include "Rendering/IEffects.h"
#include "Rendering/IGraphicsDevice.h"
#include "Rendering/ITextures.h"
#include "Rendering/IRenderTargets.h"
#include "Rendering/IBuffers.h"
#include "Rendering/IContext.h"
 
#include "Foundation/Logging/Logger.h"
 
#include <cmath>
#include <complex>
 
using std::complex;
 
namespace
{
  Cogs::Logging::Log logger = Cogs::Logging::getLogger("FourierTransform");
 
  const float twoPi = float(2.0*M_PI);
 
  // Reverse the lower logN bits of i
  inline size_t reverseBits16(size_t i, size_t logN)
  {
    size_t t = i;
    t = ((0xAAAAu & t) >> 1) | ((t << 1) & 0xAAAAu);  // 1010101010101010
    t = ((0xCCCCu & t) >> 2) | ((t << 2) & 0xCCCCu);  // 1100110011001100
    t = ((0xF0F0u & t) >> 4) | ((t << 4) & 0xF0F0u);  // 1111000011110000
    t = ((0xFF00u & t) >> 8) | ((t << 8) & 0xFF00u);  // 1111111100000000
    return t >> (16 - logN);
  }
 
  // Calculate complex twiddle-factor e^{-2\pi n/N}
  inline complex<float> twiddleFactor(const int n, const int N)
  {
    const float arg = (-twoPi*float(n)) / float(N);
    return complex<float>(std::cos(arg), std::sin(arg));
  }
 
  struct Globals{
    int fourierPass;
  };
 
} // of anonymous namespace
 
Cogs::TextureHandle Cogs::FourierTransform2D::createBakedTable(IGraphicsDevice* device, const int log2N)
{
  int N = 1 << log2N;
 
  std::vector<glm::vec4> table(log2N*N);
 
  // First pass with permutations
  for (int k = 0; k < N; k++) {
    size_t iReversedIndex0 = reverseBits16(k &(~1), log2N);
    size_t iReversedIndex1 = reverseBits16(k | 1, log2N);
    bool iEven = (k & 1) == 0;
    complex<float> iW = (iEven ? 1.f : -1.f);
    table[k].x = float(iReversedIndex0);
    table[k].y = float(iReversedIndex1);
    table[k].z = iW.real();
    table[k].w = iW.imag();
  }
  // Subsequent passes
  for (int l = 1; l < log2N; l++) {
    int blockSize = 2 << l;
    int halfBlockSize = 1 << l;
    int halfBlockMask = halfBlockSize - 1;
    for (int k = 0; k < N; k++) {
      int iBlockNumber = k / blockSize;
      int iBlockIndex = k & halfBlockMask;
      bool iEven = (k & halfBlockSize) == 0;
      int iTwiddleIndex = int(iBlockIndex << (log2N - l - 1));
      int iIndex0 = blockSize*iBlockNumber + iBlockIndex;
      int iIndex1 = iIndex0 + halfBlockSize;
      complex<float> iW = (iEven ? 1.f : -1.f)*twiddleFactor(-iTwiddleIndex, N);
      table[N*l + k].x = float(iIndex0);
      table[N*l + k].y = float(iIndex1);
      table[N*l + k].z = iW.real();
      table[N*l + k].w = iW.imag();
    }
  }
 
  return device->getTextures()->loadTexture(reinterpret_cast<const unsigned char*>(table.data()),
                                            N, log2N, TextureFormat::R32G32B32A32_FLOAT, 0);
 
}
 
void Cogs::FourierTransform2D::initialize(IGraphicsDevice* device, GPGPUQuadRenderer& gpgpuQuadRenderer)
{
  this->device = device;
  IEffects* effects = device->getEffects();
 
  PreprocessorDefinitions definitions;
  passRadix2x2 = effects->loadEffect("Terrain/GPGPUPassThroughVS.hlsl",
                                     "Terrain/FourierTransformPS.hlsl",
                                     definitions);
  if (!HandleIsValid(passRadix2x2)) {
    LOG_ERROR(logger, "Error loading passRadix2x2 effect.");
    return;
  }
  VertexFormatHandle handle = gpgpuQuadRenderer.vertexFormat();
  passRadix2x2InputLayout = device->getBuffers()->loadInputLayout(&handle, 1, passRadix2x2);
 
  definitions.push_back({ "DOUBLE", "1" });
  passRadix2x2Double = effects->loadEffect("Terrain/GPGPUPassThroughVS.hlsl",
                                           "Terrain/FourierTransformPS.hlsl",
                                           definitions);
 
  if (!HandleIsValid(passRadix2x2Double)) {
    LOG_ERROR(logger, "Error loading passRadix2x2Double effect.");
    return;
  }
  passRadix2x2DoubleInputLayout = device->getBuffers()->loadInputLayout(&handle, 1, passRadix2x2Double);
 
  constantBuffer = device->getBuffers()->loadBuffer(nullptr, sizeof(int), Usage::Dynamic, AccessMode::Write, BindFlags::ConstantBuffer);
}
 
void Cogs::FourierTransform2D::setSize(int log2N)
{
  this->log2N = log2N;
 
  tableTex = createBakedTable(device, log2N);
 
  ITextures* textures = device->getTextures();
  IRenderTargets* renderTargets = device->getRenderTargets();
 
  for (int k = 0; k < 2; k++) {
    scratchTex[k] = textures->loadTexture(nullptr, 1 << log2N, 1 << log2N, TextureFormat::R32G32_FLOAT, TextureFlags::RenderTarget);
    scratchTarget[k] = renderTargets->createRenderTarget(scratchTex[k]);
 
    scratchDoubleTex[k] = textures->loadTexture(nullptr, 1 << log2N, 1 << log2N, TextureFormat::R32G32B32A32_FLOAT, TextureFlags::RenderTarget);
    scratchDoubleTarget[k] = renderTargets->createRenderTarget(scratchDoubleTex[k]);
  }
}
 
 
void  Cogs::FourierTransform2D::inverseFourierTransform(RenderContext& renderContext,
                                                        GPGPUQuadRenderer& gpgpuQuadRenderer,
                                                        RenderTargetHandle& target,
                                                        TextureHandle& source,
                                                        bool doubleData)
{
  IContext* context = renderContext.context;
 
  context->setViewport(0, 0, float(1 << log2N), float(1 << log2N));
  context->setEffect(doubleData ? passRadix2x2Double : passRadix2x2);
 
  for (int pass = 0; pass < log2N; pass++) {
    context->setRenderTarget(pass + 1 == log2N ? target : (doubleData ? scratchDoubleTarget[pass & 1] : scratchTarget[pass & 1]), DepthStencilHandle::InvalidHandle);
    gpgpuQuadRenderer.bind(context);
    {
      MappedBuffer<Globals> constants(context, constantBuffer, MapMode::WriteDiscard);
      if (constants) {
        constants->fourierPass = pass;
      }
    }
    context->setConstantBuffer("Globals", constantBuffer);
    context->setTexture("bakedTable", 0, tableTex);
    context->setTexture("source", 1, pass == 0 ? source : (doubleData ? scratchDoubleTex[(pass + 1) & 1] : scratchTex[(pass + 1) & 1]));
    context->setInputLayout(doubleData ? passRadix2x2DoubleInputLayout : passRadix2x2InputLayout);
 
    gpgpuQuadRenderer.draw(context);
  }
 
  context->setRenderTarget(RenderTargetHandle::NoHandle, DepthStencilHandle::NoHandle);
}