Questions about the implementation

[Glinttsd]

Hi, thanks for your great work!

I am trying to exploit how the TP kernel is implemented. I noticed that there is a CUDA file named subkernel_per_interaction_multirep.cuh (which is attached here) generated during the execution. I believe that is where the GPU kernel is implemented.

Regarding this CUDA file, I have two questions for its forward process:

1. The sparse computation part to compute the CG coefficients (Algorithm 2 in the paper) seems not unrolled. The source code is provided below. The for loop is not tagged with #pragma unroll to unroll the loop.

  // PERFORM CG DECOMPOSITION CALCULATION 
  {%- for i in range(tensor.nnz) %}
      {%- set coord1, coord2, coord3, value = tensor.tuples[i] %}
      L3_local_vec[{{coord3}}] += {{value}} * {{instructions[interaction_index].path_weight}} * L1_local_vec[{{coord1}}] * L2_local_vec[{{coord2}}];
  {%- endfor %}

2. The calculated results L3_local_vec from the above code seems been written into the shared memory smem_gemm_L3, and then multiplies the weight matrix. It seems to be contrary to the paper's description:

Finally, the output Z is accumulated to shared memory after multiplication by W .

The source code is provided below.

// WRITE TO SMEM_GEMM_L3 
#pragma unroll 
for(int L3_irrep_index = 0; L3_irrep_index < {{L3.irrep_lengths[w]}}; L3_irrep_index++){
    smem_gemm_L3[(threadIdx.x * {{L3.irrep_lengths[w]}}) + L3_irrep_index] = L3_local_vec[L3_irrep_index]; 
}

// WAIT FOR WEIGHTS TO HAVE ARRIVED
cooperative_groups::wait(group); 

group.sync();

// PERFORM MATMUL 
int i = threadIdx.x;
if(i < L3_mults){
    (some calculations)
}
group.sync();

Any help would be greatly appreciated! Thank you again for the great work!

[vbharadwaj-bk]

Hi @Glinttsd, this implementation is deprecated and no longer in our repo (it was abandoned early in development). Giving a short response here and Austin will follow up.

You want to take a look at loop_unroll_tp.cuh in the templates folder for the low-level subkernels, and loop_unroll_batch.cuh and loop_unroll_conv.cuh for the wrappers that orchestrate the subkernels according to the computation schedule. That said, the approach is still very similar to what you have pasted above.

1. That for loop is actually jinja directive (notice the {% %}; it is "unrolled" when the template runs by writing out the full stream of instructions. Try heading to implementations/LoopUnrollTP.py and uncommenting these two lines to see what the kernel looks like:

   OpenEquivariance/openequivariance/implementations/LoopUnrollTP.py

   Line 98 in 7db972c

   # with open("scratch.txt", "w") as f:

2. For UVU instructions, the multiplication happens in the registers. For UVW instructions, yes, we write out the intermediate to shared memory before performing the weight matrix multiplication. I don't think this is contrary to the paper, we just elided mentioning that for certain instructions, there is an intermediate write-out.

[asglover]

In the paper, we do mention that, for UVW, we write to shared memory first:

The matrix multiplication by the weights at the end
of Algorithm 2 depends on the structure of W. When
W is square and diagonal (kernel B), multiplication
proceeds asynchronously in parallel across all threads.
When W is a general dense matrix, we temporarily
store zreg to shared memory and perform a warp-level
matrix-multiplication across all threads.

[asglover]

Out of curiosity, what is your end goal for "exploiting the how TP kernel is implemented"? Is there a particular application you are trying to accelerate that we don't currently support? Or a similar problem?

[Glinttsd]

Hi, thank you for your quick response!

Me and my team are recently focusing on the FPGA acceleration of the tensor product computation. That's why I am exploiting and learning the implementation of TP kernel.

I see that I've missed some statements in the paper. Thank you for pointing this out.

I will also try the methods that @vbharadwaj-bk mentioned.

[asglover]

Very cool!

The code will be much more readable when the templates are instantiated. I'll try to post an example of a simple kernel (instantiated, not a template) in a bit.

[Glinttsd]

That will be very helpful. Thanks!

[asglover]

Linked is the example from our readme fully instantiated from the templates! It is a 'uvu' connection mode, so the weight multiplication will happen thread locally.

https://gist.github.com/asglover/d6cc3ace0a09a6e338c6e03e9653b17a

I also included the code to save the kernel, so you can play around and try different "tensor product problem" specifications and see how it changes the kernels.

[vbharadwaj-bk]

Glad to see the conversation here! Since there are no actionable items, shall we convert this to a discussion? You can continue to chat about the specifics of the kernel there as your FPGA work progresses.

[Glinttsd]

Yes, please. Thanks for that!