<commit_before>#include "zlib.h"
#include "catch.hpp"
#include "helpers.h"
#include "png.h"
#include "gif_lib.h"
#include "test_weighting_helpers.h"
#include "trim_whitespace.h"
#include "string.h"
#include "lcms2.h"

// TODO: https://github.com/pornel/pngquant/

uint8_t gif_bytes_32[]
    = { 0x47, 0x49, 0x46, 0x38, 0x39, 0x61, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x21, 0xF9, 0x04,
        0x01, 0x00, 0x00, 0x00, 0x00, 0x2C, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x02 };
uint8_t gif_bytes_43[] = { 0x47, 0x49, 0x46, 0x38, 0x39, 0x61, 0x01, 0x00, 0x01, 0x00, 0x80, 0x00, 0x00, 0xFF, 0xFF,
                           0xFF, 0x00, 0x00, 0x00, 0x21, 0xF9, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2C, 0x00, 0x00,
                           0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x02, 0x02, 0x44, 0x01, 0x00, 0x3B };
static int read_func(GifFileType * gif, GifByteType * buffer, int bytes_please)
{
    int gif_byte_count = sizeof(gif_bytes_43);
    int * cursor = (int *)gif->UserData;
    if (*cursor + bytes_please > gif_byte_count) {
        fprintf(stderr, "More bytes requested (%i) by giflib than available (%i)", bytes_please,
                gif_byte_count - *cursor);
        bytes_please = gif_byte_count - *cursor;
    }
    memcpy(buffer, &gif_bytes_43[*cursor], bytes_please);
    *cursor += bytes_please;
    return bytes_please;
}

TEST_CASE("Test overlapped slurp", "")
{
    int cursor_a = 0, cursor_b = 0, error = 0;
    GifFileType * a = DGifOpen(&cursor_a, read_func, &error);
    REQUIRE(error == D_GIF_SUCCEEDED);

    REQUIRE(DGifSlurp(a) == GIF_OK);
    REQUIRE(a->Error == D_GIF_SUCCEEDED);

    GifFileType * b = DGifOpen(&cursor_b, read_func, &error);
    REQUIRE(error == D_GIF_SUCCEEDED);
    REQUIRE(DGifSlurp(b) == GIF_OK);

    int err = 0;
    DGifCloseFile(a, &err);
    DGifCloseFile(b, &err);
}

TEST_CASE("Load png", "[fastscaling]")
{

    bool success = false;

    uint8_t image_bytes_literal[]
        = { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, 0x00,
            0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1F, 0x15, 0xC4, 0x89, 0x00,
            0x00, 0x00, 0x0A, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9C, 0x63, 0x00, 0x01, 0x00, 0x00, 0x05, 0x00, 0x01,
            0x0D, 0x0A, 0x2D, 0xB4, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82 };
    png_size_t image_bytes_count = sizeof(image_bytes_literal);
    png_const_voidp image_bytes = &image_bytes_literal;

    png_image image;

    /* Only the image structure version number needs to be set. */
    memset(&image, 0, sizeof image);
    image.version = PNG_IMAGE_VERSION;
    image.opaque = NULL;

    if (png_image_begin_read_from_memory(&image, image_bytes, image_bytes_count)) {
        png_bytep buffer;

        /* Change this to try different formats!  If you set a colormap format
         * then you must also supply a colormap below.
         */
        image.format = PNG_FORMAT_BGRA;

        buffer = (png_bytep)malloc(PNG_IMAGE_SIZE(image));

        if (buffer != NULL) {
            if (png_image_finish_read(&image, NULL /*background*/, buffer, 0 /*row_stride*/,
                                      NULL /*colormap for PNG_FORMAT_FLAG_COLORMAP */)) {

                success = true;

                //                if (png_image_write_to_file(&image, argv[2],
                //                                            0/*convert_to_8bit*/, buffer, 0/*row_stride*/,
                //                                            NULL/*colormap*/))
                //                    result = 0;
                //
                //                else
                //                    fprintf(stderr, "pngtopng: write %s: %s\n",
                //                            image.message);
                //
                //                free(buffer);
            }

            else {
                fprintf(stderr, "png_image_finish_read: %s\n", image.message);

                /* This is the only place where a 'free' is required; libpng does
                 * the cleanup on error and success, but in this case we couldn't
                 * complete the read because of running out of memory.
                 */
                png_image_free(&image);
            }
            free(buffer);
        }

        else
            fprintf(stderr, "pngtopng: out of memory: %lu bytes\n", (unsigned long)PNG_IMAGE_SIZE(image));
    }

    else
        /* Failed to read the first argument: */
        fprintf(stderr, "png_image_begin_read_from_memory: %s\n", image.message);

    REQUIRE(success);
}

TEST_CASE("Load png from URL", "[fastscaling]")
{

    bool success = false;

    size_t bytes_count = 0;
    flow_c * c = flow_context_create();

    uint8_t * bytes = get_bytes_cached(c, &bytes_count, "http://s3.amazonaws.com/resizer-images/sun_256.png", __FILE__);
    REQUIRE_FALSE(bytes == NULL);
    png_size_t image_bytes_count = bytes_count;
    png_const_voidp image_bytes = bytes;

    png_image image;

    /* Only the image structure version number needs to be set. */
    memset(&image, 0, sizeof image);
    image.version = PNG_IMAGE_VERSION;
    image.opaque = NULL;

    if (png_image_begin_read_from_memory(&image, image_bytes, image_bytes_count)) {
        png_bytep buffer;

        /* Change this to try different formats!  If you set a colormap format
         * then you must also supply a colormap below.
         */
        image.format = PNG_FORMAT_BGRA;

        buffer = FLOW_calloc_array(c, PNG_IMAGE_SIZE(image), png_byte);

        if (buffer != NULL) {
            if (png_image_finish_read(&image, NULL /*background*/, buffer, 0 /*row_stride*/,
                                      NULL /*colormap for PNG_FORMAT_FLAG_COLORMAP */)) {

                int nonzero = (int)nonzero_count((uint8_t *)buffer, PNG_IMAGE_SIZE(image));
                if (nonzero > 0) {
                    printf("nonzero buffer: %d of %d", nonzero, PNG_IMAGE_SIZE(image));
                }

                struct flow_bitmap_bgra * source
                    = flow_bitmap_bgra_create_header(c, (unsigned int)(image.width), (unsigned int)(image.height));
                REQUIRE_FALSE(source == NULL);
                source->fmt = flow_pixel_format::flow_bgra32;
                source->stride = PNG_IMAGE_ROW_STRIDE(image);
                printf("png stride (%d), calculated (%d)\n", source->stride,
                       source->w * flow_pixel_format_bytes_per_pixel(source->fmt));
                source->alpha_meaningful = true;
                source->pixels = buffer;

                int target_width = 300;
                int target_height = 200;

                struct flow_bitmap_bgra * canvas
                    = flow_bitmap_bgra_create(c, target_width, target_height, true, flow_bgra32);

                REQUIRE_FALSE(canvas == NULL);
                struct flow_RenderDetails * details
                    = flow_RenderDetails_create_with(c, flow_interpolation_filter::flow_interpolation_filter_Robidoux);
                details->interpolate_last_percent = 2.1f;
                details->minimum_sample_window_to_interposharpen = 1.5;
                details->havling_acceptable_pixel_loss = 0.26f;

                REQUIRE_FALSE(details == NULL);
                //                details->sharpen_percent_goal = 50;
                //                details->post_flip_x = flipx;
                //                details->post_flip_y = flipy;
                //                details->post_transpose = transpose;
                // details->enable_profiling = profile;

                // Should we even have Renderer_* functions, or just 1 call that does it all?
                // If we add memory use estimation, we should keep flow_Renderer

                if (!flow_RenderDetails_render(c, details, source, canvas)) {

                    char error[255];
                    flow_context_error_message(c, error, 255);
                    printf("%s", error);
                    exit(77);
                }
                printf("Rendered!");
                flow_RenderDetails_destroy(c, details);

                flow_bitmap_bgra_destroy(c, source);

                // TODO, write out PNG here

                struct flow_io * buf = flow_io_create_for_output_buffer(c, c);
                // memset(canvas->pixels, 0, canvas->stride * canvas->h);
                if (!flow_bitmap_bgra_write_png(c, canvas, buf)) {
                    // FLOW_error_return(context);
                    FAIL("Failed to write png");
                } else {
                    if (!flow_io_write_output_buffer_to_file(c, buf, "outpng.png")) {
                        FAIL("Failed to copy to disk");
                    }
                    success = true;
                }

                flow_bitmap_bgra_destroy(c, canvas);

            } else {
                fprintf(stderr, "png_image_finish_read: %s\n", image.message);

                /* This is the only place where a 'free' is required; libpng does
                 * the cleanup on error and success, but in this case we couldn't
                 * complete the read because of running out of memory.
                 */
                png_image_free(&image);
            }
        }

        else
            fprintf(stderr, "pngtopng: out of memory: %lu bytes\n", (unsigned long)PNG_IMAGE_SIZE(image));
    }

    else
        /* Failed to read the first argument: */
        fprintf(stderr, "png_image_begin_read_from_memory: %s\n", image.message);

    flow_context_destroy(c);
    REQUIRE(success);
}

