Update the rank and world size because we add DP and run across nodes and therefore more types of ranks and corresponding groups. 2. delete the timer output to screen.

Author: Hexu Zhao

cuda_rasterizer/forward.cu

@@ -178,9 +178,7 @@ __global__ void preprocessCUDA(int P, int D, int M,
 	float4* conic_opacity,
 	const dim3 grid,
 	uint32_t* tiles_touched,
-	bool prefiltered,
-	int local_rank,
-	int world_size)
+	bool prefiltered)
 {
 	auto idx = cg::this_grid().thread_rank();
 	if (idx >= P)
@@ -468,9 +466,7 @@ void FORWARD::preprocess(int P, int D, int M,
 	float4* conic_opacity,
 	const dim3 grid,
 	uint32_t* tiles_touched,
-	bool prefiltered,
-	int local_rank,
-	int world_size)
+	bool prefiltered)
 {
 	preprocessCUDA<NUM_CHANNELS> << <(P + ONE_DIM_BLOCK_SIZE - 1) / ONE_DIM_BLOCK_SIZE, ONE_DIM_BLOCK_SIZE >> > (
 		P, D, M,
@@ -497,8 +493,6 @@ void FORWARD::preprocess(int P, int D, int M,
 		conic_opacity,
 		grid,
 		tiles_touched,
-		prefiltered,
-		local_rank,
-		world_size
+		prefiltered
 		);
 }

cuda_rasterizer/forward.h

@@ -45,9 +45,7 @@ namespace FORWARD
 		float4* conic_opacity,
 		const dim3 grid,
 		uint32_t* tiles_touched,
-		bool prefiltered,
-		int local_rank,
-		int world_size);
+		bool prefiltered);
 
 	// Main rasterization method.
 	void render(

cuda_rasterizer/rasterizer_impl.cu

@@ -80,9 +80,7 @@ __global__ void duplicateWithKeys(
 	uint32_t* gaussian_values_unsorted,
 	int* radii,
 	bool* compute_locally,
-	dim3 grid,
-	int local_rank,
-	int world_size)
+	dim3 grid)
 {
 	auto idx = cg::this_grid().thread_rank();
 	if (idx >= P)
@@ -289,9 +287,9 @@ __global__ void updateTileTouched(
 	tiles_touched[idx] = cnt;
 }
 
-void save_log_in_file(int iteration, int local_rank, int world_size, std::string log_folder, const char* filename_prefix, const char* log_content) {
-	char* filename = new char[128];
-	sprintf(filename, "%s/%s_ws=%d_rk=%d.log", log_folder.c_str(), filename_prefix, world_size, local_rank);
+void save_log_in_file(int iteration, int global_rank, int world_size, std::string log_folder, const char* filename_prefix, const char* log_content) {
+	char* filename = new char[256];
+	sprintf(filename, "%s/%s_ws=%d_rk=%d.log", log_folder.c_str(), filename_prefix, world_size, global_rank);
 	std::ofstream outfile;
 	outfile.open(filename, std::ios_base::app);
 	outfile << "iteration: " << iteration << ", " << log_content << "\n";
@@ -307,7 +305,7 @@ void save_log_in_file(int iteration, int local_rank, int world_size, std::string
 std::tuple<int, int, int, int, int, bool, bool, std::string, std::string, std::string>
  prepareArgs(const pybind11::dict &args) {
     std::string mode = args["mode"].cast<std::string>();
-    std::string local_rank_str = args["local_rank"].cast<std::string>();
+    std::string global_rank_str = args["global_rank"].cast<std::string>();
     std::string world_size_str = args["world_size"].cast<std::string>();
     std::string iteration_str = args["iteration"].cast<std::string>();
     std::string log_interval_str = args["log_interval"].cast<std::string>();
@@ -317,7 +315,7 @@ std::tuple<int, int, int, int, int, bool, bool, std::string, std::string, std::s
     // std::string dist_division_mode_str = args["dist_division_mode"].cast<std::string>();
 	std::string dist_division_mode_str = "";
 
-    int local_rank = std::stoi(local_rank_str);
+    int global_rank = std::stoi(global_rank_str);
     int world_size = std::stoi(world_size_str);
     int iteration = std::stoi(iteration_str);
     int log_interval = std::stoi(log_interval_str);
@@ -328,7 +326,7 @@ std::tuple<int, int, int, int, int, bool, bool, std::string, std::string, std::s
 	cudaError_t status = cudaGetDevice(&device);
 
 	// Pack and return the variables in a tuple
-    return std::make_tuple(local_rank, world_size, iteration, log_interval, device,
+    return std::make_tuple(global_rank, world_size, iteration, log_interval, device,
 			zhx_debug, zhx_time,
 			mode, dist_division_mode_str, log_folder_str);
 }
@@ -359,14 +357,14 @@ int CudaRasterizer::Rasterizer::preprocessForward(
 	bool debug,//raster_settings
 	const pybind11::dict &args)
 {
-	auto [local_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);
+	auto [global_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);
 	char* log_tmp = new char[500];
 
 	// print out the environment variables
 	if (mode == "train" && zhx_debug && iteration % log_interval == 1) {
-		sprintf(log_tmp, "world_size: %d, local_rank: %d, iteration: %d, log_folder: %s, zhx_debug: %d, zhx_time: %d, device: %d, log_interval: %d, dist_division_mode: %s", 
-				world_size, local_rank, iteration, log_folder.c_str(), zhx_debug, zhx_time, device, log_interval, dist_division_mode.c_str());
-		save_log_in_file(iteration, local_rank, world_size, log_folder, "cuda", log_tmp);
+		sprintf(log_tmp, "world_size: %d, global_rank: %d, iteration: %d, log_folder: %s, zhx_debug: %d, zhx_time: %d, device: %d, log_interval: %d, dist_division_mode: %s", 
+				world_size, global_rank, iteration, log_folder.c_str(), zhx_debug, zhx_time, device, log_interval, dist_division_mode.c_str());
+		save_log_in_file(iteration, global_rank, world_size, log_folder, "cuda", log_tmp);
 	}
 
 	MyTimerOnGPU timer;
@@ -410,17 +408,15 @@ int CudaRasterizer::Rasterizer::preprocessForward(
 		conic_opacity,
 		tile_grid,
 		tiles_touched_temp_buffer,
-		prefiltered,
-		local_rank,
-		world_size
+		prefiltered
 	), debug)
 	timer.stop("10 preprocess");
 
 	int num_rendered = 0;//TODO: should I calculate this here?
 
 	// Print out timing information
 	if (zhx_time && iteration % log_interval == 1) {
-		timer.printAllTimes(iteration, world_size, local_rank, log_folder, true);
+		timer.printAllTimes(iteration, world_size, global_rank, log_folder, true);
 	}
 	delete log_tmp;
 	// free temporary buffer for tiles_touched. TODO: remove it. 
@@ -454,7 +450,7 @@ void CudaRasterizer::Rasterizer::preprocessBackward(
 	bool debug,
 	const pybind11::dict &args)
 {
-	auto [local_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);
+	auto [global_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);
 
 	MyTimerOnGPU timer;
 	const float focal_y = height / (2.0f * tan_fovy);
@@ -488,7 +484,7 @@ void CudaRasterizer::Rasterizer::preprocessBackward(
 
 	// Print out timing information
 	if (zhx_time && iteration % log_interval == 1) {
-		timer.printAllTimes(iteration, world_size, local_rank, log_folder, false);
+		timer.printAllTimes(iteration, world_size, global_rank, log_folder, false);
 	}
 }
 
@@ -537,7 +533,7 @@ int CudaRasterizer::Rasterizer::renderForward(
 	bool debug,
 	const pybind11::dict &args)
 {
-	auto [local_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);	
+	auto [global_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);	
 	char* log_tmp = new char[500];
 
 	MyTimerOnGPU timer;
@@ -593,9 +589,7 @@ int CudaRasterizer::Rasterizer::renderForward(
 		binningState.point_list_unsorted,
 		radii,
 		compute_locally,
-		tile_grid,
-		local_rank,
-		world_size)
+		tile_grid)
 	CHECK_CUDA(, debug)
 	timer.stop("40 duplicateWithKeys");
 
@@ -675,7 +669,7 @@ int CudaRasterizer::Rasterizer::renderForward(
 	//////////////////////////// Logging && Save Statictis ////////////////////////////////////////////
 	// DEBUG: print out timing information
 	if (mode == "train" && zhx_time && iteration % log_interval == 1) {
-		timer.printAllTimes(iteration, world_size, local_rank, log_folder, false);
+		timer.printAllTimes(iteration, world_size, global_rank, log_folder, false);
 	}
 
 	// DEBUG: print out the number of Gaussians contributing to each pixel.
@@ -731,7 +725,7 @@ int CudaRasterizer::Rasterizer::renderForward(
 				ave_n_contrib2loss,
 				contrib2loss_ratio);
 
-			save_log_in_file(iteration, local_rank, world_size, log_folder, "n_contrib", log_tmp);
+			save_log_in_file(iteration, global_rank, world_size, log_folder, "n_contrib", log_tmp);
 			global_sum_n_rendered += n_rendered;
 			global_sum_n_considered += sum_n_considered;
 			global_sum_n_contrib2loss += sum_n_contrib2loss;
@@ -742,8 +736,8 @@ int CudaRasterizer::Rasterizer::renderForward(
 		float global_ave_n_considered_per_pix = global_sum_n_considered / (float)total_pixels;
 		float global_ave_n_contrib2loss_per_pix = global_sum_n_contrib2loss / (float)total_pixels;
 
-		sprintf(log_tmp, "local_rank: %d, world_size: %d, num_tiles: %d, num_pixels: %d, num_rendered: %d, global_ave_n_rendered_per_pix: %f, global_ave_n_considered_per_pix: %f, global_ave_n_contrib2loss_per_pix: %f", 
-			(int)local_rank,
+		sprintf(log_tmp, "global_rank: %d, world_size: %d, num_tiles: %d, num_pixels: %d, num_rendered: %d, global_ave_n_rendered_per_pix: %f, global_ave_n_considered_per_pix: %f, global_ave_n_contrib2loss_per_pix: %f", 
+			(int)global_rank,
 			(int)world_size,
 			(int)num_local_tiles,
 			(int)total_pixels,
@@ -752,7 +746,7 @@ int CudaRasterizer::Rasterizer::renderForward(
 			global_ave_n_considered_per_pix, 
 			global_ave_n_contrib2loss_per_pix
 		);
-		save_log_in_file(iteration, local_rank, world_size, log_folder, "n_contrib", log_tmp);
+		save_log_in_file(iteration, global_rank, world_size, log_folder, "n_contrib", log_tmp);
 
 		delete[] cpu_ranges;
 		delete[] cpu_n_considered;
@@ -784,7 +778,7 @@ void CudaRasterizer::Rasterizer::renderBackward(
 	bool debug,
 	const pybind11::dict &args)
 {
-	auto [local_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);
+	auto [global_rank, world_size, iteration, log_interval, device, zhx_debug, zhx_time, mode, dist_division_mode, log_folder] = prepareArgs(args);
 
 	MyTimerOnGPU timer;
 
@@ -825,6 +819,6 @@ void CudaRasterizer::Rasterizer::renderBackward(
 
 	// Print out timing information
 	if (zhx_time && iteration % log_interval == 1) {
-		timer.printAllTimes(iteration, world_size, local_rank, log_folder, false);
+		timer.printAllTimes(iteration, world_size, global_rank, log_folder, false);
 	}
 }

cuda_rasterizer/timers.cu

@@ -51,12 +51,12 @@ public:
                 return a.first < b.first;
             }
         );
-        if (print_prefix) {
-            std::cout << prefix << std::endl;
-        }
-        for (const auto& pair : sortedTimes) {
-            std::cout << pair.first << " time: " << elapsedMilliseconds(pair.first, "sum") << " ms" << std::endl;
-        }
+        // if (print_prefix) {
+        //     std::cout << prefix << std::endl;
+        // }
+        // for (const auto& pair : sortedTimes) {
+        //     std::cout << pair.first << " time: " << elapsedMilliseconds(pair.first, "sum") << " ms" << std::endl;
+        // }
         //save in file
         FILE *fp;
         fp = fopen(filename, "a");

diff_gaussian_rasterization/__init__.py

@@ -174,13 +174,13 @@ def render_gaussians(
     )
 
 def get_extended_compute_locally(cuda_args, image_height, image_width):
-    local_rank = int(cuda_args["local_rank"])
+    mp_rank = int(cuda_args["mp_rank"])
     dist_global_strategy = [int(x) for x in cuda_args["dist_global_strategy"].split(",")]
 
     num_tile_y = (image_height + 16 - 1) // 16 #TODO: this is dangerous because 16 may change.
     num_tile_x = (image_width + 16 - 1) // 16
-    tile_l = max(dist_global_strategy[local_rank]-num_tile_x-1, 0)
-    tile_r = min(dist_global_strategy[local_rank+1]+num_tile_x+1, num_tile_y*num_tile_x)
+    tile_l = max(dist_global_strategy[mp_rank]-num_tile_x-1, 0)
+    tile_r = min(dist_global_strategy[mp_rank+1]+num_tile_x+1, num_tile_y*num_tile_x)
 
     extended_compute_locally = torch.zeros(num_tile_y*num_tile_x, dtype=torch.bool, device="cuda")
     extended_compute_locally[tile_l:tile_r] = True
@@ -368,20 +368,20 @@ def render_gaussians(self, means2D, conic_opacity, rgb, depths, radii, compute_l
     def get_local2j_ids(self, means2D, radii, cuda_args):
 
         raster_settings = self.raster_settings
-        world_size = int(cuda_args["world_size"])
-        local_rank = int(cuda_args["local_rank"])
+        mp_world_size = int(cuda_args["mp_world_size"])
+        mp_rank = int(cuda_args["mp_rank"])
 
         # TODO: make it more general.
         dist_global_strategy = [int(x) for x in cuda_args["dist_global_strategy"].split(",")]
-        assert len(dist_global_strategy) == world_size+1, "dist_global_strategy should have length WORLD_SIZE+1"
+        assert len(dist_global_strategy) == mp_world_size+1, "dist_global_strategy should have length WORLD_SIZE+1"
         assert dist_global_strategy[0] == 0, "dist_global_strategy[0] should be 0"
         dist_global_strategy = torch.tensor(dist_global_strategy, dtype=torch.int, device=means2D.device)
 
         args = (
             raster_settings.image_height,
             raster_settings.image_width,
-            local_rank,
-            world_size,
+            mp_rank,
+            mp_world_size,
             means2D,
             radii,
             dist_global_strategy,
@@ -391,7 +391,7 @@ def get_local2j_ids(self, means2D, radii, cuda_args):
         local2j_ids_bool = _C.get_local2j_ids_bool(*args) # local2j_ids_bool is (P, world_size) bool tensor
 
         local2j_ids = []
-        for rk in range(world_size):
+        for rk in range(mp_world_size):
             local2j_ids.append(local2j_ids_bool[:, rk].nonzero())
 
         return local2j_ids, local2j_ids_bool

rasterize_points.cu

@@ -453,8 +453,8 @@ __global__ void getTouchedIdsBool(
 torch::Tensor GetLocal2jIdsBoolCUDA(
 	int image_height,
 	int image_width,
-	int local_rank,
-	int world_size,
+	int mp_rank,
+	int mp_world_size,
 	const torch::Tensor& means2D,
 	const torch::Tensor& radii,
 	const torch::Tensor& dist_global_strategy,
@@ -465,13 +465,13 @@ torch::Tensor GetLocal2jIdsBoolCUDA(
 	const int W = image_width;
 	bool avoid_pixel_all2all = args["avoid_pixel_all2all"].cast<bool>();
 
-	torch::Tensor local2jIdsBool = torch::full({P, world_size}, false, means2D.options().dtype(torch::kBool));
+	torch::Tensor local2jIdsBool = torch::full({P, mp_world_size}, false, means2D.options().dtype(torch::kBool));
 
 	getTouchedIdsBool << <(P + ONE_DIM_BLOCK_SIZE - 1) / ONE_DIM_BLOCK_SIZE, ONE_DIM_BLOCK_SIZE >> >(
 		P,
 		H,
 		W,
-		world_size,
+		mp_world_size,
 		reinterpret_cast<float2*>(means2D.contiguous().data<float>()),
 		radii.contiguous().data<int>(),
 		dist_global_strategy.contiguous().data<int>(),

rasterize_points.h

@@ -170,8 +170,8 @@ RenderGaussiansBackwardCUDA(
 torch::Tensor GetLocal2jIdsBoolCUDA(
 	int image_height,
 	int image_width,
-	int local_rank,
-	int world_size,
+	int mp_rank,
+	int mp_world_size,
 	const torch::Tensor& means2D,
 	const torch::Tensor& radii,
 	const torch::Tensor& dist_global_strategy,