/normxcorr/trunk

#if defined(_WIN32) || defined(_WIN64)
# include <windows.h>
    typedef UINT8 uint8_t;
    typedef UINT16 uint16_t;
    typedef UINT32 uint32_t;
    typedef INT8 int8_t;
    typedef INT16 int16_t;
    typedef INT32 int32_t;
#else
# include <stdint.h>
#endif

#include <stdio.h>
#include <stdlib.h>

#include <cuda.h>
#include <cuda_runtime.h>

#include <cublas.h>
#include <cufft.h>

#include <mex.h>

#define BLOCK_SIZE 64

//#include "normxcorr_hw_pack.h"
#include "normxcorr_hw_msg.h"

#include "normxcorr_hw_kernel.cu"


typedef enum {
    ACTION_SETUP = 1,
    ACTION_PREPARE = 2,
    ACTION_COMPUTE_BASE = 10,
    ACTION_COMPUTE_FRAGMENT = 11,
} TAction;

typedef enum {
    ERROR_CUFFT = 1,
    ERROR_CUDA_MALLOC = 2,
    ERROR_MALLOC = 3
} TError;

struct STProcessingState {
    cufftComplex *cuda_base_buffer;	// Stored FFT's of the template image
    cufftComplex *cuda_data_buffer;	// Main computational buffer
    cufftReal *cuda_temp_buffer;	// Temporary buffer for FFT inputs
    cufftReal *cuda_result_buffer;	// Temporary buffer for FFT outputs
    float *cuda_final_buffer;		// Ultimate output
    uint8_t *cuda_input_buffer;		// Input buffer
    
    float *cuda_lsum_buffer;
    float *cuda_denom_buffer;

    int *grid_size;
    uint16_t *cuda_nonzero_items;
    uint16_t *cuda_nonzero_buffer;

    int ncp;			// Number of control points
    int fft_size;		// Matrix Size for FFT (base_size + input_size - 1)
    int fft_size2;		// size * size
    int fft_alloc_size;		// cuda optimized size2
    int fft_inner_size;		// size * (size/2 + 1), R2C/C2R

    int fft_initialized;	// Flag indicating if CUFFT plan is initialized
    cufftHandle cufft_plan;
    cufftHandle cufft_r2c_plan;
    cufftHandle cufft_c2r_plan;
};

typedef struct STProcessingState TProcessingState;
static TProcessingState *pstate = NULL;


static void fftFree(TProcessingState *ps) {
    if (ps->fft_initialized) {
	cufftDestroy(ps->cufft_r2c_plan);
	cufftDestroy(ps->cufft_c2r_plan);
//	cufftDestroy(ps->cufft_plan);
//	cublasShutdown();
	ps->fft_initialized = false;
    }

    if (ps->cuda_base_buffer) {
	free(ps->grid_size);
	free(ps->cuda_nonzero_items);
	
	cudaFree(ps->cuda_lsum_buffer);
	cudaFree(ps->cuda_denom_buffer);
	cudaFree(ps->cuda_nonzero_buffer);
    
	cudaFree(ps->cuda_temp_buffer);
	cudaFree(ps->cuda_final_buffer);
	cudaFree(ps->cuda_result_buffer);
	cudaFree(ps->cuda_data_buffer);
	cudaFree(ps->cuda_base_buffer);
	cudaFree(ps->cuda_input_buffer);
	
	ps->cuda_base_buffer = NULL;
    }
}