// This test is windows only
// TEST_CASE("Test png writing; repro uninitialized value", "[libpng]"){
//    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL,
//                                                  NULL); // makepng_error, makepng_warning);
//
//    uint32_t w = 300;
//    uint32_t h = 200;
//    uint32_t stride = 1200;
//
//    uint8_t * pixels = (uint8_t *)calloc(1, stride * h);
//    uint8_t ** rows = (uint8_t **)calloc(1, sizeof(uint8_t *) * h);
//
//    unsigned int y;
//    for (y = 0; y < h; ++y) {
//        rows[y] = ((uint8_t *)pixels + (stride * y));
//    }
//
//    png_set_compression_level(png_ptr, Z_BEST_SPEED);
//    png_set_text_compression_level(png_ptr, Z_DEFAULT_COMPRESSION);
//
//    png_init_io(png_ptr, fopen("/dev/null","wb"));
//
//    png_infop info_ptr = NULL;
//    info_ptr = png_create_info_struct(png_ptr);
//    if (info_ptr == NULL)
//        png_error(png_ptr, "OOM allocating info structure"); // TODO: comprehend png error handling
//    {
//        png_set_rows(png_ptr, info_ptr, rows);
//
//        int color_type = PNG_COLOR_TYPE_RGB_ALPHA;
//        int transform = PNG_TRANSFORM_BGR;
//
//        png_set_IHDR(png_ptr, info_ptr, (png_uint_32)w, (png_uint_32)h, 8, color_type,
//                     PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
//
//        png_set_sRGB_gAMA_and_cHRM(png_ptr, info_ptr, PNG_sRGB_INTENT_PERCEPTUAL);
//
//        png_write_png(png_ptr, info_ptr, transform, NULL);
//
//        rows = NULL;
//        png_destroy_write_struct(&png_ptr, &info_ptr);
//    }
//    free(pixels);
//    free(rows);
//}

TEST_CASE("Test png writing", "[libpng]")
{
    printf("\n Testing libpng version %s\n", PNG_LIBPNG_VER_STRING);
    printf("   with zlib   version %s\n", ZLIB_VERSION);
    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL,
                                                  NULL); // makepng_error, makepng_warning);

    uint32_t w = 300;
    uint32_t h = 200;
    uint32_t stride = 1200;

    uint8_t * pixels = (uint8_t *)calloc(1, stride * h);
    uint8_t ** rows = (uint8_t **)calloc(1, sizeof(uint8_t *) * h);

    unsigned int y;
    for (y = 0; y < h; ++y) {
        rows[y] = ((uint8_t *)pixels + (stride * y));
    }

    png_set_compression_level(png_ptr, Z_BEST_SPEED);
    png_set_text_compression_level(png_ptr, Z_DEFAULT_COMPRESSION);

    png_init_io(png_ptr, fopen("/dev/null", "wb"));
    png_infop info_ptr = NULL;
    info_ptr = png_create_info_struct(png_ptr);
    png_set_rows(png_ptr, info_ptr, rows);
    int color_type = PNG_COLOR_TYPE_RGB_ALPHA;
    int transform = PNG_TRANSFORM_BGR;
    png_set_IHDR(png_ptr, info_ptr, (png_uint_32)w, (png_uint_32)h, 8, color_type,
                 PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
    png_set_sRGB_gAMA_and_cHRM(png_ptr, info_ptr, PNG_sRGB_INTENT_PERCEPTUAL);
//Uninitialized read happens here:
    png_write_png(png_ptr, info_ptr, transform, NULL);

    png_destroy_write_struct(&png_ptr, &info_ptr);
    free(pixels);
    free(rows);
}
<commit_msg>Don't try to use dev/null on windows. It's not there.<commit_after>#include "zlib.h"
#include "catch.hpp"
#include "helpers.h"
#include "png.h"
#include "gif_lib.h"
#include "test_weighting_helpers.h"
#include "trim_whitespace.h"
#include "string.h"
#include "lcms2.h"

// TODO: https://github.com/pornel/pngquant/

uint8_t gif_bytes_32[]
    = { 0x47, 0x49, 0x46, 0x38, 0x39, 0x61, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x21, 0xF9, 0x04,
        0x01, 0x00, 0x00, 0x00, 0x00, 0x2C, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x02 };
uint8_t gif_bytes_43[] = { 0x47, 0x49, 0x46, 0x38, 0x39, 0x61, 0x01, 0x00, 0x01, 0x00, 0x80, 0x00, 0x00, 0xFF, 0xFF,
                           0xFF, 0x00, 0x00, 0x00, 0x21, 0xF9, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2C, 0x00, 0x00,
                           0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x02, 0x02, 0x44, 0x01, 0x00, 0x3B };
static int read_func(GifFileType * gif, GifByteType * buffer, int bytes_please)
{
    int gif_byte_count = sizeof(gif_bytes_43);
    int * cursor = (int *)gif->UserData;
    if (*cursor + bytes_please > gif_byte_count) {
        fprintf(stderr, "More bytes requested (%i) by giflib than available (%i)", bytes_please,
                gif_byte_count - *cursor);
        bytes_please = gif_byte_count - *cursor;
    }
    memcpy(buffer, &gif_bytes_43[*cursor], bytes_please);
    *cursor += bytes_please;
    return bytes_please;
}

TEST_CASE("Test overlapped slurp", "")
{
    int cursor_a = 0, cursor_b = 0, error = 0;
    GifFileType * a = DGifOpen(&cursor_a, read_func, &error);
    REQUIRE(error == D_GIF_SUCCEEDED);

    REQUIRE(DGifSlurp(a) == GIF_OK);
    REQUIRE(a->Error == D_GIF_SUCCEEDED);

    GifFileType * b = DGifOpen(&cursor_b, read_func, &error);
    REQUIRE(error == D_GIF_SUCCEEDED);
    REQUIRE(DGifSlurp(b) == GIF_OK);

    int err = 0;
    DGifCloseFile(a, &err);
    DGifCloseFile(b, &err);
}

TEST_CASE("Load png", "[fastscaling]")
{

    bool success = false;

    uint8_t image_bytes_literal[]
        = { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, 0x00,
            0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1F, 0x15, 0xC4, 0x89, 0x00,
            0x00, 0x00, 0x0A, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9C, 0x63, 0x00, 0x01, 0x00, 0x00, 0x05, 0x00, 0x01,
            0x0D, 0x0A, 0x2D, 0xB4, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82 };
    png_size_t image_bytes_count = sizeof(image_bytes_literal);
    png_const_voidp image_bytes = &image_bytes_literal;

    png_image image;

    /* Only the image structure version number needs to be set. */
    memset(&image, 0, sizeof image);
    image.version = PNG_IMAGE_VERSION;
    image.opaque = NULL;

    if (png_image_begin_read_from_memory(&image, image_bytes, image_bytes_count)) {
        png_bytep buffer;

        /* Change this to try different formats!  If you set a colormap format
         * then you must also supply a colormap below.
         */
        image.format = PNG_FORMAT_BGRA;

        buffer = (png_bytep)malloc(PNG_IMAGE_SIZE(image));

        if (buffer != NULL) {
            if (png_image_finish_read(&image, NULL /*background*/, buffer, 0 /*row_stride*/,
                                      NULL /*colormap for PNG_FORMAT_FLAG_COLORMAP */)) {

                success = true;

                //                if (png_image_write_to_file(&image, argv[2],
                //                                            0/*convert_to_8bit*/, buffer, 0/*row_stride*/,
                //                                            NULL/*colormap*/))
                //                    result = 0;
                //
                //                else
                //                    fprintf(stderr, "pngtopng: write %s: %s\n",
                //                            image.message);
                //
                //                free(buffer);
            }

            else {
                fprintf(stderr, "png_image_finish_read: %s\n", image.message);

                /* This is the only place where a 'free' is required; libpng does
                 * the cleanup on error and success, but in this case we couldn't
                 * complete the read because of running out of memory.
                 */
                png_image_free(&image);
            }
            free(buffer);
        }

        else
            fprintf(stderr, "pngtopng: out of memory: %lu bytes\n", (unsigned long)PNG_IMAGE_SIZE(image));
    }

    else
        /* Failed to read the first argument: */
        fprintf(stderr, "png_image_begin_read_from_memory: %s\n", image.message);

    REQUIRE(success);
}

TEST_CASE("Load png from URL", "[fastscaling]")
{

    bool success = false;

    size_t bytes_count = 0;
    flow_c * c = flow_context_create();

    uint8_t * bytes = get_bytes_cached(c, &bytes_count, "http://s3.amazonaws.com/resizer-images/sun_256.png", __FILE__);
    REQUIRE_FALSE(bytes == NULL);
    png_size_t image_bytes_count = bytes_count;
    png_const_voidp image_bytes = bytes;

    png_image image;

    /* Only the image structure version number needs to be set. */
    memset(&image, 0, sizeof image);
    image.version = PNG_IMAGE_VERSION;
    image.opaque = NULL;

    if (png_image_begin_read_from_memory(&image, image_bytes, image_bytes_count)) {
        png_bytep buffer;

        /* Change this to try different formats!  If you set a colormap format
         * then you must also supply a colormap below.
         */
        image.format = PNG_FORMAT_BGRA;

        buffer = FLOW_calloc_array(c, PNG_IMAGE_SIZE(image), png_byte);

        if (buffer != NULL) {
            if (png_image_finish_read(&image, NULL /*background*/, buffer, 0 /*row_stride*/,
                                      NULL /*colormap for PNG_FORMAT_FLAG_COLORMAP */)) {

                int nonzero = (int)nonzero_count((uint8_t *)buffer, PNG_IMAGE_SIZE(image));
                if (nonzero > 0) {
                    printf("nonzero buffer: %d of %d", nonzero, PNG_IMAGE_SIZE(image));
                }

                struct flow_bitmap_bgra * source
                    = flow_bitmap_bgra_create_header(c, (unsigned int)(image.width), (unsigned int)(image.height));
                REQUIRE_FALSE(source == NULL);
                source->fmt = flow_pixel_format::flow_bgra32;
                source->stride = PNG_IMAGE_ROW_STRIDE(image);
                printf("png stride (%d), calculated (%d)\n", source->stride,
                       source->w * flow_pixel_format_bytes_per_pixel(source->fmt));
                source->alpha_meaningful = true;
                source->pixels = buffer;

                int target_width = 300;
                int target_height = 200;

                struct flow_bitmap_bgra * canvas
                    = flow_bitmap_bgra_create(c, target_width, target_height, true, flow_bgra32);

                REQUIRE_FALSE(canvas == NULL);
                struct flow_RenderDetails * details
                    = flow_RenderDetails_create_with(c, flow_interpolation_filter::flow_interpolation_filter_Robidoux);
                details->interpolate_last_percent = 2.1f;
                details->minimum_sample_window_to_interposharpen = 1.5;
                details->havling_acceptable_pixel_loss = 0.26f;

                REQUIRE_FALSE(details == NULL);
                //                details->sharpen_percent_goal = 50;
                //                details->post_flip_x = flipx;
                //                details->post_flip_y = flipy;
                //                details->post_transpose = transpose;
                // details->enable_profiling = profile;

                // Should we even have Renderer_* functions, or just 1 call that does it all?
                // If we add memory use estimation, we should keep flow_Renderer

                if (!flow_RenderDetails_render(c, details, source, canvas)) {

                    char error[255];
                    flow_context_error_message(c, error, 255);
                    printf("%s", error);
                    exit(77);
                }
                printf("Rendered!");
                flow_RenderDetails_destroy(c, details);

                flow_bitmap_bgra_destroy(c, source);

                // TODO, write out PNG here

                struct flow_io * buf = flow_io_create_for_output_buffer(c, c);
                // memset(canvas->pixels, 0, canvas->stride * canvas->h);
                if (!flow_bitmap_bgra_write_png(c, canvas, buf)) {
                    // FLOW_error_return(context);
                    FAIL("Failed to write png");
                } else {
                    if (!flow_io_write_output_buffer_to_file(c, buf, "outpng.png")) {
                        FAIL("Failed to copy to disk");
                    }
                    success = true;
                }

                flow_bitmap_bgra_destroy(c, canvas);

            } else {
                fprintf(stderr, "png_image_finish_read: %s\n", image.message);

                /* This is the only place where a 'free' is required; libpng does
                 * the cleanup on error and success, but in this case we couldn't
                 * complete the read because of running out of memory.
                 */
                png_image_free(&image);
            }
        }

        else
            fprintf(stderr, "pngtopng: out of memory: %lu bytes\n", (unsigned long)PNG_IMAGE_SIZE(image));
    }

    else
        /* Failed to read the first argument: */
        fprintf(stderr, "png_image_begin_read_from_memory: %s\n", image.message);

    flow_context_destroy(c);
    REQUIRE(success);
}

