From a16d91d95e0ea1eeb858eac82bddb002b4c9a707 Mon Sep 17 00:00:00 2001 From: bertiet Date: Fri, 2 Sep 2022 09:24:05 +0000 Subject: [PATCH] 2D Extension of 1D Cooley-Tukey FFT. Refactoring original not necesssary as extension makes use of spatial variable separability. Included also is bash script for performing sweeps for data analysis --- python/fft/perf_analysis.py | 7 +- python/fft/sweep2d.sh | 35 ++++++ src/fft/complex.cpp | 50 +++----- src/fft/complex.hpp | 12 +- src/tools/FourierTransform.hpp | 6 +- src/tools/FourierTransform2D.hpp | 202 +++++++++++++++++++++++++++++++ 6 files changed, 262 insertions(+), 50 deletions(-) create mode 100755 python/fft/sweep2d.sh create mode 100644 src/tools/FourierTransform2D.hpp diff --git a/python/fft/perf_analysis.py b/python/fft/perf_analysis.py index 6490a22..c042538 100644 --- a/python/fft/perf_analysis.py +++ b/python/fft/perf_analysis.py @@ -72,7 +72,7 @@ def getFFTInfoFromLog(log_file): if args.csv_out: print("Writing CSV headers") with open(args.csv_out, "w") as f: - f.write(f"Input-size,Batch-size,Radix-size,DFT Batch-size,Cycles,FLOPS Estimate,FLOPS/Cycle,GFLOPS/sec,Memory Including Gaps (bytes),Peak Live Memory Step, Peak Live Memory (bytes)\n") + f.write(f"Input-size,Batch-size,Radix-size,DFT Batch-size,Cycles,FLOPS Estimate,FLOPS/Cycle,GFLOPS/sec,Memory Including Gaps (bytes)\n") exit() else: raise RuntimeError("Can't write headers: no output file specified.") @@ -83,9 +83,6 @@ def getFFTInfoFromLog(log_file): total_memory = sum(tile.memory.total.includingGaps for tile in report.compilation.tiles) print(f"Total memory use (bytes): {total_memory}") - peak_name, peak_live_memory = getPeakLivenessStep(report) - print(f"Program step consuming peak memory: {peak_name} {peak_live_memory}") - size, bs, radix, cycles, flops, dft_batch_size = getFFTInfoFromLog(args.log_file) flops_per_cycle = flops/cycles if flops else None gflops_per_second = flops_per_cycle * args.clock_speed_ghz if flops_per_cycle else None @@ -101,4 +98,4 @@ def getFFTInfoFromLog(log_file): # Collate everything into one line of CSV and append to file if specififed: if args.csv_out: with open(args.csv_out, "a") as f: - f.write(f"{size},{bs},{radix},{dft_batch_size},{cycles},{flops},{flops_per_cycle},{gflops_per_second},{total_memory},{peak_name},{peak_live_memory}\n") + f.write(f"{size},{bs},{radix},{dft_batch_size},{cycles},{flops},{flops_per_cycle},{gflops_per_second},{total_memory}\n") diff --git a/python/fft/sweep2d.sh b/python/fft/sweep2d.sh new file mode 100755 index 0000000..0359989 --- /dev/null +++ b/python/fft/sweep2d.sh @@ -0,0 +1,35 @@ +#!/bin/bash + +# Copyright (c) 2022 Graphcore Ltd. All rights reserved. + +RUN_DIR="2d_profiles" +mkdir -p $RUN_DIR +CSV_FILE=${RUN_DIR}/"2d_csv_results_2.txt" + +for SIZE in 256 512 1024 2048 +do + for BS in 1 4 8 16 32 64 + do + for RADIX in 2 4 8 16 32 64 128 + do + RUN_NAME="fft_2d_${SIZE}_bs${BS}_radix${RADIX}" + export POPLAR_ENGINE_OPTIONS="{\"autoReport.all\":\"true\", \"autoReport.directory\":\"${RUN_DIR}/${RUN_NAME}\", \"profiler.includeFlopEstimates\":\"true\"}" + echo "Running size: ${SIZE} batch-size: ${BS} radix: ${RADIX}" + mkdir -p ${RUN_DIR}/${RUN_NAME} + ./multi-tool FourierTransform2D --fft-size ${SIZE} --batch-size ${BS} --radix-size ${RADIX} > ${RUN_DIR}/${RUN_NAME}/run_log.txt & + done + wait + done +done + +# Overwrite the CSV file writing new headers: +python3 ../python/fft/perf_analysis.py --report-file fake --log-file fake --csv-out ${CSV_FILE} --csv-write-headers + +# Append all run results to CSV file: +for DIR in $RUN_DIR/* +do + echo Processing path $DIR + LOG_FILE=$DIR/run_log.txt + REPORT_FILE=$DIR/ipu_utils_engine/profile.pop + python3 ../python/fft/perf_analysis.py --report-file ${REPORT_FILE} --log-file ${LOG_FILE} --csv-out ${CSV_FILE} +done diff --git a/src/fft/complex.cpp b/src/fft/complex.cpp index 8ac80ba..dcb5ca4 100644 --- a/src/fft/complex.cpp +++ b/src/fft/complex.cpp @@ -55,13 +55,10 @@ namespace complex { namespace pe = popops::expr; auto complexMulExprRe = pe::Sub(pe::Mul(pe::_1, pe::_2), pe::Mul(pe::_3, pe::_4)); auto complexMulExprIm = pe::Add(pe::Mul(pe::_1, pe::_2), pe::Mul(pe::_3, pe::_4)); - - // Can only do the second expression in-place: - auto tmpReal = popops::map(graph, complexMulExprRe, {real, v.real, imag, v.imag}, - prog, debugPrefix + "/complex_mul_re"); + popops::mapInPlace(graph, complexMulExprRe, {real, v.real, imag, v.imag}, + prog, debugPrefix + "/complex_mul_re"); popops::mapInPlace(graph, complexMulExprIm, {imag, v.real, real, v.imag}, prog, debugPrefix + "/complex_mul_im"); - real = tmpReal; } ComplexTensor multiply(poplar::Graph& graph, @@ -129,17 +126,12 @@ namespace complex { graph.setTileMapping(matrix.real, graph.getTileMapping(matmulMapping)); graph.setTileMapping(matrix.imag, graph.getTileMapping(matmulMapping)); - poplar::OptionFlags matmulOptions; - if (availableMemoryProportion > 0.f) { - matmulOptions.set("availableMemoryProportion", std::to_string(availableMemoryProportion)); - } - poplar::Tensor partial = poplin::matMul(graph, matrix.real, realBatch, prog, - elemType, debugStr + "/real_matmul", matmulOptions); + elemType, debugStr + "/real_matmul"); poplin::matMulAcc(graph, partial, 1.f, matrix.imag, imagBatch, prog, - debugStr + "/imag_matmul", matmulOptions); + debugStr + "/imag_matmul"); // FLOP estimates for matrix multiplies: flopEstimate += 2 * matrix.dim(0) * matrix.dim(1) * realBatch.dim(1) * 2; @@ -219,25 +211,21 @@ namespace complex { auto result_odd = fftSubResult.transpose().slice(batchSize, 2*batchSize, 0); ipu_utils::logger()->debug("Twiddle coeff shape: {} and multiply shape: {}", w.shape(), result_odd.shape()); - // Copy the DFT results to a linear layout if there are enough - // elements for this to make sense (this heuristic is very approximate): - if (result_even.real.numElements() > graph.getTarget().getNumTiles()) { - ipu_utils::logger()->debug("Re-mapping DFT result ({} > {}).", - result_even.real.numElements(), graph.getTarget().getNumTiles()); - auto result_even_remapped = ComplexTensor(graph, result_even.elementType(), result_even.shape(), "dft_even_remapped"); - result_even_remapped.mapLinearly(graph); - prog.add(copy(result_even, result_even_remapped)); - result_even = result_even_remapped; - - auto result_odd_remapped = ComplexTensor(graph, result_even.elementType(), result_even.shape(), "dft_even_remapped"); - result_odd_remapped.mapLinearly(graph); - prog.add(copy(result_odd, result_odd_remapped)); - result_odd = result_odd_remapped; - } + // Copy the DFT results to a linear layout: + auto result_even_remapped = ComplexTensor(graph, result_even.elementType(), result_even.shape(), "dft_even_remapped"); + result_even_remapped.mapLinearly(graph); + prog.add(copy(result_even, result_even_remapped)); + result_even = result_even_remapped; + + auto result_odd_remapped = ComplexTensor(graph, result_even.elementType(), result_even.shape(), "dft_even_remapped"); + result_odd_remapped.mapLinearly(graph); + prog.add(copy(result_odd, result_odd_remapped)); + result_odd = result_odd_remapped; // Element-wise multiply odd components by coefficients: - result_odd.multiplyInPlace(graph, w, prog, "twiddle"); - auto tmp = result_odd; + auto tmp = multiply(graph, w, result_odd, prog, "twiddle"); + //result_odd.multiplyInPlace(graph, w, prog, "twiddle"); + //auto tmp = result_odd; // FLOP estimate for complex multiply: flopEstimate += 6 * tmp.real.numElements(); @@ -266,7 +254,7 @@ namespace complex { ComplexTensor FFTBuilder::inverseFourierMatrices( std::size_t length, poplar::Type elemType) { - const double twoPi_over_length = (2.0L / length) * 3.141592653589793238462643383279502884L; + const float twoPi_over_length = (2.0 / length) * 3.14159265358979323846; std::vector real(length * length, 0.f); std::vector imag(length * length, 0.f); for (std::size_t row = 0; row < length; ++row) { @@ -292,7 +280,7 @@ namespace complex { throw std::logic_error("FFT size must be a multiple of 2."); } auto baseSize = N / 2; - const double s = ((2.0L * (N-1)) / N) * 3.141592653589793238462643383279502884L; + const float s = ((2.0 * (N-1)) / N) * 3.14159265358979323846; std::vector real(baseSize, 0.f); std::vector imag(baseSize, 0.f); diff --git a/src/fft/complex.hpp b/src/fft/complex.hpp index 0be0d3e..834d9a4 100644 --- a/src/fft/complex.hpp +++ b/src/fft/complex.hpp @@ -105,29 +105,23 @@ class FFTBuilder { FFTBuilder(poplar::Graph &graph, poplar::program::Sequence &sequence, const std::string debugName) - : graph(graph), prog(sequence), debugPrefix(debugName), - availableMemoryProportion(-1.f), flopEstimate(0) {} - - /// Set the proportion of memory available for the inner DFT matrix-multiplies. - void setAvailableMemoryProportion(float proportion) { availableMemoryProportion = proportion; } + : graph(graph), prog(sequence), debugPrefix(debugName), flopEstimate(0) {} /// Build the compute graph that applies FFT to the given complex vector. /// The program will be appended to the sequence specified in construction /// of this object. ComplexTensor fft1d(ComplexTensor input, std::size_t radix = 0); - + poplar::program::Sequence& getProgram() { return prog; } std::size_t getFlopEstimate() const { return flopEstimate; } private: - float availableMemoryProportion; - std::size_t flopEstimate; - // Utility functions used in construction of the FFT graph program. ComplexTensor multiplyMatrixByVectorBatch(const ComplexTensor matrix, ComplexTensor vectors); ComplexTensor dft1d(ComplexTensor fourierMatrix, ComplexTensor even, ComplexTensor odd); std::pair splitEvenOdd(ComplexTensor input); ComplexTensor inverseFourierMatrices(std::size_t length, poplar::Type elemType); ComplexTensor twiddleCoefficients(std::size_t N, poplar::Type elemType); + std::size_t flopEstimate; }; } // namespace complex \ No newline at end of file diff --git a/src/tools/FourierTransform.hpp b/src/tools/FourierTransform.hpp index 8482595..da1179b 100644 --- a/src/tools/FourierTransform.hpp +++ b/src/tools/FourierTransform.hpp @@ -43,7 +43,6 @@ struct FourierTransform : input.mapLinearly(graph); ipu_utils::logger()->info("Building FFT of input-size {} batch-size {} radix-size {}", size, batchSize, radixSize); - builder.setAvailableMemoryProportion(availableMemoryProportion); auto output = builder.fft1d(input, radixSize); ipu_utils::logger()->info("FFT estimated FLOP count: {}", builder.getFlopEstimate()); @@ -106,9 +105,7 @@ struct FourierTransform : "Batch size for 1D FFT (i.e. number of input vectors).") ("radix-size", po::value(&radixSize)->default_value(0), "Choose radix size (base case size at which DFT matrix-multiply is performed). The default (0) automatically " - "sets the radix to half the input size (i.e. no FFT recursion).") - ("available-memory-proportion", po::value(&availableMemoryProportion)->default_value(-1.f), - "Set the memory proportion available for the inner DFT matrix multiplies. Default: use the Poplar default."); + "sets the radix to half the input size (i.e. no FFT recursion)."); } void init(const boost::program_options::variables_map& args) override { @@ -129,7 +126,6 @@ struct FourierTransform : std::size_t size; std::size_t batchSize; std::size_t radixSize; - float availableMemoryProportion; std::vector realData; std::vector imagData; }; diff --git a/src/tools/FourierTransform2D.hpp b/src/tools/FourierTransform2D.hpp new file mode 100644 index 0000000..9c52088 --- /dev/null +++ b/src/tools/FourierTransform2D.hpp @@ -0,0 +1,202 @@ +// Copyright (c) 2022 Graphcore Ltd. All rights reserved. + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "../fft/complex.hpp" +#include "../fft/utils.hpp" + +#include + +/// Example computes a 2D Fourier transform using the Cooley-Tukey algorithm for fast +/// Fourier transforms (FFT). Then discrete Fourier transform (DFT) matrix is factorised +/// into a base matrix multiply of some dimension (the radix size) followed by 'twiddles' or +/// 'butterflies' that compute the second linear transformation in the factorisation (without +/// the computational cost of the original large DFT matrix multiply). +struct FourierTransform2D : + public ipu_utils::BuilderInterface, public ToolInterface +{ + FourierTransform2D() {} + virtual ~FourierTransform2D() {} + + void build(poplar::Graph& graph, const poplar::Target&) override { + popops::addCodelets(graph); + poplin::addCodelets(graph); + poplar::program::Sequence prog; + + poplar::program::Sequence fftSeq; + complex::FFTBuilder builder(graph, fftSeq, "fft_builder"); + // keep this. 1d input + auto input = complex::ComplexTensor(graph, poplar::FLOAT, {batchSize, size}, "a"); + //ipu_utils::logger()->info("Shape: {}", input.imag.shape()); + auto inputMatrixX = complex::ComplexTensor(graph, poplar::FLOAT, {batchSize, size, size}, "b"); + + input.mapLinearly(graph); + + ipu_utils::logger()->info("Building FFT of input-size {} batch-size {} radix-size {}", size, batchSize, radixSize); + + // inputMatrix.mapLinearlyy(graph); + + auto graphFunctionOutput = builder.fft1d(input, radixSize); + auto graphFunc = graph.addFunction(builder.getProgram()); + + // Call graph function in a loop: + + auto inputMatrix = complex::ComplexTensor( + inputMatrixX.real.dimShuffle({1,0,2}), + inputMatrixX.imag.dimShuffle({1,0,2}) + ); + + ipu_utils::logger()->info("Input Shape: {}", inputMatrix.imag.shape()); + + + for (int n = 0u; n < size; n++) { + + //slice rows of input matrix + auto inputSliceReal = inputMatrix.real.slice(n, n+1, 0); + auto inputSliceImag = inputMatrix.imag.slice(n, n+1, 0); + auto inputSlice = complex::ComplexTensor(inputSliceReal, inputSliceImag); + // Map input slice within loop so whole matrix isn't remapped each time + inputSlice.mapLinearly(graph); + + auto outputSliceReal = inputMatrix.real.slice(n, n+1, 0); + auto outputSliceImag = inputMatrix.imag.slice(n, n+1, 0); + auto outputSlice = complex::ComplexTensor(outputSliceReal, outputSliceImag); + + // 1. Add program to sequence Copy input slice from matrix into input + // 2. call graph function + // 3. Add program to sequence to copy output to slice of result matrix (input and result + prog.add(poplar::program::Copy(inputSlice.imag, input.imag)); + prog.add(poplar::program::Copy(inputSlice.real, input.real)); + prog.add(poplar::program::Call(graphFunc)); + prog.add(poplar::program::Copy(graphFunctionOutput.imag, outputSlice.imag)); + prog.add(poplar::program::Copy(graphFunctionOutput.real, outputSlice.real)); + } + + auto inputRealTranspose = inputMatrix.real.dimShuffle({2,1,0}); + auto inputImagTranspose = inputMatrix.imag.dimShuffle({2,1,0}); + + for (int n = 0u; n < size; n++) { + + //slice rows of transposed input matrix + auto inputSliceReal = inputRealTranspose.slice(n, n+1, 0); + auto inputSliceImag = inputImagTranspose.slice(n, n+1, 0); + auto inputSlice = complex::ComplexTensor(inputSliceReal, inputSliceImag); + // Map input slice within loop so whole matrix isn't remapped each time + inputSlice.mapLinearly(graph); + + auto outputSliceReal = inputRealTranspose.slice(n, n+1, 0); + auto outputSliceImag = inputImagTranspose.slice(n, n+1, 0); + auto outputSlice = complex::ComplexTensor(outputSliceReal, outputSliceImag); + + prog.add(poplar::program::Copy(inputSlice.imag, input.imag)); + prog.add(poplar::program::Copy(inputSlice.real, input.real)); + prog.add(poplar::program::Call(graphFunc)); + prog.add(poplar::program::Copy(graphFunctionOutput.imag, outputSlice.imag)); + prog.add(poplar::program::Copy(graphFunctionOutput.real, outputSlice.real)); + } + + ipu_utils::logger()->info("FFT estimated FLOP count: {}", builder.getFlopEstimate()); + auto cycleCount = poplar::cycleCount(graph, prog, 0, poplar::SyncType::INTERNAL); + + auto outputMatrix = complex::ComplexTensor( + inputMatrix.real.dimShuffle({1,0,2}), + inputMatrix.imag.dimShuffle({1,0,2})); + + graph.createHostWrite("input_real", inputMatrixX.real); + graph.createHostWrite("input_imag", inputMatrixX.imag); + graph.createHostRead("output_real", outputMatrix.real); + graph.createHostRead("output_imag", outputMatrix.imag); + + graph.createHostRead("cycle_count", cycleCount); + + getPrograms().add("fft", prog); + + } + + void execute(poplar::Engine& engine, const poplar::Device& device) override { + // Create input values and write to the device: + // auto step = 1.f / (size*size); + // Each item in a batch is identical: + auto x = 0u; + for (auto b = 0u; b < batchSize; ++b) { + for (auto i = 0u; i < (size*size); ++i) { + // auto x = i * step; + x += 1; + realData[(b*size*size) + i] = x; + imagData[(b*size*size) + i] = x; + } + } + + ipu_utils::writeTensor(engine, "input_real", realData); + ipu_utils::writeTensor(engine, "input_imag", imagData); + if (size < 8u && batchSize < 4u) { + for (auto b = 0u; b < batchSize; ++b) { + ipu_utils::logger()->info("FFT input[{}] Re:\n{}\n", b, slice(realData, b * size * size, (b + 1) * size * size)); + ipu_utils::logger()->info("FFT input[{}] Im:\n{}\n", b, slice(imagData, b * size * size, (b + 1) * size * size)); + } + } + + ipu_utils::logger()->info("Running program"); + getPrograms().run(engine, "fft"); + + ipu_utils::readTensor(engine, "output_real", realData); + ipu_utils::readTensor(engine, "output_imag", imagData); + + uint64_t cycleCount = 0u; + ipu_utils::readScalar(engine, "cycle_count", cycleCount); + ipu_utils::logger()->info("FFT completed in {} cycles.", cycleCount); + if (size < 8u && batchSize < 4u) { + for (auto b = 0u; b < batchSize; ++b) { + auto sliceOutputReal = slice(realData, b * size * size, (b + 1) * size * size); + auto sliceOutputImag = slice(imagData, b * size * size, (b + 1) * size * size); + ipu_utils::logger()->info("FFT result[{}] Re:\n{}\n", b, sliceOutputReal); + ipu_utils::logger()->info("FFT result[{}] Im:\n{}\n", b, sliceOutputImag); + + } + } + } + + void addToolOptions(boost::program_options::options_description& desc) override { + namespace po = boost::program_options; + desc.add_options() + ("fft-size", po::value(&size)->default_value(1024), + "Size of square input matrix to 2D FFT.") + ("batch-size", po::value(&batchSize)->default_value(1), + "Batch size for FFT (i.e. number of input vectors).") + ("radix-size", po::value(&radixSize)->default_value(0), + "Choose radix size (base case size at which DFT matrix-multiply is performed). The default (0) automatically " + "sets the radix to half the input size (i.e. no FFT recursion)."); + } + + void init(const boost::program_options::variables_map& args) override { + if (size % 2) { + throw std::logic_error("FFT input size must be a multiple of 2."); + } + if (radixSize == 0) { + radixSize = size / 2; + } + realData.resize(size * size * batchSize); + imagData.resize(size * size * batchSize); + } + +private: + std::size_t size; + std::size_t batchSize; + std::size_t radixSize; + std::vector realData; + std::vector imagData; +};