Как улучшить разделение с плавающей запятой на графических процессорах с поддержкой FMA?

Современные алгоритмы деления на основе FMA обычно в значительной степени зависят от работы Питера Маркштейна, который много работал над этой проблемой. Каноническая ссылка:

Питер Маркштейн, «IA-64 и элементарные функции: скорость и точность», Prentice-Hall 2000.

Приведенный ниже код C также основан на работе Маркштейна. По сути, он состоит из быстрого пути, который обрабатывает распространенные случаи как можно быстрее, и медленного пути, который обрабатывает все надоедливые частные случаи.

Используемые концепции довольно просты, но требуют тщательного математического анализа, чтобы гарантировать правильно округленный результат. Первым шагом является вычисление точной обратной величины делителя. Это требует точного вычисления ошибки аппроксимации, для чего FMA является лучшим инструментом. Уточнение обратной зависимости от (аппаратного) приближения обычно принимает форму одной итерации Ньютона-Рафсона с квадратичной сходимостью или одной итерации Галлея с кубической сходимостью, обе из которых идеально отображаются в FMA.

Умножение обратной величины делителя на делимое дает приближение к частному. Это также уточняется на основе вычисления остатка на основе FMA. На этом этапе обычно работают с версией делимого, уменьшенной до единицы, чтобы защититься от переполнения и потери значимости в промежуточных вычислениях и обеспечить максимально широкий диапазон операндов в быстром пути алгоритма деления. Это также означает, что в конце необходимо выполнить заключительный шаг масштабирования для корректировки величины частного.

Предпочтительное расположение двух путей в коде состоит в том, чтобы сначала безоговорочно выполнить вычисление быстрого пути, затем проверить в конце, выполняются ли условия для использования быстрого пути, и, если нет, вызвать код медленного пути. Это позволяет вычислять необходимые условия одновременно с вычислением быстрого пути и позволяет обрисовывать в общих чертах вычисление медленного пути, что облегчает размещение этого кода на редко используемой странице памяти, где собраны различные коды медленного пути или обработки исключений.

Пожалуйста, относитесь к приведенному ниже коду как к алгоритмическому наброску. Включенный тест уровня «дыма», основанный на случайных тестовых векторах, далек от строгой структуры тестирования. Условия для вызова медленного пути не были ни максимально жесткими, ни настолько жесткими, насколько это необходимо, ни наиболее эффективной компоновкой. Уточнение обратной зависимости использует один шаг Ньютона-Рафсона, но этого может быть недостаточно для данной обратной аппроксимации графического процессора, и может потребоваться замена итерации Галлея за счет дополнительного FMA. Наконец, я пропустил построение всего медленного пути, так как это превысило бы ограничения ответа Stackoverflow как с точки зрения усилий по проектированию, так и с точки зрения количества текстового описания.

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

float uint32_as_float (uint32_t a)
{
    float r;
    memcpy (&r, &a, sizeof r);
    return r;
}

uint32_t float_as_uint32 (float a)
{
    uint32_t r;
    memcpy (&r, &a, sizeof r);
    return r;
}

float rcp_approx (float a)
{
    return 1.0f / a; // fake it for now
}

float fp32_div_slowpath (float dividend, float divisor)
{
    return dividend / divisor; // fake it for now
}