#ifdef _MSC_VER
// /dev/null doesn't work on windows. Don't try this test.
#else

TEST_CASE("Test png writing", "[libpng]")
{
    printf("\n Testing libpng version %s\n", PNG_LIBPNG_VER_STRING);
    printf("   with zlib   version %s\n", ZLIB_VERSION);
    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL,
                                                  NULL); // makepng_error, makepng_warning);

    uint32_t w = 300;
    uint32_t h = 200;
    uint32_t stride = 1200;

    uint8_t * pixels = (uint8_t *)calloc(1, stride * h);
    uint8_t ** rows = (uint8_t **)calloc(1, sizeof(uint8_t *) * h);

    unsigned int y;
    for (y = 0; y < h; ++y) {
        rows[y] = ((uint8_t *)pixels + (stride * y));
    }

    png_set_compression_level(png_ptr, Z_BEST_SPEED);
    png_set_text_compression_level(png_ptr, Z_DEFAULT_COMPRESSION);

    png_init_io(png_ptr, fopen("/dev/null", "wb"));
    png_infop info_ptr = NULL;
    info_ptr = png_create_info_struct(png_ptr);
    png_set_rows(png_ptr, info_ptr, rows);
    int color_type = PNG_COLOR_TYPE_RGB_ALPHA;
    int transform = PNG_TRANSFORM_BGR;
    png_set_IHDR(png_ptr, info_ptr, (png_uint_32)w, (png_uint_32)h, 8, color_type,
                 PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
    png_set_sRGB_gAMA_and_cHRM(png_ptr, info_ptr, PNG_sRGB_INTENT_PERCEPTUAL);
//Uninitialized read happens here:
    png_write_png(png_ptr, info_ptr, transform, NULL);

    png_destroy_write_struct(&png_ptr, &info_ptr);
    free(pixels);
    free(rows);
}
#endif
<|endoftext|>


<commit_before>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/gpu_fusible.h"

#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"

namespace xla {
namespace gpu {

namespace {
void AppendParams(const HloInstruction& instr,
                  std::vector<HloInstruction*>* params) {
  if (instr.opcode() == HloOpcode::kFusion) {
    params->insert(std::end(*params), std::begin(instr.fused_parameters()),
                   std::end(instr.fused_parameters()));
  } else {
    for (HloInstruction* operand : instr.operands()) {
      params->push_back(operand);
    }
  }
}
}  // namespace

bool LayoutsAreReduceInputFusionFriendly(const HloInstruction& producer,
                                         const HloInstruction& reduce) {
  std::vector<HloInstruction*> params;
  AppendParams(producer, &params);
  AppendParams(reduce, &params);
  int64 max_rank = -1;
  const Layout* max_rank_layout;
  for (HloInstruction* param : params) {
    if (param->shape().IsArray() && param->shape().rank() > max_rank) {
      max_rank = param->shape().rank();
      max_rank_layout = &param->shape().layout();
    }
  }
  return absl::c_all_of(params, [&](HloInstruction* param) {
    return (!param->shape().IsArray()) || (param->shape().rank() < max_rank) ||
           (LayoutUtil::Equal(param->shape().layout(), *max_rank_layout));
  });
}

bool IsReduceInputFusion(const HloInstruction& instr) {
  if (instr.IsMultiOutputFusion()) {
    for (const HloInstruction* operand :
         instr.fused_expression_root()->operands()) {
      if (IsReductionToVector(*operand)) {
        CHECK(instr.fusion_kind() == HloInstruction::FusionKind::kInput)
            << " Multi-output fusion rooted at reduction-to-vector ops must be "
               "of kind kInput: "
            << instr.ToString();
        return true;
      }
    }
  } else if (instr.opcode() == HloOpcode::kFusion &&
             IsReductionToVector(*instr.fused_expression_root())) {
    CHECK(instr.fusion_kind() == HloInstruction::FusionKind::kInput)
        << " Fusion rooted at reduction-to-vector op must be of kind kInput: "
        << instr.ToString();
    return true;
  }
  return false;
}

bool IsInputFusibleReduction(const HloInstruction& instr) {
  return IsReduceInputFusion(instr) || IsReductionToVector(instr);
}

bool ShapesCompatibleForMultiOutputFusion(const HloInstruction& instr1,
                                          const HloInstruction& instr2) {
  // Returns the instructions that determines the emitter used for lowering,
  // sometimes referred to as "the real hero".
  auto get_real_hero =
      [&](const HloInstruction* instr) -> const HloInstruction* {
    if (instr->opcode() == HloOpcode::kFusion) {
      auto fused_expression_root = instr->fused_expression_root();
      if (instr->IsMultiOutputFusion()) {
        // If possible, we want to pick a reduction-to-vector operand of the
        // fusion root, because it has the most constraints.
        for (const auto* inst : fused_expression_root->operands()) {
          if (IsReductionToVector(*inst)) {
            return inst;
          }
        }
        return fused_expression_root->operands()[0];
      }
      return fused_expression_root;
    }
    return instr;
  };

  // Multi-output fusion kernels share a common parallel loop. The loop
  // dimenstions are determined by instruction shapes.
  auto get_loop_shape = [&](const HloInstruction* element_instr) {
    // Special-case reduction-to-vector ops: The loop dimensions are determined
    // by the shape of the first operand.
    if (IsReductionToVector(*element_instr)) {
      return element_instr->operand(0)->shape();
    }
    return element_instr->shape();
  };

  // All shapes of the root tuple of multi-output fusions should agree, i.e. all
  // root ops should have equal output shapes. An exception are
  // reduction-to-vector ops. Here the input shapes of the reduction (first
  // operand shape) and the reduction dimensions need to match.
  auto* instr_1 = get_real_hero(&instr1);
  auto* instr_2 = get_real_hero(&instr2);
  // TODO(tjoerg): Relax the shape constraint. The datatype does not matter.
  if (IsReductionToVector(*instr_1) && IsReductionToVector(*instr_2) &&
      (!ShapeUtil::Equal(instr_1->shape(), instr_2->shape()) ||
       instr_1->dimensions() != instr_2->dimensions())) {
    return false;
  }
  // The elementwise output shapes must be the same (including layout).
  // TODO(tjoerg): Further relax the constraint. The datatype does not matter.
  return ShapeUtil::EqualIgnoringFpPrecision(get_loop_shape(instr_1),
                                             get_loop_shape(instr_2));
}

bool IsFusible(const HloInstruction& instr) {
  // Don't fuse get-tuple-element on GPU: We can, but it's slower than not
  // fusing.  We never generate kernels for unfused GTEs.  Instead, if an
  // unfused GTE is an input to a kernel (including a fusion kernel), we
  // compute the address of the GTE at the top of the kernel.  Often we know the
  // address of the GTE result statically, so we can do this without chasing any
  // pointers.
  return (instr.IsElementwise() && instr.operand_count() > 0) ||
         instr.opcode() == HloOpcode::kBitcast ||
         instr.opcode() == HloOpcode::kBroadcast ||
         instr.opcode() == HloOpcode::kConcatenate ||
         instr.opcode() == HloOpcode::kDynamicSlice ||
         instr.opcode() == HloOpcode::kDynamicUpdateSlice ||
         instr.opcode() == HloOpcode::kFusion ||
         instr.opcode() == HloOpcode::kGather ||
         instr.opcode() == HloOpcode::kIota || instr.opcode() == HloOpcode::kPad ||
         instr.opcode() == HloOpcode::kReduce ||
         instr.opcode() == HloOpcode::kReduceWindow ||
         instr.opcode() == HloOpcode::kReshape ||
         instr.opcode() == HloOpcode::kReverse ||
         instr.opcode() == HloOpcode::kScatter ||
         instr.opcode() == HloOpcode::kSlice ||
         instr.opcode() == HloOpcode::kTranspose;
}

}  // namespace gpu
}  // namespace xla
<commit_msg>Fix formatting issues in IsFusible<commit_after>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/gpu_fusible.h"

