Optimized offset calculator in indexing ops

Signed-off-by: majing <[email protected]>
intel · Sep 19, 2024 · 94937b0 · 94937b0
1 parent da7692e
commit 94937b0
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/src/comm/TensorInfo.h b/src/comm/TensorInfo.h
@@ -137,16 +137,13 @@ struct IndexToOffset {
       return linearId;
     }
 
-#pragma unroll
-    for (int dim = XPU_MAX_TENSORINFO_DIMS - 1; dim >= 0; --dim) {
-      if (dim < info.dims) {
-        auto divider = at::detail::IntDivider<IndexType>(info.sizes[dim]);
-        auto divmod = divider.divmod(linearId);
-        linearId = divmod.div;
-        offset += divmod.mod * info.strides[dim];
-      }
+    for (int dim = info.dims - 1; dim > 0; --dim) {
+      IndexType curDimIndex = linearId % info.sizes[dim];
+      IndexType curDimOffset = curDimIndex * info.strides[dim];
+      offset += curDimOffset;
+      linearId /= info.sizes[dim];
     }
-    return offset;
+    return offset + linearId * info.strides[0];
   }
 };