#include <unistd.h>
static int fftInit(TProcessingState *ps) {
    cufftResult cufft_err;
    cudaError cuda_err;
    
    fftFree(ps);

//    cublasInit();
//    cufftPlan2d(&ps->cufft_plan, ps->fft_size, ps->fft_size, CUFFT_C2C);

    cufft_err = cufftPlan2d(&ps->cufft_r2c_plan, ps->fft_size, ps->fft_size, CUFFT_R2C);
    if (cufft_err) {
	reportError("Problem initializing c2r plan, cufft code: %i", cufft_err);
	return ERROR_CUFFT;
    }	
    
    cufft_err = cufftPlan2d(&ps->cufft_c2r_plan, ps->fft_size, ps->fft_size, CUFFT_C2R);
    if (cufft_err) {
	reportError("Problem initializing r2c plan, cufft code: %i", cufft_err);
	cufftDestroy(ps->cufft_r2c_plan);
	return ERROR_CUFFT;
    }

    ps->fft_initialized = true;

    cuda_err = cudaMalloc((void**)&ps->cuda_base_buffer, ps->ncp * ps->fft_alloc_size * sizeof(cufftComplex));
    if (cuda_err) {
	reportError("Device memory allocation of %u*%u*cufftComplex bytes for cuda_base_buffer is failed", ps->ncp, ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }

    cuda_err = cudaMalloc((void**)&ps->cuda_data_buffer, ps->fft_alloc_size * sizeof(cufftComplex));
    if (cuda_err) {
	reportError("Device memory allocation of %u*cufftComplex bytes for cuda_data_buffer is failed", ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }

    cuda_err = cudaMalloc((void**)&ps->cuda_input_buffer, ps->ncp * ps->fft_alloc_size * sizeof(uint8_t));
    if (cuda_err) {
	reportError("Device memory allocation of %u*%u*uint8 bytes for cuda_input_buffer is failed", ps->ncp, ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }

    cuda_err = cudaMalloc((void**)&ps->cuda_final_buffer, ps->ncp * ps->fft_alloc_size * sizeof(float));
    if (cuda_err) {
	reportError("Device memory allocation of %u*%u*float bytes for cuda_final_buffer is failed", ps->ncp, ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }
    cudaMemset((void*)ps->cuda_final_buffer, 0, ps->ncp * ps->fft_alloc_size * sizeof(float));

    cuda_err = cudaMalloc((void**)&ps->cuda_result_buffer, ps->fft_alloc_size * sizeof(cufftReal));
    if (cuda_err) {
	reportError("Device memory allocation of %u*cufftReal bytes for cuda_result_buffer is failed", ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }

    cuda_err = cudaMalloc((void**)&ps->cuda_temp_buffer, ps->fft_alloc_size * sizeof(cufftReal));
    if (cuda_err) {
	reportError("Device memory allocation of %u*cufftReal bytes for cuda_temp_buffer is failed", ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }
    cudaMemset((void*)ps->cuda_temp_buffer, 0, ps->fft_alloc_size * sizeof(cufftReal));

    cuda_err = cudaMalloc((void**)&ps->cuda_lsum_buffer, ps->ncp * ps->fft_alloc_size * sizeof(float));
    if (cuda_err) {
	reportError("Device memory allocation of %u*%u*float bytes for cuda_lsum_buffer is failed", ps->ncp, ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }

    cuda_err = cudaMalloc((void**)&ps->cuda_denom_buffer, ps->ncp * ps->fft_alloc_size * sizeof(float));
    if (cuda_err) {
	reportError("Device memory allocation of %u*%u*float bytes for cuda_denom_buffer is failed", ps->ncp, ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }

    cuda_err = cudaMalloc((void**)&ps->cuda_nonzero_buffer, ps->ncp * ps->fft_alloc_size * sizeof(uint16_t));
    if (cuda_err) {
	reportError("Device memory allocation of %u*%u*uint16 bytes for cuda_nonzero_buffer is failed", ps->ncp, ps->fft_alloc_size);
	fftFree(ps);
	return ERROR_CUDA_MALLOC;
    }
    cudaMemset((void*)ps->cuda_nonzero_buffer, 0, ps->ncp * ps->fft_alloc_size * sizeof(uint16_t));
    
    ps->cuda_nonzero_items = (uint16_t*)malloc(ps->ncp * sizeof(uint16_t));
    if (!ps->cuda_nonzero_items) {
	reportError("Host memory allocation of %u*uint16 bytes for cuda_nonzero_items is failed", ps->ncp);
	fftFree(ps);
	return ERROR_MALLOC;
    }
    
    ps->grid_size = (int*)malloc(ps->ncp*sizeof(int));
    if (!ps->grid_size) {
	reportError("Host memory allocation of %u*int bytes for grid_size is failed", ps->ncp);
	fftFree(ps);
	return ERROR_MALLOC;
    }
    
    return 0;
}

static void fftPrepare(TProcessingState *ps) {
    if (ps->fft_initialized) {
	    // Since template and current image have different neighbourhoud sizes
	cudaMemset((void*)ps->cuda_temp_buffer, 0, ps->fft_alloc_size * sizeof(cufftReal));
    }
}



void pstateFree(TProcessingState *ps) {
    if (ps) {
	fftFree(ps);
	free(ps);
    }
}

TProcessingState *pstateInit() {
    TProcessingState *ps;
    
    ps = (TProcessingState*)malloc(sizeof(TProcessingState));
    if (ps) {
	ps->ncp = 0;
	ps->cuda_base_buffer = NULL;
	ps->fft_initialized = false;
    }
    return ps;
}

static inline void *fftUploadBaseData(TProcessingState *ps, int icp, const mxArray *data, const mxArray *lsum, const mxArray *denom, const mxArray *nonzero) {
    uint8_t *dataPtr;

    if (!ps->fft_initialized) {
	reportError("cuFFT engine is not initialized yet");
	return NULL;
    }
    
    int size = ps->fft_size;
    int size2 = size*size;
    int alloc_size = ps->fft_alloc_size;

    int N = mxGetM(data);
    int N2 = N * N;

    dim3 input_block_dim(N, 1, 1);
    dim3 input_grid_dim(N, 1, 1);

    uint8_t *cudaInputPtr = ps->cuda_input_buffer + icp * alloc_size;
    cufftComplex *cudaPtr = ps->cuda_base_buffer + icp * alloc_size;
    cufftReal *cudaRealPtr = ps->cuda_temp_buffer;

    dataPtr = (uint8_t*)mxGetData(data);
    cudaMemcpy(cudaInputPtr, dataPtr, N2*sizeof(uint8_t), cudaMemcpyHostToDevice);
    vecPack<<<input_grid_dim, input_block_dim>>>(cudaRealPtr, size, cudaInputPtr, N);

    cufftExecR2C(ps->cufft_r2c_plan, cudaRealPtr, cudaPtr);

// Loading various stuff
    cudaMemcpy(ps->cuda_lsum_buffer + icp * alloc_size, mxGetData(lsum), size2*sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(ps->cuda_denom_buffer + icp * alloc_size, mxGetData(denom), size2*sizeof(float), cudaMemcpyHostToDevice);

    N = mxGetM(nonzero);

    ps->cuda_nonzero_items[icp] = N;

    if (N%BLOCK_SIZE) ps->grid_size[icp] = 1 + (N / BLOCK_SIZE);
    else ps->grid_size[icp] = N / BLOCK_SIZE;
    
    cudaMemcpy(ps->cuda_nonzero_buffer + icp * alloc_size, mxGetData(nonzero), ps->cuda_nonzero_items[icp]*sizeof(uint16_t), cudaMemcpyHostToDevice);

    return cudaPtr;
}

static inline mxArray *fftCompute(TProcessingState *ps, int icp, const mxArray *data, float sum, float denom) {
    uint8_t *dataPtr;
    double *ar;
    mxArray *res;

    int size = ps->fft_size;
    int size2 = size * size;
    int alloc_size = ps->fft_alloc_size;

    int N = mxGetM(data);
    int N2 = N * N;

    dim3 input_block_dim(N, 1, 1);
    dim3 input_grid_dim(N, 1, 1);

    dim3 block_dim(size / 2 + 1, 1, 1);
    dim3 grid_dim(size, 1, 1);

    uint8_t *cudaInputPtr = ps->cuda_input_buffer + icp * alloc_size;
    cufftComplex *cudaPtr = ps->cuda_base_buffer + icp * alloc_size;
    cufftReal *cudaRealPtr = ps->cuda_temp_buffer;
    cufftComplex *cudaDataPtr = ps->cuda_data_buffer;
    float *cudaResultPtr = ps->cuda_final_buffer;

    dataPtr = (uint8_t*)mxGetData(data);
    cudaMemcpy(cudaInputPtr, dataPtr, N2*sizeof(uint8_t), cudaMemcpyHostToDevice);
    vecPack<<<input_grid_dim, input_block_dim>>>(cudaRealPtr, size, cudaInputPtr, N);

    cufftExecR2C(ps->cufft_r2c_plan, cudaRealPtr, cudaDataPtr);

    vecMul<<<grid_dim,block_dim>>>(cudaDataPtr, cudaPtr, size/2+1);

    cudaRealPtr = ps->cuda_result_buffer;
    cufftExecC2R(ps->cufft_c2r_plan, cudaDataPtr, cudaRealPtr);

    uint16_t nz_items = ps->cuda_nonzero_items[icp];
    uint16_t *nz = ps->cuda_nonzero_buffer + icp*alloc_size;
    
    float *cudaDenom = ps->cuda_denom_buffer + icp*alloc_size;
    float *cudaLSum = ps->cuda_lsum_buffer + icp*alloc_size;

    int grids;

//    printf("%i (%i %i)\n", nz_items, icp, ps->ncp);

    if (nz_items) {
	grids = ps->grid_size[icp];
    
        dim3 output_block_dim(BLOCK_SIZE, 1, 1);
        dim3 output_grid_dim(grids, 1, 1);

	vecCompute<<<output_grid_dim, output_block_dim>>>(
	    nz, cudaResultPtr,
	    cudaRealPtr, 1./(size2 * (N2 - 1)),
	    cudaLSum, sum / (N2 * (N2 - 1)),
	    cudaDenom, denom
	);
    }
    
    res = mxCreateNumericMatrix(size, size, mxSINGLE_CLASS, mxREAL);
    ar = mxGetPr(res);

    cudaMemcpy(ar, cudaResultPtr, size2*sizeof(cufftReal), cudaMemcpyDeviceToHost);

    return res;
}


/*
static inline double fftDownloadData(TProcessingState *ps, mxArray *data) {
  cudaMemcpy( input_single, rhs_complex_d, sizeof(cufftComplex)*N*M, cudaMemcpyDeviceToHost);
}    
*/


static void selfClean() {
    reportMessage("Unloading normxcorr_hw");
    
    if (pstate) {
	pstateFree(pstate);
	pstate = NULL;
    }
}


void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
    int err;
    int64_t *errPtr;
    
    int deviceCount;
    cudaDeviceProp deviceProp;
    
    mxArray *idMatrix;
    int32_t id, *idPtr;

    TProcessingState *ps;

    TAction action;

    int iprop;
    
    const mxArray *input;
    const mxArray *base;
    const mxArray *lsum;
    const mxArray *denom;
    const mxArray *nonzero;

    double input_sum;
    double input_denom;

    unsigned int icp;

    if (!nrhs) {
	reportMessage("Initializing normxcorr_hw instance");

	if (nlhs != 1) {
	    reportError("You should accept a single result from initialization call");
	    return;
	}
	
	if (pstate) {
	    reportError("Only a single calculation process is supported at the moment");
	    return;
	}


	// Initialising, for now a single client is supported only
	idMatrix = mxCreateNumericMatrix(1, 1, mxINT32_CLASS, mxREAL);
	if (!idMatrix) {
	    reportError("Initialization is failed");
	    return;
	}

	// Detecting cuda devices
	cudaGetDeviceCount(&deviceCount);
	if (deviceCount) {
	    cudaGetDeviceProperties(&deviceProp, 0);
	    if ((deviceProp.major > 1)||((deviceProp.major == 1)&&(deviceProp.minor > 2))) {
		id = 1;
	    } else { // Hardware capabilities are bellow 1.3
		id = 0;
	    }
	} else { // No cuda device, using software
	    id = -1;
	}
	
	
	if (id > 0) {
	    pstate = pstateInit();
	    if (!pstate) {
		mxDestroyArray(idMatrix);
	        reportError("State structure initialization is failed");
	        return;
	    }
	} else {
	    pstate = NULL;
	}

	idPtr = (int32_t*)mxGetData(idMatrix);
	idPtr[0] = id;
	
	plhs[0] = idMatrix;
	
	mexAtExit(selfClean);

	return;
    } else {
/*
	idMatrix = (mxArray*)prhs[0];
	if ((mxGetClassID(idMatrix) != mxINT32_CLASS)||(mxGetM(idMatrix) != 1)||(mxGetN(idMatrix) != 1)) {
	    reportError("Invalid parameter is supplied in place of process identificator");
	    return;
	}

	idPtr = (int32_t*)mxGetData(idMatrix);
	if (!idPtr) {
	    reportError("Mex is not able to obtain process identificator");
	    return;
	}
	
	id = *idPtr;
	if (id != 1) {
	    reportError("Invalid process identificator is supplied");
	    return;
	}

        if (!pstate) {
	    reportError("The interface is not initialized");
	    return;
	}
*/
    }

	// Clean request
    if (nrhs == 1) {
	reportMessage("Cleaning normxcorr_hw instance");

	pstateFree(pstate);
	pstate = NULL;
	return;
    }

    ps = pstate;
    
    action = (TAction)int(mxGetScalar((mxArray*)prhs[1]));

//    reportMessage("Executing normxcorr_hw action: %u", action);

    switch (action) {
     case ACTION_COMPUTE_FRAGMENT:
/*
	if (nrhs != 6) {
	    reportError("This action expects 4 arguments, but %i is passed", nrhs - 2);
	    return;
        }
	
	if (nlhs != 1) {
	    reportError("This action expects single ouput, but %i is passed", nlhs);
	}
*/
	
	icp = (unsigned int)mxGetScalar(prhs[2]);
/*	if (icp >= ps->ncp) {
	    reportError("The control point (%i) is out of range (0-%u)", icp, ps->ncp - 1);
	    return;
	}
*/
	input = prhs[3];
/*    
	if (mxGetNumberOfDimensions(input) != 2) {
	    reportError("Invalid dimensionality of input matrix, 2D matrix is expected");
	    return;
	}
	
	if (mxGetClassID(input) != mxUINT8_CLASS) {
	    reportError("Invalid matrix. The data type (%s) is not supported", mxGetClassName(input));
	    return;
	}

*/
	input_sum = mxGetScalar(prhs[4]);
	input_denom = mxGetScalar(prhs[5]);

	plhs[0] = fftCompute(ps, icp, input, input_sum, input_denom);
     break;
     case ACTION_COMPUTE_BASE:
	if (nrhs != 7) {
	    reportError("This action expects 5 arguments, but %i is passed", nrhs - 2);
	    return;
        }

	icp = (unsigned int)mxGetScalar(prhs[2]);
	if (icp >= ps->ncp) {
	    reportError("The control point (%i) is out of range (0-%u)", icp, ps->ncp - 1);
	    return;
	}

	base = prhs[3];
    
	if (mxGetNumberOfDimensions(base) != 2) {
	    reportError("Invalid dimensionality of base matrix, 2D matrix is expected");
	    return;
	}

	if (mxGetClassID(base) != mxUINT8_CLASS) {
	    reportError("Invalid matrix. The data type (%s) is not supported", mxGetClassName(base));
	    return;
	}

	lsum = prhs[4];
	denom = prhs[5];
	nonzero = prhs[6];

	iprop = ps->fft_size;
	if (
	    (mxGetNumberOfDimensions(lsum) != 2)||
	    (mxGetNumberOfDimensions(denom) != 2)||
	    (mxGetClassID(lsum) != mxSINGLE_CLASS)||
	    (mxGetClassID(denom) != mxSINGLE_CLASS)||
	    (mxGetClassID(nonzero) != mxUINT16_CLASS)||
	    (mxGetN(lsum) != iprop)||(mxGetM(lsum) != iprop)||
	    (mxGetN(denom) != iprop)||(mxGetM(denom) != iprop)
	    
	) {
	    reportError("Invalid properties for base initialization are specified");
	    return;
	}

	fftUploadBaseData(ps, icp, base, lsum, denom, nonzero);
     break;
     case ACTION_SETUP:
	if (nrhs != 4) {
	    reportError("SETUP action expects 'ncp' and 'corrsize' parameters");
	    return;
	}

	iprop = (int)mxGetScalar(prhs[2]);
	ps->ncp = iprop + 1;

	iprop = (int)mxGetScalar(prhs[3]);
	ps->fft_size = 6 * iprop + 1;
	ps->fft_size2 = ps->fft_size * ps->fft_size;
	ps->fft_inner_size = ps->fft_size * (ps->fft_size / 2 + 1);

	if (ps->fft_size2 % BLOCK_SIZE) {
	    ps->fft_alloc_size = ((ps->fft_size2 / BLOCK_SIZE) + 1) * BLOCK_SIZE;
	} else {
	    ps->fft_alloc_size = ps->fft_size2;
	}

	err = fftInit(ps);

	if (nlhs == 1) {
	    idMatrix = mxCreateNumericMatrix(1, 1, mxINT64_CLASS, mxREAL);
	    if (idMatrix) {
		errPtr = (int64_t*)mxGetData(idMatrix);
		errPtr[0] = err;
		plhs[0] = idMatrix;
	    } else {
		reportError("Initialization of result matrix is failed");
	        return;
	    }
	}
     break;
     case ACTION_PREPARE:
        fftPrepare(ps);
     break;
     default:
        reportError("Unknown request %i", action);
    }
    
}