balancing.cu

#include "../balancing.cuh"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cuda.h>
#include <cuda_runtime.h>

int main(int argc, char* args[]) {
	int n_worker = atoi(args[1]);
	int n_expert = atoi(args[2]);
	int cap_v = atoi(args[3]);
	int tot_expert = n_worker * n_expert;

	long* lec = new long[tot_expert];
	for (int i = 0; i < tot_expert; ++i) {
		lec[i] = i;
	}
	long* g_lec;
	cudaMalloc(&g_lec, sizeof(long) * tot_expert);
	cudaMemcpy(g_lec, lec, sizeof(long) * tot_expert, cudaMemcpyHostToDevice);

	int* cap = new int[n_expert];
	for (int i = 0; i < n_expert; ++i) {
		cap[i] = cap_v;
	}
	int* g_cap;
	cudaMalloc(&g_cap, sizeof(int) * n_expert);
	cudaMemcpy(g_cap, cap, sizeof(int) * n_expert, cudaMemcpyHostToDevice);

	long* eca = new long[tot_expert];
	long* g_eca;
	cudaMalloc(&g_eca, sizeof(long) * tot_expert);

	auto smgr = getCudaStreamManager(0);
	fmoe_cuda_limit_by_capacity_impl(g_lec, g_cap, g_eca, n_expert, n_worker, smgr);

	cudaMemcpy(cap, g_cap, sizeof(int) * n_expert, cudaMemcpyDeviceToHost);
	cudaMemcpy(eca, g_eca, sizeof(long) * tot_expert, cudaMemcpyDeviceToHost);

	printf("%d\n", cap[0]);
	for (int i = 0; i < tot_expert; ++i) {
		printf("%ld %ld\n", lec[i], eca[i]);
	}
}