TIL_RadeonFilter.C

Go to the documentation of this file.

/*
 * Copyright (c) 2021
 *    Side Effects Software Inc.  All rights reserved.
 *
 * Redistribution and use of Houdini Development Kit samples in source and
 * binary forms, with or without modification, are permitted provided that the
 * following conditions are met:
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. The name of Side Effects Software may not be used to endorse or
 *    promote products derived from this software without specific prior
 *    written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY SIDE EFFECTS SOFTWARE 'AS IS' AND ANY EXPRESS
 * OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN
 * NO EVENT SHALL SIDE EFFECTS SOFTWARE BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
 * OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 * EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include "TIL_RadeonFilter.h"
#include <FS/UT_DSO.h>
#include <UT/UT_EnvControl.h>
#include <UT/UT_ErrorLog.h>
#include <UT/UT_Exit.h>
#include <UT/UT_Debug.h>
#include <UT/UT_StopWatch.h>
#include <UT/UT_ParallelUtil.h>
#include <UT/UT_DSOVersion.h>
#include <UT/UT_WorkArgs.h>
#include <stdio.h>
#include <stdarg.h>

#if defined(RIF_USE_METAL)
    #define BACKEND_TYPE RIF_BACKEND_API_METAL
#else
    #define BACKEND_TYPE RIF_BACKEND_API_OPENCL
#endif

namespace
{
    // Keywords
    static constexpr UT_StringLit       theRadius("radius");
    static constexpr UT_StringLit       theThreshold("threshold");
    static constexpr UT_StringLit       theWeight("weight");
    static constexpr UT_StringLit       theUseHDR("useHDR");
    static constexpr UT_StringLit       theModelPath("modelPath");

    static constexpr UT_StringLit       theNName("N");
    static constexpr UT_StringLit       thePzName("Pz");
    static constexpr UT_StringLit       theDefaultModelPath("./models");
}

// Uncomment this to get more information when running filters
//#define TRACE_CALLS
#if defined(TRACE_CALLS)
    #define DEBUG_PRINT(...) do { \
        UT_WorkBuffer   tmp; \
        tmp.format(__VA_ARGS__); \
        fprintf(stderr, "%s(%d): %s\n", __FILE__, __LINE__, tmp.buffer()); \
    } while (false) \
    /* end macro */
#else
    #define DEBUG_PRINT(...) ((void)0)
#endif

namespace
{
    static void SYS_PRINTF_CHECK_ATTRIBUTE(2, 3)
    error(UT_StringHolder &msg, const char *fmt, ...)
    {
        UT_WorkBuffer   tmp;
        va_list args;
        va_start(args, fmt);
        tmp.vsprintf(fmt, args);
        va_end(args);
        msg = UT_StringHolder(tmp);
        DEBUG_PRINT("{}", msg);
    }

