// Copyright (c) 1993-2023, NVIDIA CORPORATION. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//  * Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//  * Neither the name of NVIDIA CORPORATION nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``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 THE COPYRIGHT OWNER OR
// CONTRIBUTORS 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 <assert.h>
#include <nvvm.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>

/* Two levels of indirection to stringify LIBDEVICE_MAJOR_VERSION and
 * LIBDEVICE_MINOR_VERSION correctly. */
#define getLibDeviceName()               _getLibDeviceName(LIBDEVICE_MAJOR_VERSION, LIBDEVICE_MINOR_VERSION)
#define _getLibDeviceName(MAJOR, MINOR)  __getLibDeviceName(MAJOR, MINOR)
#define __getLibDeviceName(MAJOR, MINOR) ("/libdevice/libdevice." #MAJOR #MINOR ".bc")

#define getLibnvvmHome()            _getLibnvvmHome(LIBNVVM_HOME)
#define _getLibnvvmHome(NVVM_HOME)  __getLibnvvmHome(NVVM_HOME)
#define __getLibnvvmHome(NVVM_HOME) (#NVVM_HOME)

typedef enum {
    PTXGEN_SUCCESS                    = 0x0000,
    PTXGEN_FILE_IO_ERROR              = 0x0001,
    PTXGEN_BAD_ALLOC_ERROR            = 0x0002,
    PTXGEN_LIBNVVM_COMPILATION_ERROR  = 0x0004,
    PTXGEN_LIBNVVM_ERROR              = 0x0008,
    PTXGEN_INVALID_USAGE              = 0x0010,
    PTXGEN_LIBNVVM_HOME_UNDEFINED     = 0x0020,
    PTXGEN_LIBNVVM_VERIFICATION_ERROR = 0x0040
} PTXGenStatus;

typedef enum { PTXGEN_INPUT_PROGRAM, PTXGEN_INPUT_LIBDEVICE } PTXGENInput;

static PTXGenStatus getLibDevicePath(char **buffer)
{
    assert(buffer);

    const char *libnvvmPath = getLibnvvmHome();
    if (libnvvmPath == NULL) {
        fprintf(stderr, "The environment variable LIBNVVM_HOME undefined\n");
        return PTXGEN_LIBNVVM_HOME_UNDEFINED;
    }

    const char *libdevice = getLibDeviceName();
    *buffer               = malloc(strlen(libnvvmPath) + strlen(libdevice) + 1);
    if (*buffer == NULL) {
        fprintf(stderr, "Failed to allocate memory\n");
        return PTXGEN_BAD_ALLOC_ERROR;
    }

    // Concatenate libnvvmPath with libdevice.
    *buffer = strcat(strcpy(*buffer, libnvvmPath), libdevice);

    return PTXGEN_SUCCESS;
}

static PTXGenStatus addFileToProgram(const char *filename, nvvmProgram prog, PTXGENInput inputType)
{
    assert(filename);

    // Open the input file.
    FILE *f = fopen(filename, "rb");
    if (f == NULL) {
        fprintf(stderr, "Failed to open %s\n", filename);
        return PTXGEN_FILE_IO_ERROR;
    }

    // Allocate a buffer for the input.
    struct stat fileStat;
    fstat(fileno(f), &fileStat);
    char *buffer = malloc(fileStat.st_size);
    if (buffer == NULL) {
        fprintf(stderr, "Failed to allocate memory\n");
        return PTXGEN_BAD_ALLOC_ERROR;
    }
    const size_t size = fread(buffer, 1, fileStat.st_size, f);
    if (ferror(f)) {
        fprintf(stderr, "Failed to read %s\n", filename);
        fclose(f);
        free(buffer);
        return PTXGEN_FILE_IO_ERROR;
    }
    fclose(f);

    // Add the input to the libNVVM program instance.
    nvvmResult result;
    if (inputType == PTXGEN_INPUT_LIBDEVICE)
        result = nvvmLazyAddModuleToProgram(prog, buffer, size, filename);
    else
        result = nvvmAddModuleToProgram(prog, buffer, size, filename);
    if (result != NVVM_SUCCESS) {
        fprintf(stderr, "Failed to add the module %s to the compilation unit\n", filename);
        free(buffer);
        return PTXGEN_LIBNVVM_ERROR;
    }

    free(buffer);
    return PTXGEN_SUCCESS;
}