float fp32_div (float dividend, float divisor)
{
    const uint32_t FP32_MANT_MASK = 0x007fffffu;
    const uint32_t FP32_ONE = 0x3f800000u;
    const uint32_t FP32_SIGN_MASK = 0x80000000u;
    const uint32_t FP32_SIGN_EXPO_MASK = 0xff800000u;
    const uint32_t FP32_HI_LIMIT = 0x7e800000u; // 0x1.0p+126
    const uint32_t FP32_LO_LIMIT = 0x00800001u; // 0x1.0p-126 + ulp

    // fast path
    float recip_approx = rcp_approx (divisor); 
    float recip_err = fmaf (recip_approx, -divisor, 1.0f);
    float dividend_mant = uint32_as_float ((float_as_uint32 (dividend) & FP32_MANT_MASK) | FP32_ONE);
    float dividend_scale = uint32_as_float (float_as_uint32 (dividend) & FP32_SIGN_EXPO_MASK);
    float refined_recip = fmaf (recip_approx, recip_err, recip_approx);
    float quotient_mant = refined_recip * dividend_mant;
    float quotient_residual = fmaf (quotient_mant, -divisor, dividend_mant);
    float refined_quotient_mant = fmaf (refined_recip, quotient_residual, quotient_mant);
    float refined_quotient_residual = fmaf (refined_quotient_mant, -divisor, dividend_mant);
    float final_quotient_mant = fmaf (refined_recip, refined_quotient_residual, refined_quotient_mant);
    float final_quotient = final_quotient_mant * dividend_scale;

    // check if we need to apply the slow path and invoke it if that is the case
    uint32_t id = float_as_uint32 (divisor) & ~FP32_SIGN_MASK;
    uint32_t iq = float_as_uint32 (final_quotient) & ~FP32_SIGN_MASK;
    if ((id > FP32_HI_LIMIT) || ((iq - FP32_LO_LIMIT) > (FP32_HI_LIMIT - FP32_LO_LIMIT))) {
        final_quotient = fp32_div_slowpath (dividend, divisor);
    }

    return final_quotient;
}

// Fixes via: Greg Rose, KISS: A Bit Too Simple. http://eprint.iacr.org/2011/007
static unsigned int z=362436069,w=521288629,jsr=362436069,jcong=123456789;
#define znew (z=36969*(z&0xffff)+(z>>16))
#define wnew (w=18000*(w&0xffff)+(w>>16))
#define MWC  ((znew<<16)+wnew)
#define SHR3 (jsr^=(jsr<<13),jsr^=(jsr>>17),jsr^=(jsr<<5)) /* 2^32-1 */
#define CONG (jcong=69069*jcong+13579)                     /* 2^32 */
#define KISS ((MWC^CONG)+SHR3)

int main (void)
{
    uint64_t count = 0;
    float divisor, dividend, res, ref;

    do {
        if (!(count & 0xffffffULL)) printf ("\rcount: %llu", count);

        dividend = uint32_as_float (KISS);
        divisor  = uint32_as_float (KISS);
        res = fp32_div (dividend, divisor);
        ref = dividend / divisor;
        if (float_as_uint32 (res) != float_as_uint32 (ref)) {
            printf ("error @ dividend=% 16.6a divisor=% 16.6a (% 15.8e, % 15.8e): res=% 16.6a ref=% 16.6a (% 15.8e, % 15.8e)\n", 
                    dividend, divisor, dividend, divisor, res, ref, res, ref);
            return EXIT_FAILURE;
        }

        count++;
    } while (count);
    return EXIT_SUCCESS;
}

Некоторые дополнительные примечания после повторного ознакомления с проблемой. Предпочтительно, если аппаратная инструкция обратной аппроксимации рассматривает субнормаль как ноль. Это позволяет отказаться от проверки величины делителя при вызове кода медленного пути, поскольку любой делитель больше 2 ¹²⁶ приведет к обратному значению нуля и частному бесконечности, что приведет к запуску проверки отношения.

Что касается проверки отношения, может быть принят любой результат быстрого пути, который представляет собой нормальный результат, за исключением наименьшего нормального результата, который может представлять собой неправильное округление от субнормального результата. Другими словами, частные, абсолютное значение которых находится в [0x1.000002p-126, 0x1.fffffep+127], не должны инициировать повторное вычисление кодом медленного пути.

Нет необходимости в шаге уточнения частного, если используется достаточно точное обратное выражение. Исчерпывающий тест, который охватывает все делимые в [1,2] и все делители в [1,2] и, следовательно, все возможные комбинации шаблонов мантиссы (мантиссы), выполним на современном оборудовании, требующем 2⁴⁶ тестовых случаев. Даже если скалярный код выполняется на одном ядре ЦП, он выполняется менее чем за четыре дня без каких-либо сообщений об ошибках.

