Using scatter_reduce function did not work for CUDA/LLVM backend

[juhyeonkim95]

Hi,

I am trying to use scatter_reduce function to aovs, like this.

dr::scatter_reduce(ReduceOp::Add, aovs, throughput[0], index);

This worked on scalar mode, but failed on CUDA/LLVM mode, with the following error.

Not implemented, please use scatter_reduce() variant that takes a array target argument.

Could you please help me regarding this problem? I have tried several ways but could not resolve the problem.

Thanks in advance.

[merlinND]

Hello @juhyeonkim95,

What are the types of aovs and throughput[0]?
Assuming aovs is an array of pointers to DrJit arrays (e.g. Float aovs[6];), you should pass a specific AoV array as the target of scatter_reduce:

dr::scatter_reduce(ReduceOp::Add, aovs[0], throughput[0], index);

[juhyeonkim95]

They are from the default path integrator.

Float *aovs and Spectrum throughput

I want to scatter a certain channel of throughput to aovs with specific index, so I think it shouldn't be aovs[0].

So, assuming if we are using RGB mode, what I want to do is

aovs[3*index[0]+0] += throughput[0][0]
aovs[3*index[1]+0] += throughput[0][1]
aovs[3*index[2]+0] += throughput[0][2]

for R channel, and same for other channels like,

aovs[3*index[0]+1] += throughput[1][0]
aovs[3*index[1]+1] += throughput[1][1]
aovs[3*index[2]+1] += throughput[1][2]

aovs[3*index[0]+2] += throughput[2][0]
aovs[3*index[1]+2] += throughput[2][1]
aovs[3*index[2]+2] += throughput[2][2]

I implemented the above using the following code.

dr::scatter_reduce(ReduceOp::Add, aovs, throughput[0], 3*index+0);
dr::scatter_reduce(ReduceOp::Add, aovs, throughput[1], 3*index+1);
dr::scatter_reduce(ReduceOp::Add, aovs, throughput[2], 3*index+2);

This worked for scalar mode, but failed for CUDA/LLVM mode...

[rtabbara]

Hi @juhyeonkim95,

If I understand your requirement correctly, I don't think you need scatter_reduce at all here and can simply do

aovs[3*index] 	        += throughput[0];
aovs[3*index + 1]	+= throughput[1];
aovs[3*index + 2]	+= throughput[2];

where index is just a scalar value so that 3*index is the particular AOV channel you want to write into.

As Merlin suggested, because the Float type when targeting LLVM/CUDA variants is a JitArray, you need your target argument to be a Float not a pointer. Passing a Float* to scatter reduce means you're passing a pointer to an array which is why you're getting the error.

[juhyeonkim95]

Hi,

Thank you for the suggestion, but index is not a scalar value but UInt32, so aovs[3*index] += throughput[0]; causes a compile error.

I think I should use scatter_reduce because I want to create a histogram on aovs.

[rtabbara]

Ah, in this case I think you will have to manually split up your operations across each AOV channel. Again though, I don't think scatter_reduce is what you're after and you can instead use a mask

// Compute the array that maps throughput to a given AOV index
UInt32 index = ...;

// Here, I'm assuming you've allocated 3 * (max_index + 1) AOV channels
for (size_t i = 0; i <= max_index; ++i) {
	// Get corresponding AOV channels
	Float& r = aovs[3*i];
	Float& g = aovs[3*i + 1];
	Float& b = aovs[3*i + 2];

	Mask mask = index == i;
	dr::masked(r, mask) += throughput[0];
	dr::masked(g, mask) += throughput[1];
	dr::masked(b, mask) += throughput[2];
}

[juhyeonkim95]

Thank you for the suggestion.

This worked, but I need to disable loop record by following code to get a correct result.
dr.set_flag(dr.JitFlag.LoopRecord, False)

Also because it uses for loop, as max_index increases, it becomes really slow. max_index is usually set to 1000-10000.
Actually that is the reason why I did not use for loop, but built-in scatter_reduce function that can perform an efficient atomic operation.

Furthermore, for scalar_rgb mode, the code using scatter_reduce exactly gave the result that I wanted, so I think using scatter_reduce is not a wrong choice.

dr::scatter_reduce(ReduceOp::Add, aovs, throughput[0], index);
dr::scatter_reduce(ReduceOp::Add, aovs, throughput[1], index);
dr::scatter_reduce(ReduceOp::Add, aovs, throughput[2], index);

The only problem is that this only worked for scalar mode and not for CUDA/LLVM modes.

[rtabbara]

Also because it uses for loop, as max_index increases, it becomes really slow. max_index is usually set to 1000-10000.

If you want max_index to range between 1000-10,000 then that implies you're after between 3000-30,000 AOV channels and for a reasonably sized film resolution, I would be a bit concerned about hitting memory limits.

Nonethless, supposing you could allocate enough memory for all your AOV channels, I think the confusion here is the use of index to map to (an RGB set of) AOV channels, and not offsets of an array. In LLVM/CUDA mode, when calling

dr::scatter_reduce(ReduceOp::Add, aovs, throughput[0], index)

here, Float* aovs = [aov_ch_1, aov_ch_2, ... aov_ch_n] where each of aov_ch_i is an array. So you're trying to use index to specify which AOV channel you want to write into but you don't have any information as to where within that AOV channel a given throughput value should be written. e.g. throughput[0][10] maps to aov_index=5 but where in the channel aovs[3*5][?] should it be written?

The reason why this works in scalar mode is that now you're processing a ray at a time and your array size for each channel is 1. Hence, Float* becomes a normal C-type float* and there's no ambiguity as to where you're writing each throughput value in memory.

I think one alternative approach is to avoid using AOVs entirely, and just define a single Color<Float, 3> that packs your data as

Color<Float,3> flattened = [aov_ch_1_index0, aov_ch_1_index1,..., aov_ch1_indexN, aov_ch_2_index0, ...]

Then your index would be

size_t width = res_x * res_y;

// Map throughput to aov index
UInt32 aov_index = ...;
UInt32 offsets = dr::arange<UInt32>(width);

// Combination of which aov channel I want and also where in that channel
UInt32 index = width * aov_index + offsets;

and you could then call

dr::scatter_reduce(ReduceOp::Add, flattened, throughput, index)

[juhyeonkim95]

Thank you for the suggestion. I am travelling now, so I will test these code later and let you know again.