Skip to content

Commit c618791

Browse files
author
Siavash Nazari
committed
[MLIR][Python] Add shard Dialect Python Bindings
1 parent 37af81f commit c618791

File tree

5 files changed

+128
-0
lines changed

5 files changed

+128
-0
lines changed

‎mlir/python/CMakeLists.txt‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -336,6 +336,15 @@ declare_mlir_dialect_python_bindings(
336336
dialects/memref.py
337337
DIALECT_NAME memref)
338338

339+
declare_mlir_dialect_python_bindings(
340+
ADD_TO_PARENT MLIRPythonSources.Dialects
341+
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
342+
TD_FILE dialects/ShardOps.td
343+
SOURCES
344+
dialects/shard.py
345+
DIALECT_NAME shard
346+
GEN_ENUM_BINDINGS)
347+
339348
declare_mlir_dialect_python_bindings(
340349
ADD_TO_PARENT MLIRPythonSources.Dialects
341350
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"

‎mlir/python/mlir/dialects/ShardOps.td‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
//===-- ShardOps.td - Entry point for ShardOps bindings ---------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef PYTHON_BINDINGS_SHARD_OPS
10+
#define PYTHON_BINDINGS_SHARD_OPS
11+
12+
include "mlir/Dialect/Shard/IR/ShardOps.td"
13+
14+
#endif

‎mlir/python/mlir/dialects/shard.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2+
# See https://llvm.org/LICENSE.txt for license information.
3+
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4+
5+
from ._shard_ops_gen import *
6+
from ._shard_enum_gen import *

‎mlir/test/python/dialects/shard.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# RUN: %PYTHON %s | FileCheck %s
2+
3+
from mlir.ir import *
4+
from mlir.dialects import shard
5+
from mlir.dialects import func
6+
7+
8+
def constructAndPrintInModule(f):
9+
print("\nTEST:", f.__name__)
10+
with Context(), Location.unknown():
11+
module = Module.create()
12+
with InsertionPoint(module.body):
13+
f()
14+
print(module)
15+
return f
16+
17+
18+
# CHECK-LABEL: TEST: testShardGrid
19+
@constructAndPrintInModule
20+
def testShardGrid():
21+
# Test creating shard grids with different shapes
22+
grid2d = shard.GridOp("grid_2d", [2, 2])
23+
grid1d = shard.GridOp("grid_1d", [4])
24+
grid_dynamic = shard.GridOp("grid_dynamic", [2, -1]) # -1 for dynamic dimension
25+
26+
# CHECK: shard.grid @grid_2d(shape = 2x2)
27+
# CHECK: shard.grid @grid_1d(shape = 4)
28+
# CHECK: shard.grid @grid_dynamic(shape = 2x?)
29+
30+
31+
# CHECK-LABEL: TEST: testCollectiveOperations
32+
@constructAndPrintInModule
33+
def testCollectiveOperations():
34+
# Create grid and types
35+
grid = shard.GridOp("grid_2x2", [2, 2])
36+
i32 = IntegerType.get_signless(32)
37+
input_type = RankedTensorType.get([4, 2], i32)
38+
gather_result_type = RankedTensorType.get([4, 4], i32)
39+
40+
# Create a function to hold the operations
41+
func_type = FunctionType.get([input_type], [input_type])
42+
test_func = func.FuncOp("test_collectives", func_type)
43+
44+
with InsertionPoint(test_func.add_entry_block()):
45+
arg = test_func.entry_block.arguments[0]
46+
47+
gather_op = shard.AllGatherOp(
48+
input=arg,
49+
grid=FlatSymbolRefAttr.get("grid_2x2"),
50+
grid_axes=ArrayAttr.get([IntegerAttr.get(i32, 1)]),
51+
gather_axis=IntegerAttr.get(i32, 1),
52+
result=gather_result_type,
53+
)
54+
55+
reduce_op = shard.AllReduceOp(
56+
input=arg,
57+
grid=FlatSymbolRefAttr.get("grid_2x2"),
58+
reduction=shard.ReductionKind.Sum,
59+
result=input_type,
60+
)
61+
62+
func.ReturnOp([reduce_op])
63+
64+
# CHECK: shard.grid @grid_2x2(shape = 2x2)
65+
# CHECK: func @test_collectives(%{{.*}}: tensor<4x2xi32>) -> tensor<4x2xi32>
66+
# CHECK: %{{.*}} = shard.all_gather %{{.*}} on @grid_2x2 grid_axes = [1] gather_axis = 1 : tensor<4x2xi32> -> tensor<4x4xi32>
67+
# CHECK: %{{.*}} = shard.all_reduce %{{.*}} on @grid_2x2 reduction = sum : tensor<4x2xi32> -> tensor<4x2xi32>

‎utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -981,6 +981,38 @@ filegroup(
981981
],
982982
)
983983

984+
##---------------------------------------------------------------------------##
985+
# Shard dialect.
986+
##---------------------------------------------------------------------------##
987+
988+
gentbl_filegroup(
989+
name = "ShardOpsPyGen",
990+
tbl_outs = {
991+
"mlir/dialects/_shard_enum_gen.py": [
992+
"-gen-python-enum-bindings",
993+
"-bind-dialect=shard",
994+
],
995+
"mlir/dialects/_shard_ops_gen.py": [
996+
"-gen-python-op-bindings",
997+
"-bind-dialect=shard",
998+
],
999+
},
1000+
tblgen = "//mlir:mlir-tblgen",
1001+
td_file = "mlir/dialects/ShardOps.td",
1002+
deps = [
1003+
"//mlir:OpBaseTdFiles",
1004+
"//mlir:ShardTdFiles",
1005+
],
1006+
)
1007+
1008+
filegroup(
1009+
name = "ShardOpsPyFiles",
1010+
srcs = [
1011+
"mlir/dialects/shard.py",
1012+
":ShardOpsPyGen",
1013+
],
1014+
)
1015+
9841016
##---------------------------------------------------------------------------##
9851017
# Shape dialect.
9861018
##---------------------------------------------------------------------------##

0 commit comments

Comments
 (0)