#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"

namespace xla {
namespace gpu {

namespace {
void AppendParams(const HloInstruction& instr,
                  std::vector<HloInstruction*>* params) {
  if (instr.opcode() == HloOpcode::kFusion) {
    params->insert(std::end(*params), std::begin(instr.fused_parameters()),
                   std::end(instr.fused_parameters()));
  } else {
    for (HloInstruction* operand : instr.operands()) {
      params->push_back(operand);
    }
  }
}
}  // namespace

bool LayoutsAreReduceInputFusionFriendly(const HloInstruction& producer,
                                         const HloInstruction& reduce) {
  std::vector<HloInstruction*> params;
  AppendParams(producer, &params);
  AppendParams(reduce, &params);
  int64 max_rank = -1;
  const Layout* max_rank_layout;
  for (HloInstruction* param : params) {
    if (param->shape().IsArray() && param->shape().rank() > max_rank) {
      max_rank = param->shape().rank();
      max_rank_layout = &param->shape().layout();
    }
  }
  return absl::c_all_of(params, [&](HloInstruction* param) {
    return (!param->shape().IsArray()) || (param->shape().rank() < max_rank) ||
           (LayoutUtil::Equal(param->shape().layout(), *max_rank_layout));
  });
}

bool IsReduceInputFusion(const HloInstruction& instr) {
  if (instr.IsMultiOutputFusion()) {
    for (const HloInstruction* operand :
         instr.fused_expression_root()->operands()) {
      if (IsReductionToVector(*operand)) {
        CHECK(instr.fusion_kind() == HloInstruction::FusionKind::kInput)
            << " Multi-output fusion rooted at reduction-to-vector ops must be "
               "of kind kInput: "
            << instr.ToString();
        return true;
      }
    }
  } else if (instr.opcode() == HloOpcode::kFusion &&
             IsReductionToVector(*instr.fused_expression_root())) {
    CHECK(instr.fusion_kind() == HloInstruction::FusionKind::kInput)
        << " Fusion rooted at reduction-to-vector op must be of kind kInput: "
        << instr.ToString();
    return true;
  }
  return false;
}

bool IsInputFusibleReduction(const HloInstruction& instr) {
  return IsReduceInputFusion(instr) || IsReductionToVector(instr);
}

bool ShapesCompatibleForMultiOutputFusion(const HloInstruction& instr1,
                                          const HloInstruction& instr2) {
  // Returns the instructions that determines the emitter used for lowering,
  // sometimes referred to as "the real hero".
  auto get_real_hero =
      [&](const HloInstruction* instr) -> const HloInstruction* {
    if (instr->opcode() == HloOpcode::kFusion) {
      auto fused_expression_root = instr->fused_expression_root();
      if (instr->IsMultiOutputFusion()) {
        // If possible, we want to pick a reduction-to-vector operand of the
        // fusion root, because it has the most constraints.
        for (const auto* inst : fused_expression_root->operands()) {
          if (IsReductionToVector(*inst)) {
            return inst;
          }
        }
        return fused_expression_root->operands()[0];
      }
      return fused_expression_root;
    }
    return instr;
  };

  // Multi-output fusion kernels share a common parallel loop. The loop
  // dimenstions are determined by instruction shapes.
  auto get_loop_shape = [&](const HloInstruction* element_instr) {
    // Special-case reduction-to-vector ops: The loop dimensions are determined
    // by the shape of the first operand.
    if (IsReductionToVector(*element_instr)) {
      return element_instr->operand(0)->shape();
    }
    return element_instr->shape();
  };

  // All shapes of the root tuple of multi-output fusions should agree, i.e. all
  // root ops should have equal output shapes. An exception are
  // reduction-to-vector ops. Here the input shapes of the reduction (first
  // operand shape) and the reduction dimensions need to match.
  auto* instr_1 = get_real_hero(&instr1);
  auto* instr_2 = get_real_hero(&instr2);
  // TODO(tjoerg): Relax the shape constraint. The datatype does not matter.
  if (IsReductionToVector(*instr_1) && IsReductionToVector(*instr_2) &&
      (!ShapeUtil::Equal(instr_1->shape(), instr_2->shape()) ||
       instr_1->dimensions() != instr_2->dimensions())) {
    return false;
  }
  // The elementwise output shapes must be the same (including layout).
  // TODO(tjoerg): Further relax the constraint. The datatype does not matter.
  return ShapeUtil::EqualIgnoringFpPrecision(get_loop_shape(instr_1),
                                             get_loop_shape(instr_2));
}

bool IsFusible(const HloInstruction& instr) {
  // Don't fuse get-tuple-element on GPU: We can, but it's slower than not
  // fusing.  We never generate kernels for unfused GTEs.  Instead, if an
  // unfused GTE is an input to a kernel (including a fusion kernel), we
  // compute the address of the GTE at the top of the kernel.  Often we know the
  // address of the GTE result statically, so we can do this without chasing any
  // pointers.
  return (instr.IsElementwise() && instr.operand_count() > 0) ||
         instr.opcode() == HloOpcode::kBitcast ||
         instr.opcode() == HloOpcode::kBroadcast ||
         instr.opcode() == HloOpcode::kConcatenate ||
         instr.opcode() == HloOpcode::kDynamicSlice ||
         instr.opcode() == HloOpcode::kDynamicUpdateSlice ||
         instr.opcode() == HloOpcode::kFusion ||
         instr.opcode() == HloOpcode::kGather ||
         instr.opcode() == HloOpcode::kIota ||
         instr.opcode() == HloOpcode::kPad ||
         instr.opcode() == HloOpcode::kReduce ||
         instr.opcode() == HloOpcode::kReduceWindow ||
         instr.opcode() == HloOpcode::kReshape ||
         instr.opcode() == HloOpcode::kReverse ||
         instr.opcode() == HloOpcode::kScatter ||
         instr.opcode() == HloOpcode::kSlice ||
         instr.opcode() == HloOpcode::kTranspose;
}

}  // namespace gpu
}  // namespace xla
<|endoftext|>


<commit_before>// Author: Enrico Guiraud, CERN 10/2021

/*************************************************************************
 * Copyright (C) 1995-2021, Rene Brun and Fons Rademakers.               *
 * All rights reserved.                                                  *
 *                                                                       *
 * For the licensing terms see $ROOTSYS/LICENSE.                         *
 * For the list of contributors see $ROOTSYS/README/CREDITS.             *
 *************************************************************************/

#ifndef ROOT_RDF_RVARIATION
#define ROOT_RDF_RVARIATION

#include "Utils.hxx" // IsRVec
#include "ColumnReaderUtils.hxx"
#include "RColumnReaderBase.hxx"
#include "RLoopManager.hxx"
#include "RVariationBase.hxx"

#include <ROOT/RStringView.hxx>
#include <ROOT/TypeTraits.hxx>
#include <Rtypes.h> // R__CLING_PTRCHECK, Long64_t

#include <array>
#include <deque>
#include <map>
#include <string>
#include <type_traits> // std::is_same, std::conditional_t
#include <utility>     // std::index_sequence
#include <vector>

class TTreeReader;

namespace ROOT {

namespace RDF {
class RDataSource;
}

namespace Internal {
namespace RDF {

using namespace ROOT::TypeTraits;

template <typename T>
bool ResultsSizeEq(const RVec<RVec<T>> &results, std::size_t expected, std::size_t nColumns)
{
   if (nColumns == 1)
      return results.size() == expected;

   return std::all_of(results.begin(), results.end(), [&expected](const RVec<T> &v) { return v.size() == expected; });
}

template <typename T>
bool ResultsSizeEq(const RVec<T> &results, std::size_t expected, std::size_t nColumns)
{
   assert(nColumns == 1);
   (void)nColumns;

   return results.size() == expected;
}

template <typename T>
std::size_t GetNVariations(const RVec<RVec<T>> &results)
{
   assert(!results.empty());
   return results[0].size();
}

template <typename T>
std::size_t GetNVariations(const RVec<T> &results)
{
   return results.size();
}

template <typename F>
class R__CLING_PTRCHECK(off) RVariation final : public RVariationBase {
   using ColumnTypes_t = typename CallableTraits<F>::arg_types;
   using TypeInd_t = std::make_index_sequence<ColumnTypes_t::list_size>;
   using ret_type = typename CallableTraits<F>::ret_type;
   // Avoid instantiating vector<bool> as `operator[]` returns temporaries in that case. Use std::deque instead.
   using ValuesPerSlot_t =
      std::conditional_t<std::is_same<ret_type, bool>::value, std::deque<ret_type>, std::vector<ret_type>>;

   F fExpression;
   ValuesPerSlot_t fLastResults;

   /// Column readers per slot and per input column
   std::vector<std::array<std::unique_ptr<RColumnReaderBase>, ColumnTypes_t::list_size>> fValues;

   template <typename... ColTypes, std::size_t... S>
   void UpdateHelper(unsigned int slot, Long64_t entry, TypeList<ColTypes...>, std::index_sequence<S...>)
   {
      // fExpression must return an RVec<T>
      const auto &results = fExpression(fValues[slot][S]->template Get<ColTypes>(entry)...);
      if (!ResultsSizeEq(results, fVariationNames.size(), fColNames.size())) {
         std::string variationName = fVariationNames[0].substr(0, fVariationNames[0].find_first_of(':'));
         throw std::runtime_error("The evaluation of the expression for variation \"" + variationName +
                                  "\" resulted in " + std::to_string(GetNVariations(results)) + " values, but " +
                                  std::to_string(fVariationNames.size()) + " were expected.");
      }
      // Assign into fLastResults without changing the addresses of its elements (we gave those addresses away in
      // GetValuePtr)
      fLastResults[slot * CacheLineStep<ret_type>()].assign(results.begin(), results.end());

      // silence "unused parameter" warnings in gcc
      (void)slot;
      (void)entry;
   }

