/normxcorr/trunk

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

#include <mex.h>

#include "normxcorr_hw.h"
#include "local_sum.h"

#include "normxcorr_hw_msg.h"
#include "normxcorr_hw_kernel.cu"


static TProcessingState *pstate = NULL;

static void fftFree(TProcessingState *ps) {
    if (ps->cuda_base_buffer) {

	cudaFree(ps->cuda_lsum_temp);
	
	cudaFree(ps->cuda_lsum_buffer);
	cudaFree(ps->cuda_denom_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;
    }

	// DS: Source of bug, that occasionaly can corrupt something ...
    if (ps->cudpp_initialized) {
	cudppDestroyPlan(ps->cudpp_plan);
	ps->cudpp_initialized = false;
    }

    if (ps->fft_initialized) {
	cufftDestroy(ps->cufft_r2c_plan);
	cufftDestroy(ps->cufft_c2r_plan);
//	cufftDestroy(ps->cufft_plan);
//	cublasShutdown();
	ps->fft_initialized = false;
    }

}

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

    int lsum_alloc_size2 = ps->lsum_alloc_size * ps->lsum_alloc_size;
    
    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;

    cudpp_config.algorithm = CUDPP_SCAN;
    cudpp_config.options = CUDPP_OPTION_FORWARD |  CUDPP_OPTION_INCLUSIVE;
    cudpp_config.op = CUDPP_ADD;
    cudpp_config.datatype = CUDPP_FLOAT;

    cudpp_err = cudppPlan(&ps->cudpp_plan, cudpp_config, ps->lsum_alloc_size, ps->lsum_alloc_size, ps->lsum_alloc_size);
    if (cudpp_err != CUDPP_SUCCESS) {
	reportError("Problem initializing CUDPP plan, cudpp code: %i", cudpp_err);
	fftFree(ps);
	return ERROR_CUDPP;
    }
    
    ps->cudpp_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) + lsum_alloc_size2 * 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) + lsum_alloc_size2 * 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_lsum_temp, 4 * lsum_alloc_size2  * sizeof(float));
    if (cuda_err) {
	reportError("Device memory allocation of 4*%u*float bytes for lsum temporary buffer is failed", lsum_alloc_size2);
	fftFree(ps);
	return ERROR_MALLOC;
    }
	// We need to zero temporary buffers as well, since we are not computing
	// cumsum of complete matrix, but non-zero part of it
    cudaMemset((void*)ps->cuda_lsum_temp, 0, 4 * lsum_alloc_size2 * sizeof(float));

    
    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;
	ps->cudpp_initialized = false;
    }

    return ps;
}


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

    if (!ps->fft_initialized) {
	reportError("cuFFT engine is not initialized yet");
	return NULL;
    }
    
    int size = ps->fft_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);

    float *lsum_temp = ps->cuda_lsum_temp;
    int step = ps->lsum_alloc_size * ps->lsum_alloc_size;

    cudaMemset((void*)(ps->cuda_lsum_temp + 2*step), 0, size * ps->lsum_alloc_size * sizeof(float));
    cudaMemset((void*)(ps->cuda_lsum_temp + 3*step), 0, size * ps->lsum_alloc_size * sizeof(float));

    vecPackBase<<<input_grid_dim, input_block_dim>>>(
	cudaInputPtr, N, 
	cudaRealPtr, size, 
	lsum_temp, lsum_temp + step, ps->lsum_alloc_size, ps->lsum_size
    );

	// In general we should expect non-zero denominals, therefore the Nonzero array is not computed
    local_sum(ps, 
	ps->cuda_lsum_buffer + icp * alloc_size, ps->cuda_denom_buffer + icp * alloc_size,
	lsum_temp + (2 * step), lsum_temp + (3 * step),
	lsum_temp, lsum_temp + step);

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

    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);

    dim3 output_block_dim(size, 1, 1);
    dim3 output_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);

    float *cudaDenom = ps->cuda_denom_buffer + icp*alloc_size;
    float *cudaLSum = ps->cuda_lsum_buffer + icp*alloc_size;

    vecCompute<<<output_grid_dim, output_block_dim>>>(
	    cudaResultPtr,
	    cudaRealPtr, 1./(size2 * (N2 - 1)),
	    cudaLSum, sum / (N2 * (N2 - 1)),
	    cudaDenom, denom,
	    size
    );
    
    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() {
    if (pstate) {
	reportMessage("Self-cleaning normxcorr_hw instance");

	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;

#ifdef VALIDATE_LSUM
    const mxArray *lsum;
    const mxArray *denom;
    const mxArray *nonzero;
#endif /* VALIDATE_LSUM */

    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 != 4)
#ifdef VALIDATE_LSUM
	    &&(nrhs != 7)
#endif /* VALIDATE_LSUM */
	) {
	    reportError("This action expects 2 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;
	}

	iprop = ps->fft_size;

#ifdef VALIDATE_LSUM
	if (nrhs == 7) {
	    lsum = prhs[4];
	    denom = prhs[5];
	    nonzero = prhs[6];
	    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;
	    }
	} else {
	    lsum = NULL;
	    denom = NULL;
	    nonzero = NULL;
	}
#endif /* VALIDATE_LSUM */
	
	fftUploadBaseData(ps, icp, base);

#ifdef VALIDATE_LSUM
	local_sum_validate(ps, icp, lsum, denom);
#endif /* VALIDATE_LSUM */

     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->corr_size = iprop;
	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_1D) {
	    ps->fft_alloc_size = ((ps->fft_size2 / BLOCK_SIZE_1D) + 1) * BLOCK_SIZE_1D;
	} else {
	    ps->fft_alloc_size = ps->fft_size2;
	}
	
	ps->subimage_size = ps->corr_size * 4 + 1;
	ps->lsum_size = ps->corr_size * 2 + 1;

	ps->lsum_temp_size = ps->subimage_size + 2*ps->lsum_size - 1;

	if (ps->fft_size % BLOCK_SIZE_2D) {
	    ps->lsum_short_aligned_size = ((ps->fft_size / BLOCK_SIZE_2D) + 1) * BLOCK_SIZE_2D;
	} else {
	    ps->lsum_short_aligned_size = ps->fft_size;
	}

	if ((ps->lsum_temp_size) % BLOCK_SIZE_2D) {
	    ps->lsum_aligned_size = (((ps->lsum_temp_size) / BLOCK_SIZE_2D) + 1) * BLOCK_SIZE_2D;
	} else {
	    ps->lsum_aligned_size = ps->lsum_temp_size;
	}

	if ((ps->lsum_temp_size + ps->lsum_size) % BLOCK_SIZE_2D) {
	    ps->lsum_alloc_size = (((ps->lsum_temp_size + ps->lsum_size) / BLOCK_SIZE_2D) + 1) * BLOCK_SIZE_2D;
	} else {
	    ps->lsum_alloc_size = ps->lsum_temp_size + ps->lsum_size;
	}

	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);
    }
    
}