msccl: add templated kernel (#945)

* msccl: add templated kernel * Use defines to improve code readability * Fix kernel indexing and review feedback
2023-11-02 17:21:53 -07:00
parent 61aed56ca7
commit f484ff17b9
6 changed files with 96 additions and 73 deletions
@@ -105,7 +105,8 @@ function(expand_collectives FILE FUNC)
          #include \"primitives.h\"
          #include \"collectives.h\"
          #include \"devcomm.h\"
-          MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(${REDOP_CURRENT}, ${DATA_TYPE});")
+          MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(${REDOP_CURRENT}, ${DATA_TYPE}, false);
+          MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(${REDOP_CURRENT}, ${DATA_TYPE}, true);")
      else()
        file(WRITE ${FILE_NAME}
          "#include \"${FILE}.h\"