[Mlir-commits] [mlir] [llvm] [mlir][mesh] Add spmdization pass (PR #80518)

Boian Petkantchin llvmlistbot at llvm.org
Mon Feb 5 09:14:26 PST 2024


================
@@ -78,6 +92,35 @@ int64_t collectiveProcessGroupSize(MeshAxesRange &&meshAxes,
   return res;
 }
 
+inline int64_t shardDimension(int64_t dim, int64_t shardCount) {
+  if (ShapedType::isDynamic(dim) || ShapedType::isDynamic(shardCount))
+    return ShapedType::kDynamic;
+
+  assert(dim % shardCount == 0);
+  return ceilDiv(dim, shardCount);
+}
+
+inline int64_t unshardDimension(int64_t dim, int64_t shardCount) {
----------------
sogartar wrote:

I renamed it.

https://github.com/llvm/llvm-project/pull/80518


More information about the Mlir-commits mailing list