[MLIR][Python] Add shard Dialect Python Bindings

[Svoch]

Add Python bindings for shard dialect. Provide means for creating constructs in this dialect in Python.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Siavash Nazari (Svoch)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/162578.diff

5 Files Affected:

• (modified) mlir/python/CMakeLists.txt (+9)
• (added) mlir/python/mlir/dialects/ShardOps.td (+14)
• (added) mlir/python/mlir/dialects/shard.py (+25)
• (added) mlir/test/python/dialects/shard.py (+73)
• (modified) utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel (+32)

diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 9f5246de6bda0..20f07440df2c3 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -336,6 +336,15 @@ declare_mlir_dialect_python_bindings(
     dialects/memref.py
   DIALECT_NAME memref)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/ShardOps.td
+  SOURCES
+    dialects/shard.py
+  DIALECT_NAME shard
+  GEN_ENUM_BINDINGS)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
diff --git a/mlir/python/mlir/dialects/ShardOps.td b/mlir/python/mlir/dialects/ShardOps.td
new file mode 100644
index 0000000000000..f8527664df67b
--- /dev/null
+++ b/mlir/python/mlir/dialects/ShardOps.td
@@ -0,0 +1,14 @@
+//===-- ShardOps.td - Entry point for ShardOps bindings ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_SHARD_OPS
+#define PYTHON_BINDINGS_SHARD_OPS
+
+include "mlir/Dialect/Shard/IR/ShardOps.td"
+
+#endif
diff --git a/mlir/python/mlir/dialects/shard.py b/mlir/python/mlir/dialects/shard.py
new file mode 100644
index 0000000000000..b64ee53bb0665
--- /dev/null
+++ b/mlir/python/mlir/dialects/shard.py
@@ -0,0 +1,25 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._shard_ops_gen import *
+from ._shard_ops_gen import _Dialect
+from ._shard_enum_gen import *
+
+try:
+    from ..ir import *
+    from ._ods_common import (
+        get_default_loc_context as _get_default_loc_context,
+        _cext as _ods_cext,
+        get_op_result_or_op_results as _get_op_result_or_op_results,
+    )
+
+    from typing import Any, List, Union
+except ImportError as e:
+    raise RuntimeError("Error loading imports from extension module") from e
+
+
+# The shard dialect currently doesn't need custom Python implementations for its operations
+# like the arith dialect does for ConstantOp. Most operations can use the generated bindings.
+# If specialized Python methods are needed for specific operations in the future,
+# they can be added here using the @_ods_cext.register_operation decorator pattern.
\ No newline at end of file
diff --git a/mlir/test/python/dialects/shard.py b/mlir/test/python/dialects/shard.py
new file mode 100644
index 0000000000000..a6455ad217fff
--- /dev/null
+++ b/mlir/test/python/dialects/shard.py
@@ -0,0 +1,73 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import shard
+from mlir.dialects import func
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: TEST: testShardGrid
+@constructAndPrintInModule
+def testShardGrid():
+    # Test creating shard grids with different shapes
+    grid2d = shard.GridOp("grid_2d", [2, 2])
+    grid1d = shard.GridOp("grid_1d", [4])
+    grid_dynamic = shard.GridOp("grid_dynamic", [2, -1])  # -1 for dynamic dimension
+
+
+# CHECK: shard.grid @grid_2d(shape = 2x2)
+# CHECK: shard.grid @grid_1d(shape = 4)
+# CHECK: shard.grid @grid_dynamic(shape = 2x?)
+
+
+# CHECK-LABEL: TEST: testCollectiveOperations
+@constructAndPrintInModule
+def testCollectiveOperations():
+    # Create grid and types
+    grid = shard.GridOp("grid_2x2", [2, 2])
+    i32 = IntegerType.get_signless(32)
+    input_type = RankedTensorType.get([4, 2], i32)
+    gather_result_type = RankedTensorType.get([4, 4], i32)
+    
+    # Create a function to hold the operations
+    func_type = FunctionType.get([input_type], [input_type])
+    test_func = func.FuncOp("test_collectives", func_type)
+    
+    with InsertionPoint(test_func.add_entry_block()):
+        arg = test_func.entry_block.arguments[0]
+        
+        # All-gather operation
+        gather_op = shard.AllGatherOp(
+            input=arg,
+            grid=FlatSymbolRefAttr.get("grid_2x2"),
+            grid_axes=ArrayAttr.get([IntegerAttr.get(i32, 1)]),
+            gather_axis=IntegerAttr.get(i32, 1),
+            result=gather_result_type
+        )
+        
+        # All-reduce operation (ReductionKind might need different construction)
+        reduce_op = shard.AllReduceOp(
+            input=arg,
+            grid=FlatSymbolRefAttr.get("grid_2x2"),
+            reduction=IntegerAttr.get(IntegerType.get_signless(32), 1),  # 1 = sum from enum
+            result=input_type
+        )
+        
+        # Return the reduced result
+        func.ReturnOp([reduce_op])
+
+
+# CHECK: shard.grid @grid_2x2(shape = 2x2)
+# CHECK: func @test_collectives(%{{.*}}: tensor<4x2xi32>) -> tensor<4x2xi32>
+# CHECK: %{{.*}} = shard.all_gather %{{.*}} on @grid_2x2 grid_axes = [1] gather_axis = 1 : tensor<4x2xi32> -> tensor<4x4xi32>
+# CHECK: %{{.*}} = shard.all_reduce %{{.*}} on @grid_2x2 reduction = sum : tensor<4x2xi32> -> tensor<4x2xi32>
+
diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index 102c4161eb74c..72af4f08bde57 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -981,6 +981,38 @@ filegroup(
     ],
 )
 
+##---------------------------------------------------------------------------##
+# Shard dialect.
+##---------------------------------------------------------------------------##
+
+gentbl_filegroup(
+    name = "ShardOpsPyGen",
+    tbl_outs = {
+        "mlir/dialects/_shard_enum_gen.py": [
+            "-gen-python-enum-bindings",
+            "-bind-dialect=shard",
+        ],
+        "mlir/dialects/_shard_ops_gen.py": [
+            "-gen-python-op-bindings",
+            "-bind-dialect=shard",
+        ],
+    },
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/ShardOps.td",
+    deps = [
+        "//mlir:OpBaseTdFiles",
+        "//mlir:ShardTdFiles",
+    ],
+)
+
+filegroup(
+    name = "ShardOpsPyFiles",
+    srcs = [
+        "mlir/dialects/shard.py",
+        ":ShardOpsPyGen",
+    ],
+)
+
 ##---------------------------------------------------------------------------##
 # Shape dialect.
 ##---------------------------------------------------------------------------##

[github-actions]

✅ With the latest revision this PR passed the Python code formatter.

[Svoch]

@PragmaTwice @makslevental, thanks for the reviews :) I believe I've addressed your comments. PTAL

[Svoch]

@makslevental - Can you PTAL once more? The tests are updated and I've added an additional verification step to make sure we're actually building correct IR. IMO, we should be good to go.

[makslevental]

@Svoch i'm hitting "approve tests" as I see you've updated the PR...

[makslevental]

LGTM! Thanks for the bindings!

[makslevental]

@Svoch since the tests have passed, let me know when you're ready to merge

[Svoch]

@Svoch since the tests have passed, let me know when you're ready to merge

@makslevental Thank you so much! Let's go whenever you're ready. Can you please merge it?

[makslevental]

Done

[github-actions]

@Svoch Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!