[Fix] NCCL API mismatch and NCCL primitive fix (#301)

* `count` in NCCL APIs refer to number of elements, not bytes;
* Fix nccl primitives. It now can handle dynamic shapes correctly.

python/hidet/distributed/group.py

@@ -81,7 +81,7 @@ def all_reduce(self, tensor: Tensor, op: str):
         assert not tensor.is_symbolic()
         assert tensor.device.is_cuda()
         addr = tensor.storage.addr
-        self._comm.all_reduce(addr, addr, tensor.nbytes, tensor.dtype, op)
+        self._comm.all_reduce(addr, addr, tensor.size, tensor.dtype, op)
 
 
 def create_nccl_group(store: Store, world_size: int, rank: int):

python/hidet/graph/ops/distributed.py

@@ -28,24 +28,21 @@ def __init__(self, x: TensorNode, op: str, comm_id: int = 0):
 
         super().__init__('all_reduce', inputs=[x], outputs=[y], attributes={'comm_id': comm_id, 'op': op})
 
-    def __str__(self):
-        return "all_reduce"
-
     def implement(self, target: Union[Target, str], working_dir: str) -> List[IRModule]:
         import hidet
         from hidet.ir.primitives.cuda.nccl import all_reduce as _all_reduce
         from hidet.lang import attrs
 
         dtype: DataType = self.inputs[0].type.dtype
         shape: Tuple[Expr, ...] = self.inputs[0].shape
-        nbytes = dtype.nbytes * prod(shape)
+        size = prod(shape)
 
         with hidet.script_module() as script_module:
 
             @hidet.script
             def launch(x: dtype[shape], y: dtype[shape]):
                 attrs.func_kind = 'public'
-                _all_reduce(x, y, nbytes, dtype, str_to_nccl_op(self.op), self.comm_id)
+                _all_reduce(x, y, size, dtype, str_to_nccl_op(self.op), self.comm_id)
 
         return [script_module.ir_module()]
 

python/hidet/ir/primitives/cuda/nccl.py

@@ -23,8 +23,12 @@ def all_reduce(sendbuff: Expr, recvbuff: Expr, count: Expr, dtype: DataType, op:
 
     comm = get_nccl_comm(comm_id)
     return BlackBoxStmt(
-        'ncclAllReduce({}, {}, {}, (ncclDataType_t){}, (ncclRedOp_t){}, '
-        '(ncclComm_t){}, (cudaStream_t){});'.format(
-            sendbuff, recvbuff, count, int(dtype_to_nccl(dtype)), int(op), comm, get_cuda_stream()
-        )
+        'ncclAllReduce({}, {}, {}, (ncclDataType_t){}, (ncclRedOp_t){}, ' '(ncclComm_t){}, (cudaStream_t){});',
+        sendbuff,
+        recvbuff,
+        count,
+        int(dtype_to_nccl(dtype)),
+        int(op),
+        comm,
+        get_cuda_stream(),
     )