[Mlir-commits] [llvm] [mlir] Add shard Dialect Python Bindings (PR #162578)

[Maksim Levental]

================
@@ -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
+ at 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?)
----------------
makslevental wrote:

these should be inside the tests themselves

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