gemini-code-assist[bot] commented on code in PR #18352:
URL: https://github.com/apache/tvm/pull/18352#discussion_r2386171518
##########
src/runtime/disco/nccl/nccl.cc:
##########
@@ -176,19 +176,19 @@ void ScatterFromWorker0(ffi::Optional<Tensor> send, bool
in_group, Tensor recv)
if (is_sender) {
CHECK(send.defined()) << "ValueError: buffer `send` must be provided when
worker_id == 0.";
Tensor buffer = send.value();
- int64_t numel = buffer.Shape()->Product();
+ int64_t numel = buffer.Shape().Product();
CHECK_EQ(numel % num_receiver, 0) << "ValueError: Scattering evenly
requires that the number "
"of elements in the buffer to be "
"divisible by the number of workers,
but got numel = "
<< numel << " and " << num_receiver << "
workers.";
DataType dtype(buffer->dtype);
int64_t numel_per_shard = numel / num_receiver;
int64_t bytes_per_shard = numel_per_shard * dtype.bytes();
- CHECK_EQ(numel_per_shard, recv.Shape()->Product())
+ CHECK_EQ(numel_per_shard, recv.Shape().Product())
<< "ValueError: The number of elements in buffer `recv` must be the
same as each shard "
"of "
"buffer `send`. `send.size` is "
- << numel << ", but `recv.size` is " << recv.Shape()->Product() << ".";
+ << numel << ", but `recv.size` is " << recv.Shape().Product() << ".";
Review Comment:
To avoid calling `recv.Shape().Product()` multiple times and to improve
readability, it's better to store its result in a local variable.
```c
const int64_t recv_numel = recv.Shape().Product();
CHECK_EQ(numel_per_shard, recv_numel)
<< "ValueError: The number of elements in buffer `recv` must be the
same as each shard "
"of "
"buffer `send`. `send.size` is "
<< numel << ", but `recv.size` is " << recv_numel << ".";
```
##########
src/runtime/disco/nccl/nccl.cc:
##########
@@ -223,19 +223,19 @@ void GatherToWorker0(Tensor send, bool in_group,
ffi::Optional<Tensor> recv) {
if (is_sender) {
CHECK(recv.defined()) << "ValueError: buffer `recv` must be provided when
worker_id == 0.";
Tensor buffer = recv.value();
- int64_t numel = buffer.Shape()->Product();
+ int64_t numel = buffer.Shape().Product();
CHECK_EQ(numel % num_receiver, 0) << "ValueError: Gathering evenly
requires that the number "
"of elements in the buffer to be "
"divisible by the number of workers,
but got numel = "
<< numel << " and " << num_receiver << "
workers.";
DataType dtype(buffer->dtype);
int64_t numel_per_shard = numel / num_receiver;
int64_t bytes_per_shard = numel_per_shard * dtype.bytes();
- CHECK_EQ(numel_per_shard, send.Shape()->Product())
+ CHECK_EQ(numel_per_shard, send.Shape().Product())
<< "ValueError: The number of elements in buffer `send` must be the
same as each shard "
"of "
"buffer `recv`. `recv.size` is "
- << numel << ", but `send.size` is " << send.Shape()->Product() << ".";
+ << numel << ", but `send.size` is " << send.Shape().Product() << ".";
Review Comment:
To avoid calling `send.Shape().Product()` multiple times and to improve
readability, it's better to store its result in a local variable.
```c
const int64_t send_numel = send.Shape().Product();
CHECK_EQ(numel_per_shard, send_numel)
<< "ValueError: The number of elements in buffer `send` must be the
same as each shard "
"of "
"buffer `recv`. `recv.size` is "
<< numel << ", but `send.size` is " << send_numel << ".";
```
##########
src/runtime/disco/nccl/nccl.cc:
##########
@@ -150,13 +150,13 @@ void BroadcastFromWorker0(ffi::Optional<Tensor> send,
bool in_group, Tensor recv
const void* send_data = [&]() -> const void* {
if (is_sender) {
CHECK(send.defined());
- CHECK(send.value().Shape()->Product() == recv.Shape()->Product());
+ CHECK(send.value().Shape().Product() == recv.Shape().Product());
return send.value()->data;
} else {
return nullptr;
}
}();
- int64_t numel = recv.Shape()->Product();
+ int64_t numel = recv.Shape().Product();
Review Comment:
To avoid calling `recv.Shape().Product()` multiple times and to improve
readability by removing the immediately-invoked lambda expression (IIFE), you
can calculate the number of elements once and use a simple if/else block.
```c
const int64_t numel = recv.Shape().Product();
const void* send_data;
if (is_sender) {
CHECK(send.defined());
CHECK(send.value().Shape().Product() == numel);
send_data = send.value()->data;
} else {
send_data = nullptr;
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