    static constexpr inline bool
    isError(rif_int status)
    {
        return status != RIF_SUCCESS;
    }

#define RIF_CALL(status, msg, call) \
    { status = call; \
      DEBUG_PRINT("CALL: {} = {}", int(status), #call); \
      if (status != RIF_SUCCESS) { \
          error(msg, "RadeonPro Error: %s\n  %s\n", \
                  rifGetErrorCodeString(status), \
                  rifGetErrorCodeDescription(status)); \
          UT_ASSERT(0); \
      } \
    } \
    /* end macro */

#define RIF_CALL_FAIL(status, msg, call) \
    { RIF_CALL(status, msg, call) \
      if (isError(status)) return false; } \
    /* end of macro */

    static rif_context
    getRIFContext(UT_StringHolder &msg, int devnum)
    {
        rif_int status = RIF_SUCCESS;
        int     ndev = 0;
        RIF_CALL(status, msg, rifGetDeviceCount(BACKEND_TYPE, &ndev));
        if (!ndev || isError(status))
            return nullptr;
        DEBUG_PRINT("Radeon {} devices", ndev);
        rif_context     context = nullptr;
        RIF_CALL(status, msg, rifCreateContext(RIF_API_VERSION, BACKEND_TYPE,
                    devnum, nullptr, &context));

        return context;
    };
}

TIL_RadeonFilter::TIL_RadeonFilter(RADEON_FILTER_TYPE type)
    : myFilterType(type)
    , myDeviceNum(0)
    , myRadius(0.01)
    , myThreshold(0)
    , myWeight(0.3)
{
}

bool
TIL_RadeonFilter::setOptions(const UT_Options &o)
{
    if (!o.importOption(theRadius.asRef(), myRadius))
        myRadius = 0.01f;
    if (!o.importOption(theThreshold.asRef(), myThreshold))
        myThreshold = 0;
    if (!o.importOption(theWeight.asRef(), myWeight))
        myWeight = 0.3;
    if (!o.importOption(theModelPath.asRef(), myModelPath))
        myModelPath = theDefaultModelPath.asHolder();

    if (myFilterType == RADEON_FILTER_TYPE::DENOISE)
    {
        // Tell the caller what extra AOVs we might want to read
        myAuxPlaneNames.append(theNName.asHolder());    // Normals
        myAuxPlaneNames.append(thePzName.asHolder());   // Depth
    }

    return true;
}

TIL_RadeonFilter::~TIL_RadeonFilter()
{
}

template <typename DT, typename ST, uint DNC, uint SNC, uint32 ISCALE, bool INVERT>
static void
copyPixels(DT *dest, const ST *src, exint npixels)
{
    float               scale = 1;
    constexpr uint      NC = DNC < SNC ? DNC : SNC;
    if (ISCALE)
        scale = INVERT ? 1.0/ISCALE : ISCALE;
    UTparallelFor(UT_BlockedRange<exint>(0, npixels),
        [&](const UT_BlockedRange<exint> &r)
        {
            exint        n = r.end() - r.begin();
            DT          *d = dest + r.begin() * DNC;
            const ST    *s = src + r.begin() * SNC;
            if (!ISCALE && DNC == SNC)
            {
                std::copy(s, s+SNC*n, d);
            }
            else
            {
                for (exint i = 0; i < n; ++i, d += DNC, s += SNC)
                {
                    if (!ISCALE)
                        std::copy(s, s+NC, d);
                    else
                    {
                        for (int ch = 0; ch < NC; ++ch)
                            d[ch] = s[ch] * scale;
                    }
                }
            }
        }
    );
}

template <typename DT, typename ST, uint32 ISCALE, bool INVERT>
static void
doCopyRaster(void *d, int dnc, const void *s, int snc, exint npix)
{
    if (dnc == 3)
    {
        if (snc == 3)
            copyPixels<DT, ST, 3, 3, ISCALE, INVERT>((DT *)d, (const ST *)s, npix);
        else
            copyPixels<DT, ST, 3, 4, ISCALE, INVERT>((DT *)d, (const ST *)s, npix);
    }
    else
    {
        if (snc == 3)
            copyPixels<DT, ST, 4, 3, ISCALE, INVERT>((DT *)d, (const ST *)s, npix);
        else
            copyPixels<DT, ST, 4, 4, ISCALE, INVERT>((DT *)d, (const ST *)s, npix);
    }
}

static const PXL_Raster &
makeFloatRaster(const PXL_Raster &src, PXL_Raster &storage)
{
    if (src.getFormat() == PXL_FLOAT32)
        return src;

    storage.setPacking(src.getPacking());
    storage.setFormat(PXL_FLOAT32);
    storage.setRes(src.getXres(), src.getYres());
    storage.init();

    float       *f = (float *)storage.getPixels();
    const void  *s = src.getPixels();
    exint        npix = src.getNumPixels();
    int          nchan = PXLpackingComponents(src.getPacking());

    switch (src.getFormat())
    {
        case PXL_INT8:
            doCopyRaster<float, uint8, 0xff, true>(f, nchan, s, nchan, npix);
            break;
        case PXL_INT16:
            doCopyRaster<float, uint16, 0xffff, true>(f, nchan, s, nchan, npix);
            break;
        case PXL_INT32:
            doCopyRaster<float, uint32, 0xffffffff, true>(f, nchan, s, nchan, npix);
            break;
        case PXL_FLOAT16:
            doCopyRaster<float, fpreal16, 0, false>(f, nchan, s, nchan, npix);
            break;
        case PXL_FLOAT32:
            doCopyRaster<float, fpreal32, 0, false>(f, nchan, s, nchan, npix);
            break;
        case PXL_MAX_DATA_FORMAT:
            UT_ASSERT(0);
            break;
    }
    return storage;
}

bool
TIL_RadeonFilter::downloadRaster(PXL_Raster &dest, const RIFImage &src)
{
    rif_int             status;
    rif_image_desc      desc;
    size_t              rsize;

    RIF_CALL_FAIL(status, myErrorString, rifImageGetInfo(*src, RIF_IMAGE_DESC,
                sizeof(desc), &desc, &rsize));
    rif_uchar   *data;
    RIF_CALL_FAIL(status, myErrorString, rifImageMap(*src, RIF_IMAGE_MAP_READ,
                (void **)&data));
    if (!data)
        return false;

    int devSize = desc.num_components;
    int destSize = PXLpackingComponents(dest.getPacking());
    UT_ASSERT(desc.image_width == dest.getXres());
    UT_ASSERT(desc.image_height == dest.getYres());
    exint       npix = dest.getNumPixels();
    switch (dest.getFormat())
    {
        case PXL_INT8:
            doCopyRaster<uint8, float, 0xff, false>(
                    dest.getPixels(), destSize, data, devSize, npix);
            break;
        case PXL_INT16:
            doCopyRaster<uint16, float, 0xffff, false>(
                    dest.getPixels(), destSize, data, devSize, npix);
            break;
        case PXL_INT32:
            doCopyRaster<uint32, float, 0xffffffff, false>(
                    dest.getPixels(), destSize, data, devSize, npix);
            break;
        case PXL_FLOAT16:
            doCopyRaster<fpreal16, float, 0, false>(
                    dest.getPixels(), destSize, data, devSize, npix);
            break;
        case PXL_FLOAT32:
            doCopyRaster<fpreal32, float, 0, false>(
                    dest.getPixels(), destSize, data, devSize, npix);
            break;
        default:
            UT_ASSERT(0);
            break;
    }
    RIF_CALL_FAIL(status, myErrorString, rifImageUnmap(*src, data));
    return true;
}

bool
TIL_RadeonFilter::uploadRaster(RIFContext &ctx,
        RIFImage &img,
        const char *label,
        const PXL_Raster &src,
        PXL_Raster &storage)
{
    const PXL_Raster    &rp = makeFloatRaster(src, storage);
    const void          *data = rp.getPixels();
    UT_ASSERT(rp.getFormat() == PXL_FLOAT32);

    rif_image_desc      desc;
    rif_int             status;
    memset(&desc, 0, sizeof(desc));
    desc.type = RIF_COMPONENT_TYPE_FLOAT32;
    desc.image_width = rp.getXres();
    desc.image_height = rp.getYres();
    desc.num_components = PXLpackingComponents(rp.getPacking());
    RIF_CALL(status, myErrorString, rifContextCreateImage(*ctx, &desc, data, img.ptr()));
    if (img && status != RIF_SUCCESS)
        img = nullptr;
    return img.isValid();
}

bool
TIL_RadeonFilter::apply(PXL_Raster *raster)
{
    RIFContext  ctx;
    ctx = getRIFContext(myErrorString, myDeviceNum);
    UT_ASSERT(ctx);
    if (!ctx)
        return false;
    DEBUG_PRINT("Apply RadeonPro");
    UT_StopWatch        timer; timer.start();
    int nchan = PXLpackingComponents(raster->getPacking());
    if (nchan != 3 && nchan != 4)
        return false;   // Only filter 3&4 channel images

    DEBUG_PRINT("Init: {}", timer.lap());
    rif_int             status;
    RIFCommandQueue     queue;
    RIF_CALL_FAIL(status, myErrorString,
            rifContextCreateCommandQueue(*ctx, queue.ptr()));
    if (!queue)
        return false;

    RIFImage    rgb;
    PXL_Raster  storage;
    const char  *colorName = (myFilterType == RADEON_FILTER_TYPE::BLOOM)
                                ? "color" : "colorImg";
    if (!uploadRaster(ctx, rgb, colorName, *raster, storage))
        return false;

    size_t              out_size;
    rif_image_desc      out_desc;
    rifImageGetInfo(*rgb, RIF_IMAGE_DESC, sizeof(out_desc), &out_desc, &out_size);

    RIFImage    out;
    RIFImage    normalsImg;
    RIFImage    depthImg;
    PXL_Raster  normalsImg_store;
    PXL_Raster  depthImg_store;
    out_desc.num_components = 3;
    out_desc.type = RIF_COMPONENT_TYPE_FLOAT32;
    RIF_CALL_FAIL(status, myErrorString, rifContextCreateImage(*ctx,
                &out_desc, nullptr, out.ptr()));

    RIFImageFilter      filter;
    if (myFilterType == RADEON_FILTER_TYPE::DENOISE)
    {
        RIF_CALL_FAIL(status, myErrorString, rifContextCreateImageFilter(*ctx,
                    RIF_IMAGE_FILTER_EAW_DENOISE, filter.ptr()));
        auto it = myAuxPlanes.find(theNName.asRef());
        if (it != myAuxPlanes.end())
        {
            if (uploadRaster(ctx, normalsImg, "normalsImg",
                        *it->second, normalsImg_store))
            {
                RIF_CALL_FAIL(status, myErrorString,
                        rifImageFilterSetParameterImage(*filter, "normalsImg", *normalsImg));
            }
        }
        it = myAuxPlanes.find(thePzName.asRef());
        if (it != myAuxPlanes.end())
        {
            if (uploadRaster(ctx, depthImg, "depthImg",
                        *it->second, depthImg_store))
            {
                RIF_CALL_FAIL(status, myErrorString,
                        rifImageFilterSetParameterImage(*filter, "depthImg", *depthImg));
            }
        }
        RIF_CALL_FAIL(status, myErrorString, rifCommandQueueAttachImageFilter(*queue,
                    *filter, *rgb, *out));
    }
    else if (myFilterType == RADEON_FILTER_TYPE::AIDENOISE)
    {
        RIF_CALL_FAIL(status, myErrorString, rifContextCreateImageFilter(*ctx,
                RIF_IMAGE_FILTER_AI_DENOISE, filter.ptr()));
        // Specifies if the color and albedo inputs are in the HDR or LDR format
        RIF_CALL_FAIL(status, myErrorString, rifImageFilterSetParameter1u(
                    *filter, theUseHDR.c_str(), 1));
        RIF_CALL_FAIL(status, myErrorString, rifImageFilterSetParameterString(
                    *filter, theModelPath.c_str(), myModelPath.c_str()));
        RIF_CALL_FAIL(status, myErrorString, rifImageFilterSetParameterImage(
            *filter, colorName, *rgb));
        RIF_CALL_FAIL(status, myErrorString, rifCommandQueueAttachImageFilter(*queue,
                    *filter, *rgb, *out));
    }
    else if (myFilterType == RADEON_FILTER_TYPE::BLOOM)
    {
        RIF_CALL_FAIL(status, myErrorString, rifContextCreateImageFilter(*ctx,
                    RIF_IMAGE_FILTER_BLOOM, filter.ptr()));
        RIF_CALL_FAIL(status, myErrorString, rifImageFilterSetParameter1f(
                    *filter, theRadius.c_str(), myRadius));
        RIF_CALL_FAIL(status, myErrorString, rifImageFilterSetParameter1f(
                    *filter, theThreshold.c_str(), myThreshold));
        RIF_CALL_FAIL(status, myErrorString, rifImageFilterSetParameter1f(
                    *filter, theWeight.c_str(), myWeight));
        RIF_CALL_FAIL(status, myErrorString, rifCommandQueueAttachImageFilter(*queue,
                    *filter, *rgb, *out));
    }
    RIF_CALL_FAIL(status, myErrorString, rifContextExecuteCommandQueue(
                    *ctx, *queue, nullptr, nullptr, nullptr));

    DEBUG_PRINT("Run Time: {}", timer.lap());
    downloadRaster(*raster, out);

    rifCommandQueueDetachImageFilter(*queue, *filter);

    return true;
}

namespace
{
    struct Factory final : public TIL_RasterFilter::Factory
    {
        Factory(RADEON_FILTER_TYPE type)
            : myType(type)
        {
        }
        const UT_StringHolder   &name() const override
        {
            static constexpr UT_StringLit   theBloom  ("radeon_bloom");
            static constexpr UT_StringLit   theDenoise("radeon_denoise");
            static constexpr UT_StringLit   theAIDenoise("radeon_aidenoise");
            switch (myType)
            {
                case RADEON_FILTER_TYPE::BLOOM:
                    return theBloom.asHolder();
                case RADEON_FILTER_TYPE::DENOISE:
                    return theDenoise.asHolder();
                case RADEON_FILTER_TYPE::AIDENOISE:
                    return theAIDenoise.asHolder();
            }
            return UT_StringHolder::theEmptyString;
        }
        const UT_StringHolder   &label() const override
        {
            static constexpr UT_StringLit   theBloom  ("AMD Radeon Bloom");
            static constexpr UT_StringLit   theDenoise("AMD Radeon Denoise");
            static constexpr UT_StringLit   theAIDenoise("AMD Radeon AI Denoise");
            switch (myType)
            {
                case RADEON_FILTER_TYPE::BLOOM:
                    return theBloom.asHolder();
                case RADEON_FILTER_TYPE::DENOISE:
                    return theDenoise.asHolder();
                case RADEON_FILTER_TYPE::AIDENOISE:
                    return theAIDenoise.asHolder();
            }
            return UT_StringHolder::theEmptyString;
        }
        UT_UniquePtr<TIL_RasterFilter>  newFilter() const override
        {
            return UTmakeUnique<TIL_RadeonFilter>(myType);
        }

        RADEON_FILTER_TYPE      myType;
    };
}

void
newRasterFilter()
{
    // Register different filters
    DEBUG_PRINT("Registering Radeon Image Filters");
    TIL_RasterFilter::registerFactory(
            UTmakeUnique<Factory>(RADEON_FILTER_TYPE::DENOISE));
    TIL_RasterFilter::registerFactory(
            UTmakeUnique<Factory>(RADEON_FILTER_TYPE::BLOOM));
    TIL_RasterFilter::registerFactory(
            UTmakeUnique<Factory>(RADEON_FILTER_TYPE::AIDENOISE));
}