[Unity][Cutlass] Fix C source generation of dense operation

[creaitr]

It seems there is an issue while generating c source of dense operation using cutlass.
Even though the dense operation contains bias parameters, the generated c code doesn't reflect that bias correctly.

Please check the change and leave feedback.

[vinx13]

please add a test case

[vinx13]

if bias is used in the pattern (func name), it should exist in the annotation. cc @yelite

[yelite]

Maybe there is pattern with bias but doesn't have annotation? IIUC this code path is shared between relax and relay. @creaiter can you share sample where the bias parameter isn't generated correctly?

[creaitr]

@yelite
Sorry for the delayed response.

I tested this before the merge of unity branch into main.
When I tried to rebuild tvm with the newest branch, there was an issue from the cutlass_fpA_intB_gemm 3rdparty, that error: identifier "__hfma2" is undefined.
This error seems related with the setting of the compute compatibility version of nvcc, and I temporally resolved it by adding the code set(CMAKE_CUDA_ARCHITECTURES "native") at the tvm/CMakeLists.txt file. (Figuring out this took some time.)
Here is a minor question that have you encountered the same error and how did you resolve it?

For the bias issue of dense operation, when I tried to compile the below module with cutlass backend,

def dense_add_bias(data, weight=None, bias=None, units=None, **kwargs):
    name = kwargs.get("name")
    kwargs.pop("name")
    if not weight:
        weight = relay.var(name + "_weight")
    if not bias:
        bias = relay.var(name + "_bias")
    data = relay.nn.dense(data, weight, units, **kwargs)
    data = relay.nn.bias_add(data, bias, axis=-1)
    return data

there were some errors at the generated .cu file.
At the following generated code, I'm suspicious about the ${bias_arg} parts.

void tvmgen_default_cutlass_main_3_(DLTensor* cutlass_3_i0, DLTensor* cutlass_3_i1, DLTensor* cutlass_3_i2, DLTensor* out0){

  using ElementInputA = float;
  using ElementInputB = float;
  using ElementOutput = float;
  using ElementComputeEpilogue = float;

  // ... omitted

  using Gemm = Operation_cutlass_tensorop_s1688gemm_128x128_16x4_tn_align1;
  int M = 1;
  int N = 2048;
  int K = 1024;
  cutlass::gemm::GemmCoord problem_size(M, N, K);
  ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
  ElementComputeEpilogue beta = ElementComputeEpilogue(0);
  void* ptr_a = (void*)(cutlass_3_i0->data);
  void* ptr_b = (void*)(cutlass_3_i1->data);
  void* ptr_bias = (void*)(${bias_arg}->data);

  void* ptr_out = (void*)(out0->data);
  
  typename Gemm::Arguments arguments{
   problem_size,
   {static_cast<ElementInputA*>(ptr_a), K}, 
   {static_cast<ElementInputB*>(ptr_b), K}, 
   {static_cast<ElementOutput*>(ptr_bias), (${bias_arg}->ndim == 1 || ${bias_arg}->shape[${bias_arg}->ndim - 2] == 1) ? 0 : N}, 
   {static_cast<ElementOutput*>(ptr_out), N}, 
   {alpha},
   1
  };

  // ... omitted
}

After applying this PR, above error was resolved for me.

Actually, I'm not familiar with this tvm project. So you can feedback me, if what I changed is the wrong part.

[vinx13]

https://github.com/apache/tvm/blob/main/cmake/modules/CUDA.cmake#L47
We have a default value CMAKE_CUDA_ARTECTURES here, does that work? Right now it requires cuda arch >= 53 to compile

[creaitr]

@vinx13
It seems test_cutlass.py#L549 is testing the dense_bias modules.
If the test case means dffierent thing, please let me know.

[creaitr]

Oh, now I see the CUDA.cmake#L47.

In my case, the CMAKE_CUDA_ARCHITECTURES is already defined as 52. Thus, it is not changed to native.
(I compiled with cmake version 3.28.1)

[churvey]

if bias is used in the pattern (func name), it should exist in the annotation. cc @yelite

It seems bias_arg_idx only annotated in relax, so partition_for_cutlass is broken in relay