[Unity][Cutlass] Fix C source generation of dense operation (apache#1…

…6476)

This commit fixes an issue that generates wrong c sources of dense operation using cutlass.

Co-authored-by: 진배 박 <jinbae@nexon.co.kr>

python/tvm/contrib/cutlass/gen_tensor_op.py

@@ -566,7 +566,10 @@ def get_flattened_batch_dim(arg_name, batch_rank):
         transposed = "transposed" in func_name or "dense" in func_name
         lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
         rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
-        bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
+        if "bias" in func_name:
+            bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", 2)
+        else:
+            bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
         residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
 
         lhs_arg = func_args[lhs_arg_idx]