   // This overload is for the case of a single column and ret_type != RVec<RVec<...>> -- the colIdx is ignored.
   template <typename U = typename ret_type::value_type>
   std::enable_if_t<!IsRVec<U>::value, void *>
   GetValuePtr(unsigned int slot, std::size_t /*colIdx*/, std::size_t varIdx)
   {
      auto &value = fLastResults[slot * CacheLineStep<ret_type>()][varIdx];
      return static_cast<void *>(&value);
   }

   // This overload is for the case of ret_type == RVec<RVec<...>>
   template <typename U = typename ret_type::value_type>
   std::enable_if_t<IsRVec<U>::value, void *> GetValuePtr(unsigned int slot, std::size_t colIdx, std::size_t varIdx)
   {
      if (fColNames.size() == 1) {
         auto &value = fLastResults[slot * CacheLineStep<ret_type>()][varIdx];
         return static_cast<void *>(&value);
      }

      auto &value = fLastResults[slot * CacheLineStep<ret_type>()][colIdx][varIdx];
      return static_cast<void *>(&value);
   }

public:
   RVariation(const std::vector<std::string> &colNames, std::string_view variationName, F expression,
              const std::vector<std::string> &variationTags, std::string_view type, const RColumnRegister &defines,
              RLoopManager &lm, const ColumnNames_t &inputColNames)
      : RVariationBase(colNames, variationName, variationTags, type, defines, lm, inputColNames),
        fExpression(std::move(expression)), fLastResults(lm.GetNSlots() * RDFInternal::CacheLineStep<ret_type>()),
        fValues(lm.GetNSlots())
   {
      for (auto i = 0u; i < lm.GetNSlots(); ++i)
         fLastResults[i * RDFInternal::CacheLineStep<ret_type>()].resize(fVariationNames.size());
   }

   RVariation(const RVariation &) = delete;
   RVariation &operator=(const RVariation &) = delete;

   void InitSlot(TTreeReader *r, unsigned int slot) final
   {
      for (auto &define : fColumnRegister.GetColumns())
         define.second->InitSlot(r, slot);
      RColumnReadersInfo info{fInputColumns, fColumnRegister, fIsDefine.data(), fLoopManager->GetDSValuePtrs(),
                              fLoopManager->GetDataSource()};
      fValues[slot] = MakeColumnReaders(slot, r, ColumnTypes_t{}, info);
      fLastCheckedEntry[slot * CacheLineStep<Long64_t>()] = -1;
   }

   /// Return the (type-erased) address of the value for the given processing slot.
   void *GetValuePtr(unsigned int slot, const std::string &column, const std::string &variation) final
   {
      const auto colIt = std::find(fColNames.begin(), fColNames.end(), column);
      assert(colIt != fColNames.end());
      const auto colIdx = std::distance(fColNames.begin(), colIt);

      const auto varIt = std::find(fVariationNames.begin(), fVariationNames.end(), variation);
      assert(varIt != fVariationNames.end());
      const auto varIdx = std::distance(fVariationNames.begin(), varIt);

      return GetValuePtr(slot, colIdx, varIdx);
   }

   /// Update the value at the address returned by GetValuePtr with the content corresponding to the given entry
   void Update(unsigned int slot, Long64_t entry) final
   {
      if (entry != fLastCheckedEntry[slot * CacheLineStep<Long64_t>()]) {
         // evaluate this filter, cache the result
         UpdateHelper(slot, entry, ColumnTypes_t{}, TypeInd_t{});
         fLastCheckedEntry[slot * CacheLineStep<Long64_t>()] = entry;
      }
   }

   const std::type_info &GetTypeId() const { return typeid(ret_type); }

   /// Clean-up operations to be performed at the end of a task.
   void FinalizeSlot(unsigned int slot) final
   {
      for (auto &v : fValues[slot])
         v.reset();
   }
};

} // namespace RDF
} // namespace Internal
} // namespace ROOT

#endif // ROOT_RDF_RVARIATION
<commit_msg>[DF] Remove useless protection against vector<bool><commit_after>// Author: Enrico Guiraud, CERN 10/2021

/*************************************************************************
 * Copyright (C) 1995-2021, Rene Brun and Fons Rademakers.               *
 * All rights reserved.                                                  *
 *                                                                       *
 * For the licensing terms see $ROOTSYS/LICENSE.                         *
 * For the list of contributors see $ROOTSYS/README/CREDITS.             *
 *************************************************************************/

#ifndef ROOT_RDF_RVARIATION
#define ROOT_RDF_RVARIATION

#include "Utils.hxx" // IsRVec
#include "ColumnReaderUtils.hxx"
#include "RColumnReaderBase.hxx"
#include "RLoopManager.hxx"
#include "RVariationBase.hxx"

#include <ROOT/RStringView.hxx>
#include <ROOT/TypeTraits.hxx>
#include <Rtypes.h> // R__CLING_PTRCHECK, Long64_t

#include <array>
#include <deque>
#include <map>
#include <string>
#include <type_traits> // std::is_same, std::conditional_t
#include <utility>     // std::index_sequence
#include <vector>

class TTreeReader;

namespace ROOT {

namespace RDF {
class RDataSource;
}

namespace Internal {
namespace RDF {

using namespace ROOT::TypeTraits;

template <typename T>
bool ResultsSizeEq(const RVec<RVec<T>> &results, std::size_t expected, std::size_t nColumns)
{
   if (nColumns == 1)
      return results.size() == expected;

   return std::all_of(results.begin(), results.end(), [&expected](const RVec<T> &v) { return v.size() == expected; });
}

template <typename T>
bool ResultsSizeEq(const RVec<T> &results, std::size_t expected, std::size_t nColumns)
{
   assert(nColumns == 1);
   (void)nColumns;

   return results.size() == expected;
}

template <typename T>
std::size_t GetNVariations(const RVec<RVec<T>> &results)
{
   assert(!results.empty());
   return results[0].size();
}

template <typename T>
std::size_t GetNVariations(const RVec<T> &results)
{
   return results.size();
}

template <typename F>
class R__CLING_PTRCHECK(off) RVariation final : public RVariationBase {
   using ColumnTypes_t = typename CallableTraits<F>::arg_types;
   using TypeInd_t = std::make_index_sequence<ColumnTypes_t::list_size>;
   using ret_type = typename CallableTraits<F>::ret_type;

   F fExpression;
   std::vector<ret_type> fLastResults;

   /// Column readers per slot and per input column
   std::vector<std::array<std::unique_ptr<RColumnReaderBase>, ColumnTypes_t::list_size>> fValues;

   template <typename... ColTypes, std::size_t... S>
   void UpdateHelper(unsigned int slot, Long64_t entry, TypeList<ColTypes...>, std::index_sequence<S...>)
   {
      // fExpression must return an RVec<T>
      const auto &results = fExpression(fValues[slot][S]->template Get<ColTypes>(entry)...);
      if (!ResultsSizeEq(results, fVariationNames.size(), fColNames.size())) {
         std::string variationName = fVariationNames[0].substr(0, fVariationNames[0].find_first_of(':'));
         throw std::runtime_error("The evaluation of the expression for variation \"" + variationName +
                                  "\" resulted in " + std::to_string(GetNVariations(results)) + " values, but " +
                                  std::to_string(fVariationNames.size()) + " were expected.");
      }
      // Assign into fLastResults without changing the addresses of its elements (we gave those addresses away in
      // GetValuePtr)
      fLastResults[slot * CacheLineStep<ret_type>()].assign(results.begin(), results.end());

      // silence "unused parameter" warnings in gcc
      (void)slot;
      (void)entry;
   }

   // This overload is for the case of a single column and ret_type != RVec<RVec<...>> -- the colIdx is ignored.
   template <typename U = typename ret_type::value_type>
   std::enable_if_t<!IsRVec<U>::value, void *>
   GetValuePtr(unsigned int slot, std::size_t /*colIdx*/, std::size_t varIdx)
   {
      auto &value = fLastResults[slot * CacheLineStep<ret_type>()][varIdx];
      return static_cast<void *>(&value);
   }

   // This overload is for the case of ret_type == RVec<RVec<...>>
   template <typename U = typename ret_type::value_type>
   std::enable_if_t<IsRVec<U>::value, void *> GetValuePtr(unsigned int slot, std::size_t colIdx, std::size_t varIdx)
   {
      if (fColNames.size() == 1) {
         auto &value = fLastResults[slot * CacheLineStep<ret_type>()][varIdx];
         return static_cast<void *>(&value);
      }

      auto &value = fLastResults[slot * CacheLineStep<ret_type>()][colIdx][varIdx];
      return static_cast<void *>(&value);
   }

public:
   RVariation(const std::vector<std::string> &colNames, std::string_view variationName, F expression,
              const std::vector<std::string> &variationTags, std::string_view type, const RColumnRegister &defines,
              RLoopManager &lm, const ColumnNames_t &inputColNames)
      : RVariationBase(colNames, variationName, variationTags, type, defines, lm, inputColNames),
        fExpression(std::move(expression)), fLastResults(lm.GetNSlots() * RDFInternal::CacheLineStep<ret_type>()),
        fValues(lm.GetNSlots())
   {
      for (auto i = 0u; i < lm.GetNSlots(); ++i)
         fLastResults[i * RDFInternal::CacheLineStep<ret_type>()].resize(fVariationNames.size());
   }

   RVariation(const RVariation &) = delete;
   RVariation &operator=(const RVariation &) = delete;

   void InitSlot(TTreeReader *r, unsigned int slot) final
   {
      for (auto &define : fColumnRegister.GetColumns())
         define.second->InitSlot(r, slot);
      RColumnReadersInfo info{fInputColumns, fColumnRegister, fIsDefine.data(), fLoopManager->GetDSValuePtrs(),
                              fLoopManager->GetDataSource()};
      fValues[slot] = MakeColumnReaders(slot, r, ColumnTypes_t{}, info);
      fLastCheckedEntry[slot * CacheLineStep<Long64_t>()] = -1;
   }

