Fix bindTensorMetaData for DID loop split

NVIDIA · Feb 8, 2025 · 5a3dcf8 · 5a3dcf8
1 parent cf1fac4
commit 5a3dcf8
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/csrc/evaluator_common.cpp b/csrc/evaluator_common.cpp
@@ -12,6 +12,7 @@
 #include <expr_evaluator.h>
 #include <instrumentation.h>
 #include <ir/utils.h>
+#include <multidevice/utils.h>
 #include <runtime/executor_kernel_arg.h>
 #include <tensor_metadata.h>
 
@@ -348,9 +349,11 @@ void PrecomputedValues::bindTensorMetaData(
       tensor.dim() == static_cast<int64_t>(logical_domain.size()),
       "Something went wrong configuring launch. Inputs do not match.");
 
-  for (const auto dim : c10::irange(logical_domain.size())) {
+  std::vector<int64_t> logical_sizes = unshardedSizes(tv, tensor.sizes());
+  for (const auto dim :
+       c10::irange(static_cast<int64_t>(logical_domain.size()))) {
     IterDomain* id = logical_domain[dim];
-    const auto dim_size = tensor.size(static_cast<int64_t>(dim));
+    const auto dim_size = logical_sizes.at(dim);
     if (id->isBroadcast()) {
       // DIDs are ignored for broadcast. See MultideviceShardingTest.Broadcast
       // and .ExpandedBroadcast.
@@ -359,13 +362,7 @@ void PrecomputedValues::bindTensorMetaData(
         bindValue(id->expandedExtent()->evaluatorIndex(), dim_size);
       }
     } else {
-      if (id->isDeviceDim()) {
-        bindValue(
-            id->extent()->evaluatorIndex(),
-            tv->getDeviceMesh().size(id->getParallelType()));
-      } else {
-        bindValue(id->extent()->evaluatorIndex(), dim_size);
-      }
+      bindValue(id->extent()->evaluatorIndex(), dim_size);
     }
   }