 
 namespace ginf {
 	template <typename T>
-	__device__ void gpuBpCollect(GpuGrid<T> *grid, int w, int h, T *msgs, T *from, int x, int y) {
+	__global__ void gpuBpCollectKernel(GpuGrid<T> *grid, int w, int h, T *msgs, T *from, T *data) {
+		GINF_DECL_X; GINF_DECL_Y;
+	
 		MatDim dim(w, h);
 		MatDim dimMsgs(w, h, GINF_NUM_DIR);
 
-		// For each state
-		for (int i = 0; i < grid->getNumLabels(); i++) {
-			from[dim.idx(x, y, i)] = 0;
-		
-			// For each neighbour in each direction, sum up the messages
-			for (int d = 0; d < GINF_NUM_DIR; d++) {
-				from[dim.idx(x, y, i)] += msgs[dimMsgs.idx(x + dDirX[d], y + dDirY[d], GINF_OPP_DIR(d), i)];
-			}
-		}
-	}
-	
-	template <typename T>
-	__device__ void gpuBpSendMsg(int iter, GpuGrid<T> *grid, int w, int h, T *msgs, T *from, T *data, T *out, int d) {
+		if (dim.isValidStrict(x, y)) {
+			// For each state
+			for (int i = 0; i < grid->getNumLabels(); ++i) {
+				from[dim(x, y, i)] = data[dim(x, y, i)];
+				
+				// For each neighbour in each direction, sum up the messages
+				for (int d = 0; d < GINF_NUM_DIR; d++) {
+					from[dim(x, y, i)] += msgs[dimMsgs(x, y, d, i)];
+				}
+			}
+		}
+	}
+	
+	template <typename T>
+	__global__ void gpuBpKernel(int iter, GpuGrid<T> *grid, int w, int h, T *msgs, T *from, T *data, int d) {
 		GINF_DECL_ALT_X(iter); GINF_DECL_Y;
-		MatDim dim(w, h);
+		
+		MatDim dim(w, h, GINF_NUM_DIR);
 		MatDim dimMsgs(w, h, GINF_NUM_DIR);
 
-	   	int nx = x + dDirX[d], ny = y + dDirY[d];
-
-		T minh = out[0] = data[dim.idx(x, y, 0)] + (from[dim.idx(x, y, 0)] - msgs[dimMsgs.idx(nx, ny, GINF_OPP_DIR(d), 0)]);
-	    	for (int i = 1; i < grid->getNumLabels(); i++) {
-	    		out[i] = data[dim.idx(x, y, i)] + (from[dim.idx(x, y, i)] - msgs[dimMsgs.idx(nx, ny, GINF_OPP_DIR(d), i)]);
-	    		minh = GINF_MIN(minh, out[i]);
-	    	}
-	    	
-	    	// Find the linear scaling constant
-		T scale = grid->getSmoothnessCost(0, 1);
-		for (int i = 1; i < grid->getNumLabels(); i++)
-			out[i] = GINF_MIN(out[i], out[i - 1] + scale);
-		for (int i = grid->getNumLabels() - 2; i >= 0; i--)
-			out[i] = GINF_MIN(out[i], out[i + 1] + scale);
-	    	
-		// Truncate the messages using truncate constant
-		minh += GINF_MIN(grid->getSmoothnessCost(0, 1) * grid->getNumLabels(),
-				grid->getSmoothnessCost(0, grid->getNumLabels() - 1));
-		T val = 0;
-		for (int i = 0; i < grid->getNumLabels(); i++) {
-			out[i] = GINF_MIN(out[i], minh);
-			val += out[i];
-		}
-		
-		// Normalize
-		val /= grid->getNumLabels();
-		for (int i = 0; i < grid->getNumLabels(); i++) {
-			out[i] -= val;
-		}
-	}
-	
-	template <typename T>
-	__global__ void gpuBpKernel(int iter, GpuGrid<T> *grid, int w, int h, T *msgs, T *from, T *data) {
-		GINF_DECL_ALT_X(iter); GINF_DECL_Y;
-		MatDim dim(w, h, GINF_NUM_DIR);
-		
-		if (dim.isValidStrict(x, y)) {
-			// Collect messages from neighbours
-			gpuBpCollect(grid, w, h, msgs, from, x, y);
-			
-			// Foreach direction, send message
+		if (dim.isValidStrict(x, y)) {
+			// Create a vector to store the message. This is stored in local memory, and so writing to it
+			// is coalesced.
 			T out[GINF_MAX_NUM_LABELS];
-			for (int d = 0; d < GINF_NUM_DIR; d++) {
-				gpuBpSendMsg(iter, grid, w, h, msgs, from, data, out, d);
-				
-				for (int i = 0; i < grid->getNumLabels(); i++) {
-					msgs[dim.idx(x, y, d, i)] = out[i];
-				}
+			
+			// Find the linear scaling constant
+			T scale = grid->getSmoothnessCost(0, 1);
+			
+			// Calculate message vector in the forward direction
+			T minh = out[0] = from[dim(x, y, 0)] - msgs[dimMsgs(x, y, d, 0)];
+			for (int i = 1; i < grid->getNumLabels(); i++) {
+				out[i] = GINF_MIN(from[dim(x, y, i)] - msgs[dimMsgs(x, y, d, i)], out[i - 1] + scale);
+				minh = GINF_MIN(minh, out[i]);
+			}
+				
+			// Truncate the messages using truncate constant
+			minh += grid->getSmoothnessCost(0, grid->getNumLabels() - 1);
+				
+			// Compute the final out vector, and sum up all the elements to perform normalization
+			T val = 0;
+			for (int i = grid->getNumLabels() - 2; i >= 0; i--) {
+				out[i] = GINF_MIN(GINF_MIN(out[i], out[i + 1] + scale), minh);
+				val += out[i];
+			}
+
+			// Find the normalizing constant
+			val /= grid->getNumLabels();
+
+			// Update the message
+			int nx = x + dDirX[d], ny = y + dDirY[d];
+			for (int i = 0; i < grid->getNumLabels(); i++) {
+				// Update using the normalized message values
+				msgs[dim(nx, ny, GINF_OPP_DIR(d), i)] = out[i] - val;
 			}
 		}
 	}

			}
			
			// For each iteration, run belief propagation
			for (int i = 0; i < numItersPerLevel; i++) {
				gpuBpKernel<<<gridSize, blockSize>>>(i, dGrid, w, h, dMsgs[t], dFrom[t], dData[t]);
			}
		}
		
		// Finally get the beliefs
		gpuBpGetBeliefKernel<<<gridSize, blockSize>>>(dGrid, dMsgs[0], dFrom[0], dResult);
		
		// Transfer the results back to CPU

	template void gpuDecodeHbp<int>(Grid<int>*, int, int, Matrix<int>*);
	
	template <typename T>
	void gpuDecodeBp(Grid<T> *grid, int numIters, Matrix<int> *result) {
		gpuDecodeHbp(grid, 1, numIters, result);
	}
	
	template void gpuDecodeBp<float>(Grid<float>*, int, Matrix<int>*);
	template void gpuDecodeBp<int>(Grid<int>*, int, Matrix<int>*);
}