   /// Return the (type-erased) address of the value for the given processing slot.
   void *GetValuePtr(unsigned int slot, const std::string &column, const std::string &variation) final
   {
      const auto colIt = std::find(fColNames.begin(), fColNames.end(), column);
      assert(colIt != fColNames.end());
      const auto colIdx = std::distance(fColNames.begin(), colIt);

      const auto varIt = std::find(fVariationNames.begin(), fVariationNames.end(), variation);
      assert(varIt != fVariationNames.end());
      const auto varIdx = std::distance(fVariationNames.begin(), varIt);

      return GetValuePtr(slot, colIdx, varIdx);
   }

   /// Update the value at the address returned by GetValuePtr with the content corresponding to the given entry
   void Update(unsigned int slot, Long64_t entry) final
   {
      if (entry != fLastCheckedEntry[slot * CacheLineStep<Long64_t>()]) {
         // evaluate this filter, cache the result
         UpdateHelper(slot, entry, ColumnTypes_t{}, TypeInd_t{});
         fLastCheckedEntry[slot * CacheLineStep<Long64_t>()] = entry;
      }
   }

   const std::type_info &GetTypeId() const { return typeid(ret_type); }

   /// Clean-up operations to be performed at the end of a task.
   void FinalizeSlot(unsigned int slot) final
   {
      for (auto &v : fValues[slot])
         v.reset();
   }
};

} // namespace RDF
} // namespace Internal
} // namespace ROOT

#endif // ROOT_RDF_RVARIATION
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/mlir_gpu/passes.h"

#include "absl/memory/memory.h"
#include "llvm/ADT/SetVector.h"
#include "mlir/Dialect/GPU/GPUDialect.h"  // from @llvm-project
#include "mlir/Dialect/GPU/ParallelLoopMapper.h"  // from @llvm-project
#include "mlir/Dialect/SCF/SCF.h"  // from @llvm-project
#include "mlir/Dialect/SCF/Transforms.h"  // from @llvm-project
#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
#include "mlir/IR/BlockAndValueMapping.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
#include "mlir/Transforms/LoopUtils.h"  // from @llvm-project
#include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"

namespace xla {
namespace mlir_gpu {
namespace {

#define GEN_PASS_CLASSES
#include "tensorflow/compiler/xla/service/mlir_gpu/passes.h.inc"

struct FusionOpRemoverPass : FusionOpRemoverPassBase<FusionOpRemoverPass> {
  void runOnFunction() override {
    getFunction().walk([&](mlir::lmhlo::FusionOp op) {
      mlir::OpBuilder builder(op);
      // FusionOp has a single region with a single block, so we can just walk
      // over it and clone operations to the outside.
      mlir::BlockAndValueMapping mapping;
      for (auto& nested_op : op.region().front().without_terminator()) {
        auto clone = builder.clone(nested_op, mapping);
        for (auto pair :
             llvm::zip(nested_op.getResults(), clone->getResults())) {
          mapping.map(std::get<0>(pair), std::get<1>(pair));
        }
      }
      op.erase();
    });
  }
};

template <typename EffectTy>
bool HasEffectsOnValue(mlir::Value value, mlir::Operation* op) {
  auto mem_effects_interface =
      mlir::dyn_cast_or_null<mlir::MemoryEffectOpInterface>(op);
  if (!mem_effects_interface) {
    return false;
  }
  llvm::SmallVector<mlir::MemoryEffects::EffectInstance, 2> effects;
  mem_effects_interface.getEffects(effects);
  return llvm::any_of(effects,
                      [op](const mlir::MemoryEffects::EffectInstance& effect) {
                        return mlir::isa<EffectTy>(effect.getEffect());
                      });
}

struct StoreForwardingPass : StoreForwardingPassBase<StoreForwardingPass> {
  mlir::StoreOp findStore(mlir::Operation* op,
                          std::function<bool(mlir::StoreOp)> matches) {
    // Search from op upwards in the current block.
    mlir::Block* block = op->getBlock();
    auto startFromIt =
        std::find_if(block->rbegin(), block->rend(),
                     [op](mlir::Operation& other) { return &other == op; });
    for (auto storeOpIt = startFromIt; storeOpIt != block->rend();
         ++storeOpIt) {
      auto storeOp = llvm::dyn_cast<mlir::StoreOp>(&*(storeOpIt));
      if (!storeOp || !matches(storeOp)) {
        continue;
      }

      return storeOp;
    }
    // No store operation found. Continue search outside of the parallel
    // loop if block is in a parallel loop.
    if (auto parallelOp =
            llvm::dyn_cast<mlir::scf::ParallelOp>(block->getParentOp())) {
      return findStore(parallelOp.getOperation(), matches);
    }
    return {};
  }

  // Recursively search defining ops for AllocOp. Return either AllocOp if it is
  // found or nullptr.
  mlir::Operation* SearchAllocOp(mlir::Value memref) {
    mlir::Operation* defOp = memref.getDefiningOp();
    while (auto subviewOp = mlir::dyn_cast_or_null<mlir::SubViewOp>(defOp)) {
      defOp = subviewOp.source().getDefiningOp();
    }
    return HasEffectsOnValue<mlir::MemoryEffects::Allocate>(memref, defOp)
               ? defOp
               : nullptr;
  }

  // Retrieves AllocOp from the cache or actually looks for it.
  mlir::Operation* GetAllocOp(
      mlir::Value memref,
      llvm::DenseMap<mlir::Value, mlir::Operation*>* memrefToAllocOp) {
    auto allocOpIt = memrefToAllocOp->find(memref);
    if (allocOpIt != memrefToAllocOp->end()) {
      return allocOpIt->second;
    }
    mlir::Operation* allocOp = SearchAllocOp(memref);
    memrefToAllocOp->insert({memref, allocOp});
    return allocOp;
  }