В практическом использовании, где это возможно, желательно заставить компилятор встраивать fp32_div(), в то же время принудительно fp_div_slowpath() в вызываемую подпрограмму.

Ниже я прототипировал версию деления с одинарной точностью на основе AVX, используя упрощения, описанные выше. Он прошел все мои тесты. Поскольку точность аппаратной обратной функции на основе AVX низкая, требуется итерация Галлея для уточнения обратной функции до полной одинарной точности, что дает результаты с максимальной ошибкой, близкой к 0,5 ulp.

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "immintrin.h"

#define MODE_RANDOM           (1)
#define MODE_MANT_EXHAUSTIVE  (2)
#define TEST_MODE             (MODE_RANDOM)

float uint32_as_float (uint32_t a) { float r; memcpy (&r, &a, sizeof r); return r; }
uint32_t float_as_uint32 (float a) { uint32_t r; memcpy (&r, &a, sizeof r); return r; }
float fp32_div_slowpath (float dividend, float divisor) { return dividend / divisor; }

/* 12-bit approximation. Subnormal arguments and results are treated as zero */
float rcp_approx_avx (float divisor)
{
    float r;
    __m256 t = _mm256_set1_ps (divisor);
    t = _mm256_rcp_ps (t);
    memcpy (&r, &t, sizeof r);
    return r;
}

float fp32_div (float dividend, float divisor)
{
    const uint32_t FP32_MANT_MASK = 0x007fffffu;
    const uint32_t FP32_ONE = 0x3f800000u;
    const uint32_t FP32_SIGN_MASK = 0x80000000u;
    const uint32_t FP32_SIGN_EXPO_MASK = 0xff800000u;
    const uint32_t FP32_QUOTIENT_HI_LIMIT = 0x7f7fffffu; // 0x1.0p+128 - ulp
    const uint32_t FP32_QUOTIENT_LO_LIMIT = 0x00800001u; // 0x1.0p-126 + ulp

    // fast path
    float recip_approx = rcp_approx_avx (divisor);
    float recip_err = fmaf (recip_approx, -divisor, 1.0f);
    float recip_err2 = fmaf (recip_err, recip_err, recip_err);
    float refined_recip = fmaf (recip_approx, recip_err2, recip_approx);
    float dividend_mant = uint32_as_float ((float_as_uint32 (dividend) & FP32_MANT_MASK) | FP32_ONE);
    float dividend_scale = uint32_as_float (float_as_uint32 (dividend) & FP32_SIGN_EXPO_MASK);
    float refined_quotient_mant = refined_recip * dividend_mant;
    float refined_quotient_residual = fmaf (refined_quotient_mant, -divisor, dividend_mant);
    float final_quotient_mant = fmaf (refined_recip, refined_quotient_residual, refined_quotient_mant);
    float final_quotient = final_quotient_mant * dividend_scale;

    // check if we need to apply the slow path and invoke it if that is the case
    uint32_t iq = float_as_uint32 (final_quotient) & ~FP32_SIGN_MASK;
    if ((iq - FP32_QUOTIENT_LO_LIMIT) > (FP32_QUOTIENT_HI_LIMIT - FP32_QUOTIENT_LO_LIMIT)) {
        final_quotient = fp32_div_slowpath (dividend, divisor);
    }
    return final_quotient;
}

// Fixes via: Greg Rose, KISS: A Bit Too Simple. http://eprint.iacr.org/2011/007
static unsigned int z=362436069,w=521288629,jsr=362436069,jcong=123456789;
#define znew (z=36969*(z&0xffff)+(z>>16))
#define wnew (w=18000*(w&0xffff)+(w>>16))
#define MWC  ((znew<<16)+wnew)
#define SHR3 (jsr^=(jsr<<13),jsr^=(jsr>>17),jsr^=(jsr<<5)) /* 2^32-1 */
#define CONG (jcong=69069*jcong+13579)                     /* 2^32 */
#define KISS ((MWC^CONG)+SHR3)

