[Mlir-commits] [mlir] [mlir][mesh] Add resharding spmdization on a 1D device mesh (PR #76179)

Boian Petkantchin llvmlistbot at llvm.org
Tue Jan 2 09:45:32 PST 2024


================
@@ -186,6 +209,29 @@ def Mesh_ShardOp : Mesh_Op<"shard", [Pure, SameOperandsAndResultType]> {
   }];
 }
 
+def Mesh_ProcessIndexOp : Mesh_Op<"process_index", [Pure, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
----------------
sogartar wrote:

I assume you mean to have the correspondence between multi index and flat/linear index.
Do you want such an operation explicitly in the IR that extracts the linear index?
If it is to be used during lowering we can have a C++ function that adds IR to flatten the index. There may be already somewhere in MLIR such a thing.

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


More information about the Mlir-commits mailing list