  void runOnFunction() override {
    llvm::DenseMap<mlir::Value, mlir::Operation*> memrefToAllocOp;

    getFunction().walk([&](mlir::LoadOp loadOp) {
      auto storeOp = findStore(loadOp, [&](mlir::StoreOp storeOp) {
        mlir::Operation* storeOpAlloc =
            GetAllocOp(storeOp.memref(), &memrefToAllocOp);
        mlir::Operation* loadOpAlloc =
            GetAllocOp(loadOp.memref(), &memrefToAllocOp);
        return storeOpAlloc && loadOpAlloc && (storeOpAlloc == loadOpAlloc);
      });
      if (!storeOp) {
        return;
      }
      auto storeIndices = storeOp.getIndices();
      auto loadIndices = loadOp.getIndices();
      if (!std::equal(storeIndices.begin(), storeIndices.end(),
                      loadIndices.begin(), loadIndices.end())) {
        return;
      }
      loadOp.replaceAllUsesWith(storeOp.getValueToStore());
      loadOp.erase();
    });
  }
};

struct DeadTempBufferRemovalPass
    : DeadTempBufferRemovalPassBase<DeadTempBufferRemovalPass> {
  bool operationConsideredDead(mlir::Operation* op) {
    for (auto result : op->getResults()) {
      if (!llvm::all_of(result.getUsers(), [&](mlir::Operation* op) {
            // Store and Dealloc is OK.
            if (llvm::isa<mlir::StoreOp, mlir::DeallocOp>(op)) {
              return true;
            }
            // Load without uses is also ok.
            if (auto loadOp = llvm::dyn_cast<mlir::LoadOp>(op)) {
              return loadOp.use_empty();
            }
            // Subview is ok if it is dead itself.
            if (llvm::isa<mlir::SubViewOp>(op)) {
              return operationConsideredDead(op);
            }
            return false;
          })) {
        return false;
      }
    }
    return true;
  }

  void recursiveErase(mlir::Operation* op,
                      llvm::SmallVectorImpl<mlir::Operation*>* erase_list) {
    for (auto result : op->getResults()) {
      for (auto user : llvm::make_early_inc_range(result.getUsers())) {
        recursiveErase(user, erase_list);
      }
    }
    erase_list->push_back(op);
  }

  void runOnFunction() override {
    llvm::SmallVector<mlir::Operation*, 8> dead_ops;
    getFunction().walk([&](mlir::Operation* op) {
      if (op->getNumResults() != 1 ||
          !HasEffectsOnValue<mlir::MemoryEffects::Allocate>(op->getResult(0),
                                                            op)) {
        return;
      }
      if (!operationConsideredDead(op)) {
        return;
      }

      // TODO(herhut): There should be a generic helper for this.
      recursiveErase(op, &dead_ops);
    });
    for (auto op : dead_ops) {
      op->erase();
    }
  }
};

struct RewriteKernelSignaturePass
    : RewriteKernelSignaturePassBase<RewriteKernelSignaturePass> {
  void runOnFunction() override {
    mlir::FuncOp func = getFunction();
    mlir::ModuleOp module = func.getParentOfType<mlir::ModuleOp>();
    getFunction().walk([&](mlir::gpu::LaunchFuncOp launchOp) {
      mlir::gpu::GPUFuncOp kernel =
          module.lookupSymbol<mlir::gpu::GPUFuncOp>(launchOp.kernel());

      if (kernel.getNumFuncArguments() !=
          func.getNumArguments() + func.getNumResults()) {
        kernel.emitError()
            << "number of kernel arguments does not match number"
            << "of arguments and results of surrounding function";
        signalPassFailure();
        return;
      }
      if (!llvm::hasSingleElement(func)) {
        func.emitError() << "surrounding function has more than one block";
        signalPassFailure();
        return;
      }

      // Compute a map from function arguments to kernel function operands.
      mlir::BlockAndValueMapping func_to_kernel;
      for (mlir::BlockArgument arg : func.getArguments()) {
        for (int i = 0, e = launchOp.getNumKernelOperands(); i < e; ++i) {
          if (launchOp.getKernelOperand(i) == arg) {
            func_to_kernel.map(arg, kernel.getArgument(i));
            break;
          }
        }
      }
      // Also add function results that are computed by the launch.
      mlir::Operation* returnOp = func.getBody().back().getTerminator();
      for (mlir::Value result : returnOp->getOperands()) {
        for (int i = 0, e = launchOp.getNumKernelOperands(); i < e; ++i) {
          if (launchOp.getKernelOperand(i) == result) {
            func_to_kernel.map(result, kernel.getArgument(i));
            break;
          }
        }
      }

      // Create a new kernel function with modified signature. It will have the
      // parameters and result types of the original funcion as its parameter
      // type and otherwise will be void.
      auto gpu_module = kernel.getParentOfType<mlir::gpu::GPUModuleOp>();
      mlir::OpBuilder kernel_builder(gpu_module.body());
      auto operand_types = llvm::to_vector<4>(llvm::concat<const mlir::Type>(
          func.getType().getInputs(), func.getType().getResults()));
      auto new_kernel = kernel_builder.create<mlir::gpu::GPUFuncOp>(
          kernel.getLoc(), kernel.getName(),
          kernel_builder.getFunctionType(operand_types, {}));
      new_kernel->setAttr(mlir::gpu::GPUDialect::getKernelFuncAttrName(),
                          kernel_builder.getUnitAttr());

      // Create a map from old kernel argument to new one.
      mlir::BlockAndValueMapping old_kernel_to_new;
      for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
        mlir::Value func_arg = func.getArgument(i);
        mlir::Value new_kernel_arg = new_kernel.getArgument(i);
        mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(func_arg);
        if (!old_kernel_arg) {
          kernel.emitOpError()
              << "argument " << i
              << " to containing function is not an argument to the kernel";
          signalPassFailure();
          return;
        }
        old_kernel_to_new.map(old_kernel_arg, new_kernel_arg);
      }
      for (int i = 0, e = returnOp->getNumOperands(); i < e; ++i) {
        mlir::Value ret_op = returnOp->getOperand(i);
        mlir::Value new_kernel_arg =
            new_kernel.getArgument(func.getNumArguments() + i);
        mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(ret_op);
        if (!old_kernel_arg) {
          kernel.emitOpError()
              << "result " << i
              << " of containing function is not an argument to the kernel";
          signalPassFailure();
          return;
        }
        old_kernel_to_new.map(old_kernel_arg, new_kernel_arg);
      }
      // Steal the body by appending the blocks and inserting a branch.
      kernel.body().cloneInto(&new_kernel.getBody(), old_kernel_to_new);
      kernel_builder.setInsertionPointToEnd(&new_kernel.body().front());
      kernel_builder.create<mlir::BranchOp>(
          new_kernel.getLoc(), &*std::next(new_kernel.body().begin()));
      // Now create a new launchOp calling the new kernel. We need to forward
      // the arguments of the surrounding function and operands to the return.
      mlir::SmallVector<mlir::Value, 4> new_operands;
      new_operands.reserve(new_kernel.getNumFuncArguments());
      new_operands.append(func.args_begin(), func.args_end());
      new_operands.append(returnOp->operand_begin(), returnOp->operand_end());
      mlir::OpBuilder launch_builder(launchOp);
      launch_builder.create<mlir::gpu::LaunchFuncOp>(
          launchOp.getLoc(), new_kernel, launchOp.getGridSizeOperandValues(),
          launchOp.getBlockSizeOperandValues(), new_operands);
      // Launch does not have results, so we can just erase it. And the kernel
      // also needs to go.
      launchOp.erase();
      kernel.erase();
    });
  }
};

struct MapParallelLoopsPass : MapParallelLoopsPassBase<MapParallelLoopsPass> {
  void runOnFunction() override {
    mlir::greedilyMapParallelSCFToGPU(getFunction().getBody());
  }
};

struct FuseInnerParallelLoopsPass
    : FuseInnerParallelLoopsPassBase<FuseInnerParallelLoopsPass> {
  void runOnFunction() override {
    getFunction().walk([](mlir::scf::ParallelOp op) {
      mlir::scf::naivelyFuseParallelOps(op.region());
    });
  }
};

struct ParallelLoopCollapsingToFirstDimPass
    : ParallelLoopCollapsingToFirstDimPassBase<
          ParallelLoopCollapsingToFirstDimPass> {
  void runOnFunction() override {
    getFunction().walk([&](mlir::scf::ParallelOp op) {
      unsigned num_loops = op.getNumLoops();
      std::vector<unsigned> combinedLoops;
      combinedLoops.reserve(num_loops);
      for (unsigned i = 0; i < num_loops; ++i) {
        combinedLoops.push_back(i);
      }
      mlir::collapseParallelLoops(op, {combinedLoops});
    });
  }
};

}  // namespace

std::unique_ptr<mlir::FunctionPass> createFusionOpRemoverPass() {
  return absl::make_unique<FusionOpRemoverPass>();
}

std::unique_ptr<mlir::FunctionPass> createStoreForwardingPass() {
  return absl::make_unique<StoreForwardingPass>();
}

std::unique_ptr<mlir::FunctionPass> createDeadTempBufferRemovalPass() {
  return absl::make_unique<DeadTempBufferRemovalPass>();
}

std::unique_ptr<mlir::FunctionPass> createRewriteKernelSignaturePass() {
  return absl::make_unique<RewriteKernelSignaturePass>();
}

std::unique_ptr<mlir::FunctionPass> createFuseInnerParallelLoopsPass() {
  return absl::make_unique<FuseInnerParallelLoopsPass>();
}

std::unique_ptr<mlir::FunctionPass> createMapParallelLoopsPass() {
  return absl::make_unique<MapParallelLoopsPass>();
}

std::unique_ptr<mlir::FunctionPass>
createParallelLoopCollapsingToFirstDimPass() {
  return absl::make_unique<ParallelLoopCollapsingToFirstDimPass>();
}

}  // namespace mlir_gpu
}  // namespace xla
<commit_msg>Use mlir::OpState::operator->() to get to methods of mlir::Operation.<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/mlir_gpu/passes.h"

#include "absl/memory/memory.h"
#include "llvm/ADT/SetVector.h"
#include "mlir/Dialect/GPU/GPUDialect.h"  // from @llvm-project
#include "mlir/Dialect/GPU/ParallelLoopMapper.h"  // from @llvm-project
#include "mlir/Dialect/SCF/SCF.h"  // from @llvm-project
#include "mlir/Dialect/SCF/Transforms.h"  // from @llvm-project
#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
#include "mlir/IR/BlockAndValueMapping.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
#include "mlir/Transforms/LoopUtils.h"  // from @llvm-project
#include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"

namespace xla {
namespace mlir_gpu {
namespace {

#define GEN_PASS_CLASSES
#include "tensorflow/compiler/xla/service/mlir_gpu/passes.h.inc"

struct FusionOpRemoverPass : FusionOpRemoverPassBase<FusionOpRemoverPass> {
  void runOnFunction() override {
    getFunction().walk([&](mlir::lmhlo::FusionOp op) {
      mlir::OpBuilder builder(op);
      // FusionOp has a single region with a single block, so we can just walk
      // over it and clone operations to the outside.
      mlir::BlockAndValueMapping mapping;
      for (auto& nested_op : op.region().front().without_terminator()) {
        auto clone = builder.clone(nested_op, mapping);
        for (auto pair :
             llvm::zip(nested_op.getResults(), clone->getResults())) {
          mapping.map(std::get<0>(pair), std::get<1>(pair));
        }
      }
      op.erase();
    });
  }
};

template <typename EffectTy>
bool HasEffectsOnValue(mlir::Value value, mlir::Operation* op) {
  auto mem_effects_interface =
      mlir::dyn_cast_or_null<mlir::MemoryEffectOpInterface>(op);
  if (!mem_effects_interface) {
    return false;
  }
  llvm::SmallVector<mlir::MemoryEffects::EffectInstance, 2> effects;
  mem_effects_interface.getEffects(effects);
  return llvm::any_of(effects,
                      [op](const mlir::MemoryEffects::EffectInstance& effect) {
                        return mlir::isa<EffectTy>(effect.getEffect());
                      });
}

struct StoreForwardingPass : StoreForwardingPassBase<StoreForwardingPass> {
  mlir::StoreOp findStore(mlir::Operation* op,
                          std::function<bool(mlir::StoreOp)> matches) {
    // Search from op upwards in the current block.
    mlir::Block* block = op->getBlock();
    auto startFromIt =
        std::find_if(block->rbegin(), block->rend(),
                     [op](mlir::Operation& other) { return &other == op; });
    for (auto storeOpIt = startFromIt; storeOpIt != block->rend();
         ++storeOpIt) {
      auto storeOp = llvm::dyn_cast<mlir::StoreOp>(&*(storeOpIt));
      if (!storeOp || !matches(storeOp)) {
        continue;
      }

      return storeOp;
    }
    // No store operation found. Continue search outside of the parallel
    // loop if block is in a parallel loop.
    if (auto parallelOp =
            llvm::dyn_cast<mlir::scf::ParallelOp>(block->getParentOp())) {
      return findStore(parallelOp.getOperation(), matches);
    }
    return {};
  }

  // Recursively search defining ops for AllocOp. Return either AllocOp if it is
  // found or nullptr.
  mlir::Operation* SearchAllocOp(mlir::Value memref) {
    mlir::Operation* defOp = memref.getDefiningOp();
    while (auto subviewOp = mlir::dyn_cast_or_null<mlir::SubViewOp>(defOp)) {
      defOp = subviewOp.source().getDefiningOp();
    }
    return HasEffectsOnValue<mlir::MemoryEffects::Allocate>(memref, defOp)
               ? defOp
               : nullptr;
  }

  // Retrieves AllocOp from the cache or actually looks for it.
  mlir::Operation* GetAllocOp(
      mlir::Value memref,
      llvm::DenseMap<mlir::Value, mlir::Operation*>* memrefToAllocOp) {
    auto allocOpIt = memrefToAllocOp->find(memref);
    if (allocOpIt != memrefToAllocOp->end()) {
      return allocOpIt->second;
    }
    mlir::Operation* allocOp = SearchAllocOp(memref);
    memrefToAllocOp->insert({memref, allocOp});
    return allocOp;
  }

