[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