// Read the nvvmProgram compilation log.
static PTXGenStatus dumpCompilationLog(nvvmProgram prog)
{
    size_t       logSize;
    PTXGenStatus status = PTXGEN_SUCCESS;
    if (nvvmGetProgramLogSize(prog, &logSize) != NVVM_SUCCESS) {
        fprintf(stderr, "Failed to get the compilation log size\n");
        status = PTXGEN_LIBNVVM_ERROR;
    }
    else {
        char *log = malloc(logSize);
        if (log == NULL) {
            fprintf(stderr, "Failed to allocate memory\n");
            status = PTXGEN_BAD_ALLOC_ERROR;
        }
        else if (nvvmGetProgramLog(prog, log) != NVVM_SUCCESS) {
            fprintf(stderr, "Failed to get the compilation log\n");
            status = PTXGEN_LIBNVVM_ERROR;
        }
        else {
            fprintf(stderr, "%s\n", log);
        }
        free(log);
    }
    return status;
}

static PTXGenStatus
generatePTX(unsigned numOptions, const char **options, unsigned numFilenames, const char **filenames)
{
    // Create the compilation unit (the libNVVM program instance).
    nvvmProgram prog;
    if (nvvmCreateProgram(&prog) != NVVM_SUCCESS) {
        fprintf(stderr, "Failed to create the compilation unit\n");
        return PTXGEN_LIBNVVM_ERROR;
    }

    // Add the module to the compilation unit.
    for (unsigned i = 0; i < numFilenames; ++i) {
        PTXGenStatus status = addFileToProgram(filenames[i], prog, PTXGEN_INPUT_PROGRAM);
        if (status != PTXGEN_SUCCESS) {
            nvvmDestroyProgram(&prog);
            return status;
        }
    }

    // Verify the compilation unit.
    if (nvvmVerifyProgram(prog, numOptions, options) != NVVM_SUCCESS) {
        fprintf(stderr, "Failed to verify the compilation unit\n");
        return PTXGEN_LIBNVVM_VERIFICATION_ERROR;
    }

    // Add libdevice to the libNVVM program instance.
    char        *libDeviceName;
    PTXGenStatus status = getLibDevicePath(&libDeviceName);
    if (status != PTXGEN_SUCCESS) {
        nvvmDestroyProgram(&prog);
        return status;
    }
    status = addFileToProgram(libDeviceName, prog, PTXGEN_INPUT_LIBDEVICE);
    free(libDeviceName);
    if (status != PTXGEN_SUCCESS) {
        nvvmDestroyProgram(&prog);
        return status;
    }

    // Display the compilation log.
    status |= dumpCompilationLog(prog);
    if (status & PTXGEN_LIBNVVM_VERIFICATION_ERROR) {
        nvvmDestroyProgram(&prog);
        return status;
    }

    // Compile the compilation unit.
    if (nvvmCompileProgram(prog, numOptions, options) != NVVM_SUCCESS) {
        fprintf(stderr, "Failed to generate PTX from the compilation unit\n");
        status |= PTXGEN_LIBNVVM_COMPILATION_ERROR;
    }
    else {
        size_t ptxSize;
        if (nvvmGetCompiledResultSize(prog, &ptxSize) != NVVM_SUCCESS) {
            fprintf(stderr, "Failed to get the PTX output size\n");
            status |= PTXGEN_LIBNVVM_ERROR;
        }
        else {
            char *ptx = malloc(ptxSize);
            if (ptx == NULL) {
                fprintf(stderr, "Failed to allocate memory\n");
                status |= PTXGEN_BAD_ALLOC_ERROR;
            }
            else if (nvvmGetCompiledResult(prog, ptx) != NVVM_SUCCESS) {
                fprintf(stderr, "Failed to get the PTX output\n");
                status |= PTXGEN_LIBNVVM_ERROR;
            }
            else {
                fprintf(stdout, "%s\n", ptx);
            }
            free(ptx);
        }
    }

    status |= dumpCompilationLog(prog);

    // Release the resources.
    nvvmDestroyProgram(&prog);

    return status;
}

static void showUsage(void)
{
    fprintf(stderr,
            "Usage: ptxgen [OPTION]... [FILE]...\n"
            "  [FILE] could be a .bc file or a .ll file\n");
}

int main(int argc, char *argv[])
{
    PTXGenStatus status = PTXGEN_SUCCESS;

    if (argc == 1) {
        showUsage();
        return PTXGEN_INVALID_USAGE;
    }

    // Extract libNVVM options and the input file names.
    unsigned     numOptions = 0, numFilenames = 0;
    const char **options   = malloc((argc - 1) * sizeof(char *));
    const char **filenames = malloc((argc - 1) * sizeof(char *));
    assert(options && filenames);
    for (int i = 1; i < argc; ++i) {
        if (argv[i][0] == '-') {
            options[numOptions] = argv[i];
            ++numOptions;
        }
        else {
            filenames[numFilenames] = argv[i];
            ++numFilenames;
        }
    }

    if (numFilenames == 0) {
        showUsage();
        status = PTXGEN_INVALID_USAGE;
    }
    else {
        // Use libNVVM to generate PTX.
        status = generatePTX(numOptions, options, numFilenames, filenames);
    }

    free(options);
    free(filenames);
    return status;
}