  void runOnFunction() override {
    llvm::DenseMap<mlir::Value, mlir::Operation*> memrefToAllocOp;

    getFunction().walk([&](mlir::LoadOp loadOp) {
      auto storeOp = findStore(loadOp, [&](mlir::StoreOp storeOp) {
        mlir::Operation* storeOpAlloc =
            GetAllocOp(storeOp.memref(), &memrefToAllocOp);
        mlir::Operation* loadOpAlloc =
            GetAllocOp(loadOp.memref(), &memrefToAllocOp);
        return storeOpAlloc && loadOpAlloc && (storeOpAlloc == loadOpAlloc);
      });
      if (!storeOp) {
        return;
      }
      auto storeIndices = storeOp.getIndices();
      auto loadIndices = loadOp.getIndices();
      if (!std::equal(storeIndices.begin(), storeIndices.end(),
                      loadIndices.begin(), loadIndices.end())) {
        return;
      }
      loadOp.replaceAllUsesWith(storeOp.getValueToStore());
      loadOp.erase();
    });
  }
};

struct DeadTempBufferRemovalPass
    : DeadTempBufferRemovalPassBase<DeadTempBufferRemovalPass> {
  bool operationConsideredDead(mlir::Operation* op) {
    for (auto result : op->getResults()) {
      if (!llvm::all_of(result.getUsers(), [&](mlir::Operation* op) {
            // Store and Dealloc is OK.
            if (llvm::isa<mlir::StoreOp, mlir::DeallocOp>(op)) {
              return true;
            }
            // Load without uses is also ok.
            if (auto loadOp = llvm::dyn_cast<mlir::LoadOp>(op)) {
              return loadOp.use_empty();
            }
            // Subview is ok if it is dead itself.
            if (llvm::isa<mlir::SubViewOp>(op)) {
              return operationConsideredDead(op);
            }
            return false;
          })) {
        return false;
      }
    }
    return true;
  }

  void recursiveErase(mlir::Operation* op,
                      llvm::SmallVectorImpl<mlir::Operation*>* erase_list) {
    for (auto result : op->getResults()) {
      for (auto user : llvm::make_early_inc_range(result.getUsers())) {
        recursiveErase(user, erase_list);
      }
    }
    erase_list->push_back(op);
  }

  void runOnFunction() override {
    llvm::SmallVector<mlir::Operation*, 8> dead_ops;
    getFunction().walk([&](mlir::Operation* op) {
      if (op->getNumResults() != 1 ||
          !HasEffectsOnValue<mlir::MemoryEffects::Allocate>(op->getResult(0),
                                                            op)) {
        return;
      }
      if (!operationConsideredDead(op)) {
        return;
      }

      // TODO(herhut): There should be a generic helper for this.
      recursiveErase(op, &dead_ops);
    });
    for (auto op : dead_ops) {
      op->erase();
    }
  }
};

struct RewriteKernelSignaturePass
    : RewriteKernelSignaturePassBase<RewriteKernelSignaturePass> {
  void runOnFunction() override {
    mlir::FuncOp func = getFunction();
    mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
    getFunction().walk([&](mlir::gpu::LaunchFuncOp launchOp) {
      mlir::gpu::GPUFuncOp kernel =
          module.lookupSymbol<mlir::gpu::GPUFuncOp>(launchOp.kernel());

      if (kernel.getNumFuncArguments() !=
          func.getNumArguments() + func.getNumResults()) {
        kernel.emitError()
            << "number of kernel arguments does not match number"
            << "of arguments and results of surrounding function";
        signalPassFailure();
        return;
      }
      if (!llvm::hasSingleElement(func)) {
        func.emitError() << "surrounding function has more than one block";
        signalPassFailure();
        return;
      }

      // Compute a map from function arguments to kernel function operands.
      mlir::BlockAndValueMapping func_to_kernel;
      for (mlir::BlockArgument arg : func.getArguments()) {
        for (int i = 0, e = launchOp.getNumKernelOperands(); i < e; ++i) {
          if (launchOp.getKernelOperand(i) == arg) {
            func_to_kernel.map(arg, kernel.getArgument(i));
            break;
          }
        }
      }
      // Also add function results that are computed by the launch.
      mlir::Operation* returnOp = func.getBody().back().getTerminator();
      for (mlir::Value result : returnOp->getOperands()) {
        for (int i = 0, e = launchOp.getNumKernelOperands(); i < e; ++i) {
          if (launchOp.getKernelOperand(i) == result) {
            func_to_kernel.map(result, kernel.getArgument(i));
            break;
          }
        }
      }

      // Create a new kernel function with modified signature. It will have the
      // parameters and result types of the original funcion as its parameter
      // type and otherwise will be void.
      auto gpu_module = kernel->getParentOfType<mlir::gpu::GPUModuleOp>();
      mlir::OpBuilder kernel_builder(gpu_module.body());
      auto operand_types = llvm::to_vector<4>(llvm::concat<const mlir::Type>(
          func.getType().getInputs(), func.getType().getResults()));
      auto new_kernel = kernel_builder.create<mlir::gpu::GPUFuncOp>(
          kernel.getLoc(), kernel.getName(),
          kernel_builder.getFunctionType(operand_types, {}));
      new_kernel->setAttr(mlir::gpu::GPUDialect::getKernelFuncAttrName(),
                          kernel_builder.getUnitAttr());

      // Create a map from old kernel argument to new one.
      mlir::BlockAndValueMapping old_kernel_to_new;
      for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
        mlir::Value func_arg = func.getArgument(i);
        mlir::Value new_kernel_arg = new_kernel.getArgument(i);
        mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(func_arg);
        if (!old_kernel_arg) {
          kernel.emitOpError()
              << "argument " << i
              << " to containing function is not an argument to the kernel";
          signalPassFailure();
          return;
        }
        old_kernel_to_new.map(old_kernel_arg, new_kernel_arg);
      }
      for (int i = 0, e = returnOp->getNumOperands(); i < e; ++i) {
        mlir::Value ret_op = returnOp->getOperand(i);
        mlir::Value new_kernel_arg =
            new_kernel.getArgument(func.getNumArguments() + i);
        mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(ret_op);
        if (!old_kernel_arg) {
          kernel.emitOpError()
              << "result " << i
              << " of containing function is not an argument to the kernel";
          signalPassFailure();
          return;
        }
        old_kernel_to_new.map(old_kernel_arg, new_kernel_arg);
      }
      // Steal the body by appending the blocks and inserting a branch.
      kernel.body().cloneInto(&new_kernel.getBody(), old_kernel_to_new);
      kernel_builder.setInsertionPointToEnd(&new_kernel.body().front());
      kernel_builder.create<mlir::BranchOp>(
          new_kernel.getLoc(), &*std::next(new_kernel.body().begin()));
      // Now create a new launchOp calling the new kernel. We need to forward
      // the arguments of the surrounding function and operands to the return.
      mlir::SmallVector<mlir::Value, 4> new_operands;
      new_operands.reserve(new_kernel.getNumFuncArguments());
      new_operands.append(func.args_begin(), func.args_end());
      new_operands.append(returnOp->operand_begin(), returnOp->operand_end());
      mlir::OpBuilder launch_builder(launchOp);
      launch_builder.create<mlir::gpu::LaunchFuncOp>(
          launchOp.getLoc(), new_kernel, launchOp.getGridSizeOperandValues(),
          launchOp.getBlockSizeOperandValues(), new_operands);
      // Launch does not have results, so we can just erase it. And the kernel
      // also needs to go.
      launchOp.erase();
      kernel.erase();
    });
  }
};

struct MapParallelLoopsPass : MapParallelLoopsPassBase<MapParallelLoopsPass> {
  void runOnFunction() override {
    mlir::greedilyMapParallelSCFToGPU(getFunction().getBody());
  }
};

struct FuseInnerParallelLoopsPass
    : FuseInnerParallelLoopsPassBase<FuseInnerParallelLoopsPass> {
  void runOnFunction() override {
    getFunction().walk([](mlir::scf::ParallelOp op) {
      mlir::scf::naivelyFuseParallelOps(op.region());
    });
  }
};

struct ParallelLoopCollapsingToFirstDimPass
    : ParallelLoopCollapsingToFirstDimPassBase<
          ParallelLoopCollapsingToFirstDimPass> {
  void runOnFunction() override {
    getFunction().walk([&](mlir::scf::ParallelOp op) {
      unsigned num_loops = op.getNumLoops();
      std::vector<unsigned> combinedLoops;
      combinedLoops.reserve(num_loops);
      for (unsigned i = 0; i < num_loops; ++i) {
        combinedLoops.push_back(i);
      }
      mlir::collapseParallelLoops(op, {combinedLoops});
    });
  }
};

}  // namespace

std::unique_ptr<mlir::FunctionPass> createFusionOpRemoverPass() {
  return absl::make_unique<FusionOpRemoverPass>();
}

std::unique_ptr<mlir::FunctionPass> createStoreForwardingPass() {
  return absl::make_unique<StoreForwardingPass>();
}

std::unique_ptr<mlir::FunctionPass> createDeadTempBufferRemovalPass() {
  return absl::make_unique<DeadTempBufferRemovalPass>();
}

std::unique_ptr<mlir::FunctionPass> createRewriteKernelSignaturePass() {
  return absl::make_unique<RewriteKernelSignaturePass>();
}

std::unique_ptr<mlir::FunctionPass> createFuseInnerParallelLoopsPass() {
  return absl::make_unique<FuseInnerParallelLoopsPass>();
}

std::unique_ptr<mlir::FunctionPass> createMapParallelLoopsPass() {
  return absl::make_unique<MapParallelLoopsPass>();
}

std::unique_ptr<mlir::FunctionPass>
createParallelLoopCollapsingToFirstDimPass() {
  return absl::make_unique<ParallelLoopCollapsingToFirstDimPass>();
}

}  // namespace mlir_gpu
}  // namespace xla
<|endoftext|>