int main (void)
{
    uint64_t count = 0;
    float divisor, dividend, res, ref;

#if TEST_MODE == MODE_RANDOM
    do {
        if (!(count & 0xffffffULL)) printf ("\rcount: %llu", count);

        dividend = uint32_as_float (KISS);
        divisor  = uint32_as_float (KISS);
        res = fp32_div (dividend, divisor);
        ref = dividend / divisor;
        if (float_as_uint32 (res) != float_as_uint32 (ref)) {
            printf ("error @ dividend=% 16.6a divisor=% 16.6a (% 15.8e, % 15.8e): res=% 16.6a ref=% 16.6a (% 15.8e, % 15.8e)\n", 
                    dividend, divisor,  dividend, divisor, res, ref, res, ref);
            return EXIT_FAILURE;
        }

        count++;
    } while (count);
#else
    for (dividend = 1.0f; dividend <= 2.0f; dividend = uint32_as_float (float_as_uint32 (dividend) + 1)) {
        printf ("\rdividend=%08x", float_as_uint32 (dividend));
        for (divisor = 1.0f; divisor <= 2.0f; divisor = uint32_as_float (float_as_uint32 (divisor) + 1)) {
            res = fp32_div (dividend, divisor);
            ref = dividend / divisor;
            if (float_as_uint32 (res) != float_as_uint32 (ref)) {
                printf ("error @ dividend=% 16.6a divisor=% 16.6a (% 15.8e, % 15.8e): res=% 16.6a ref=% 16.6a (% 15.8e, % 15.8e)\n", 
                        dividend, divisor,  dividend, divisor, res, ref, res, ref);
                return EXIT_FAILURE;
            }
        }
    }
#endif

    return EXIT_SUCCESS;
}

Соответствующий код CUDA (протестированный) для графического процессора NVIDIA выглядит следующим образом:

__noinline__ __device__ float fp32_div_slowpath (float dividend, float divisor) 
{ 
    return dividend / divisor; 
}

/* Subnormal arguments and results are treated as zero */
__forceinline__ __device__ float rcp_approx_gpu (float divisor)
{
    float r;
    asm ("rcp.approx.ftz.f32 %0,%1;\n\t" : "=f"(r) : "f"(divisor));
    return r;
}

__forceinline__ __device__ float fp32_div (float dividend, float divisor)
{
    const uint32_t FP32_MANT_MASK = 0x007fffffu;
    const uint32_t FP32_ONE = 0x3f800000u;
    const uint32_t FP32_SIGN_MASK = 0x80000000u;
    const uint32_t FP32_SIGN_EXPO_MASK = 0xff800000u;
    const uint32_t FP32_QUOTIENT_HI_LIMIT = 0x7f7fffffu; // 0x1.0p+128 - ulp
    const uint32_t FP32_QUOTIENT_LO_LIMIT = 0x00800001u; // 0x1.0p-126 + ulp

    // fast path
    float recip_approx = rcp_approx_gpu (divisor);
    float recip_err = fmaf (recip_approx, -divisor, 1.0f);
    float refined_recip = fmaf (recip_approx, recip_err, recip_approx);
    float dividend_mant = __int_as_float ((__float_as_int (dividend) & FP32_MANT_MASK) | FP32_ONE);
    float dividend_scale = __int_as_float (__float_as_int (dividend) & FP32_SIGN_EXPO_MASK);
    float refined_quotient_mant = refined_recip * dividend_mant;
    float refined_quotient_residual = fmaf (refined_quotient_mant, -divisor, dividend_mant);
    float final_quotient_mant = fmaf (refined_recip, refined_quotient_residual, refined_quotient_mant);
    float final_quotient = final_quotient_mant * dividend_scale;

    // check if we need to apply the slow path and invoke it if that is the case
    uint32_t iq = __float_as_int (final_quotient) & ~FP32_SIGN_MASK;
    if ((iq - FP32_QUOTIENT_LO_LIMIT) > (FP32_QUOTIENT_HI_LIMIT - FP32_QUOTIENT_LO_LIMIT)) {
        final_quotient = fp32_div_slowpath (dividend, divisor);
    }
    return final_quotient;
}