[Mlir-commits] [mlir] [Mesh] initialize mesh dialect (PR #68007)
Chengji Yao
llvmlistbot at llvm.org
Mon Oct 2 09:59:19 PDT 2023
https://github.com/yaochengji created https://github.com/llvm/llvm-project/pull/68007
This is the 1st PR of [Mesh sharding RFC](https://discourse.llvm.org/t/open-mlir-meeting-9-28-2023-rfc-sharding-framework-design-for-device-mesh/73695), includes
- mesh.cluster op
- mesh.shard op (the mesh.annotate op in the RFC slides, the name is modified a bit from @stellaraccident 's advice, which I think might be a bit more concise)
- MeshSharding attribute
>From 12b70771999f1208e204ff4ebc141b950f569610 Mon Sep 17 00:00:00 2001
From: Chengji Yao <yaochengji at hotmial.com>
Date: Sat, 30 Sep 2023 00:24:38 +0000
Subject: [PATCH] [Mesh] initialize mesh dialect
---
mlir/include/mlir/Dialect/CMakeLists.txt | 1 +
mlir/include/mlir/Dialect/Mesh/CMakeLists.txt | 1 +
.../mlir/Dialect/Mesh/IR/CMakeLists.txt | 8 ++
mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td | 91 ++++++++++++
mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h | 35 +++++
mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td | 132 ++++++++++++++++++
mlir/include/mlir/InitAllDialects.h | 2 +
mlir/lib/Dialect/CMakeLists.txt | 1 +
mlir/lib/Dialect/Mesh/CMakeLists.txt | 15 ++
mlir/lib/Dialect/Mesh/MeshOps.cpp | 122 ++++++++++++++++
mlir/test/Dialect/Mesh/invalid.mlir | 65 +++++++++
mlir/test/Dialect/Mesh/ops.mlir | 69 +++++++++
12 files changed, 542 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/Mesh/CMakeLists.txt
create mode 100644 mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt
create mode 100644 mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
create mode 100644 mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h
create mode 100644 mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td
create mode 100644 mlir/lib/Dialect/Mesh/CMakeLists.txt
create mode 100644 mlir/lib/Dialect/Mesh/MeshOps.cpp
create mode 100644 mlir/test/Dialect/Mesh/invalid.mlir
create mode 100644 mlir/test/Dialect/Mesh/ops.mlir
diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt
index 9082d9633339c9f..1c4569ecfa58485 100644
--- a/mlir/include/mlir/Dialect/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/CMakeLists.txt
@@ -19,6 +19,7 @@ add_subdirectory(Linalg)
add_subdirectory(LLVMIR)
add_subdirectory(Math)
add_subdirectory(MemRef)
+add_subdirectory(Mesh)
add_subdirectory(MLProgram)
add_subdirectory(NVGPU)
add_subdirectory(OpenACC)
diff --git a/mlir/include/mlir/Dialect/Mesh/CMakeLists.txt b/mlir/include/mlir/Dialect/Mesh/CMakeLists.txt
new file mode 100644
index 000000000000000..f33061b2d87cffc
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Mesh/CMakeLists.txt
@@ -0,0 +1 @@
+add_subdirectory(IR)
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt
new file mode 100644
index 000000000000000..f1d15891d029e1a
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt
@@ -0,0 +1,8 @@
+add_mlir_dialect(MeshOps mesh)
+add_mlir_doc(MeshOps MeshOps Dialects/ -gen-op-doc)
+
+set(LLVM_TARGET_DEFINITIONS MeshBase.td)
+mlir_tablegen(MeshOpsAttributes.h.inc -gen-attrdef-decls)
+mlir_tablegen(MeshOpsAttributes.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRMeshOpsAttrIncGen)
+add_mlir_doc(MeshOps MeshAttributes Dialects/ -gen-attrdef-doc)
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td b/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
new file mode 100644
index 000000000000000..05b5eb75c66d3f2
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
@@ -0,0 +1,91 @@
+//===- MeshBae.td - Mesh Dialect Dialect -------------------*- tablegen -*-===//
+//
+// 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 MLIR_DIALECT_MESH_IR_MESHBASE_TD
+#define MLIR_DIALECT_MESH_IR_MESHBASE_TD
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/BuiltinTypeInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Mesh Dialect
+//===----------------------------------------------------------------------===//
+
+def Mesh_Dialect : Dialect {
+ let name = "mesh";
+ let cppNamespace = "::mlir::mesh";
+
+ let description = [{
+ The `mesh` dialect contains a set of attributes, operations, interfaces that
+ are useful for representing sharding and communication on device mesh
+ cluster.
+ }];
+
+ let dependentDialects = [
+ "arith::ArithDialect"
+ ];
+
+ let useDefaultAttributePrinterParser = 1;
+ let hasConstantMaterializer = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Mesh Attribute
+//===----------------------------------------------------------------------===//
+
+def MeshSharding : AttrDef<Mesh_Dialect, "MeshSharding"> {
+ let mnemonic = "shard";
+
+ let parameters = (ins
+ OptionalParameter<"::mlir::SymbolRefAttr">:$cluster,
+ ArrayRefParameter<"::mlir::DenseI64ArrayAttr">:$axes
+ );
+
+ let summary = "Attribute that extends tensor type to distributed tensor type.";
+
+ let description = [{
+ The mesh.shard attribute contains two attribute in it:
+ 1. `cluster`: this attribute is a SymbolRefAttr that refers to the mesh
+ cluster where the distributed tensor is placed.
+
+ 2. `axes`: is an array composed of int64_t sub-arrays. The outer array's
+ maximum size is the `rank` of the related tensor plus one. For the i-th
+ sub-array, if its value is [x, y]:
+ - When i < `rank`, it indicates that the tensor's i-th dimension is sharded
+ along the x and y axes of the device mesh.
+ - When i == `rank`, it signifies that the tensor represents a partial sum
+ along the x and y axes. More partial types could be introduced if needed,
+ e.g. partial-max, partial-min and even partial-generic.
+
+ Example:
+
+ ```
+ mesh.cluster @mesh0(rank = 3, dim_sizes = [2, 2, 4])
+
+ // The tensor is fully replicated on @mesh0.
+ // Currently, there must be at least one sub-array present in axes, even
+ // if it's empty. Otherwise, a parsing error will occur.
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[]]>>
+
+ // The tensor is sharded on the first dimension along axis 0 of @mesh0
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>
+
+ // The tensor is sharded on the first dimension along axis 0 of @mesh0 and
+ // it is also a partial-sum along axis 1.
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[0], [], [1]]>
+
+ // Could also be used in the attribute of mesh.shard op
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
+ ```
+ }];
+ let assemblyFormat = "`<` ($cluster^ `,`)? `[` $axes `]` `>`";
+ let genVerifyDecl = 1;
+}
+
+#endif // MLIR_DIALECT_MESH_IR_MESHBASE_TD
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h
new file mode 100644
index 000000000000000..69159fc0d8cca89
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h
@@ -0,0 +1,35 @@
+//===- MeshOps.h - Mesh Dialect Operations ----------------------*- C++ -*-===//
+//
+// 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 MLIR_DIALECT_MESH_IR_MESHOPS_H
+#define MLIR_DIALECT_MESH_IR_MESHOPS_H
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+namespace mlir {
+
+namespace mesh {
+
+constexpr StringRef getMeshClusterAttrName() { return "mesh_cluster"; }
+
+} // namespace mesh
+
+} // namespace mlir
+
+#include "mlir/Dialect/Mesh/IR/MeshOpsDialect.h.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Mesh/IR/MeshOpsAttributes.h.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Mesh/IR/MeshOps.h.inc"
+
+#endif // MLIR_DIALECT_MESH_IR_MESHOPS_H
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td
new file mode 100644
index 000000000000000..bbd4891a2947517
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td
@@ -0,0 +1,132 @@
+//===-- MeshOps.td - Mesh dialect operation definitions ----*- tablegen -*-===//
+//
+// 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 MLIR_DIALECT_MESH_IR_MESHOPS_TD
+#define MLIR_DIALECT_MESH_IR_MESHOPS_TD
+
+include "mlir/Dialect/Mesh/IR/MeshBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/BuiltinTypes.td"
+include "mlir/IR/SymbolInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Mesh Dialect operations.
+//===----------------------------------------------------------------------===//
+
+class Mesh_Op<string mnemonic, list<Trait> traits = []> :
+ Op<Mesh_Dialect, mnemonic, traits> {
+}
+
+def Mesh_ClusterOp : Mesh_Op<"cluster", [Symbol]> {
+ let summary = "representing a mesh cluster";
+ let description = [{
+ The mesh.cluster operation is a symbol operation that identifies a specific
+ mesh cluster. The operation has three attributes:
+
+ 1. `sym_name`: This attribute uniquely identifies the name of the mesh
+ cluster. This name serves as a symbolic reference to the cluster throughout
+ the MLIR module, allowing for consistent referencing and easier debugging.
+
+ 2. `rank`: This attribute specifies the number of axes of the cluster. The
+ rank indicates the dimensionality of the mesh cluster and can be used to
+ determine the layout and the addressing space of the computation distributed
+ across the mesh.
+
+ 3. `dim_sizes`: This attribute represents the device assignment along the
+ axes of the cluster. Each integer in the array corresponds to the number of
+ devices along a specific axis. If an integer value is <= 0, it implies that
+ the number of devices along that axis is unknown. This flexibility allows
+ for dynamic device assignment or configurations where the exact number of
+ devices might not be determined during compile time.
+
+ Example:
+ ```
+ // A device mesh cluster with 3 axes, the totol device number is 4 * 8 * 12
+ // The dimension sizes are 4, 8, 12
+ mesh.cluster @mesh0(rank = 3, dim_sizes = [4, 8, 12])
+
+ // A device mesh cluster with 2 axes, the totol device number is unknown
+ // The first dimension size is 4 and the second is unknown
+ mesh.cluster @mesh1(rank = 2, dim_sizes = [4])
+
+ // A device mesh cluster with 2 axes, the totol device number is unknown
+ // The first dimension size is unknown and the second is 4
+ mesh.cluster @mesh2(rank = 2, dim_sizes = [0, 4])
+
+ // A device mesh cluster with 2 axes, the number of devices along both axes
+ // is unknown
+ mesh.cluster @mesh3(rank = 2)
+
+ // A func op's default mesh cluster is @mesh0
+ func.func @foo() -> () attributes { mesh_cluster = @mesh0 } {
+ ...
+ }
+
+ // Used in the mesh sharding attribute to extend the standard tensor to
+ // distributed
+ tensor<4x8xf32, #mesh.shard<[[0]] @mesh0>
+ ```
+ }];
+ let arguments = (ins
+ SymbolNameAttr:$sym_name,
+ I64Attr:$rank,
+ DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$dim_sizes
+ );
+ let assemblyFormat = "$sym_name `(` `rank` `=` $rank (`,` `dim_sizes` `=` $dim_sizes^)? `)` attr-dict";
+ let hasVerifier = 1;
+}
+
+def Mesh_ShardOp : Mesh_Op<"shard", [Pure, SameOperandsAndResultType]> {
+ let summary = "Annotate on how a tensor is sharded across a mesh cluster.";
+ let description = [{
+ The mesh.shard operation is designed to specify and guide the sharding
+ behavior of a tensor value across a mesh topology. This operation has one
+ operand and two attributes:
+
+ 1. `input`: This operand represents the tensor value that needs to be
+ annotated for sharding.
+
+ 2. `shard`: This attribute is type of `MeshSharding`, which is the core data
+ structure to represent distributed tensor in mesh cluster.
+
+ 3. `as_result`: A boolean attribute addressing the scenario when a tensor's
+ sharding annotation differs based on its context of use (either as a result
+ or an operand). If true, the sharding applies to the operation that defines
+ the tensor value. If false, the sharding pertains to specific users of the
+ tensor value, indicating how it should be considered when used as an operand
+ in subsequent operations.
+
+ Example:
+ ```
+ // The first mesh.shard op applies to %arg0, the second mesh.shard op
+ // applies for the operand of op0, the third mesh.shard op applies for the
+ // operand of op2
+ func.func @two_users(%arg0 : tensor<4x8xf32>) -> () {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
+ %1 = mesh.shard %0 to <@mesh0, [[1]]> {as_result = false} : tensor<4x8xf32>
+ %2 = mesh.shard %0 to <@mesh0, [[2]]> {as_result = false} : tensor<4x8xf32>
+ "op0"(%1) : ...
+ "op1"(%2) : ...
+ ...
+ }
+ ```
+ }];
+ let arguments = (ins
+ Builtin_RankedTensor:$src,
+ MeshSharding:$shard,
+ DefaultValuedAttr<BoolAttr, "true">:$as_result
+ );
+ let results = (outs
+ Builtin_RankedTensor:$result
+ );
+ let assemblyFormat = "$src `to` $shard attr-dict `:` type($result)";
+ let hasVerifier = 1;
+}
+
+#endif // MLIR_DIALECT_MESH_IR_MESHOPS_TD
diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index 5ec36a7f289e586..cbdc29020367a01 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -55,6 +55,7 @@
#include "mlir/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.h"
#include "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
+#include "mlir/Dialect/Mesh/IR/MeshOps.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
@@ -117,6 +118,7 @@ inline void registerAllDialects(DialectRegistry ®istry) {
LLVM::LLVMDialect,
math::MathDialect,
memref::MemRefDialect,
+ mesh::MeshDialect,
ml_program::MLProgramDialect,
nvgpu::NVGPUDialect,
NVVM::NVVMDialect,
diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt
index d9f6b0fb7e63c2a..68776a695cac4d4 100644
--- a/mlir/lib/Dialect/CMakeLists.txt
+++ b/mlir/lib/Dialect/CMakeLists.txt
@@ -19,6 +19,7 @@ add_subdirectory(Linalg)
add_subdirectory(LLVMIR)
add_subdirectory(Math)
add_subdirectory(MemRef)
+add_subdirectory(Mesh)
add_subdirectory(MLProgram)
add_subdirectory(NVGPU)
add_subdirectory(OpenACC)
diff --git a/mlir/lib/Dialect/Mesh/CMakeLists.txt b/mlir/lib/Dialect/Mesh/CMakeLists.txt
new file mode 100644
index 000000000000000..fc3f7c96f05c77f
--- /dev/null
+++ b/mlir/lib/Dialect/Mesh/CMakeLists.txt
@@ -0,0 +1,15 @@
+add_mlir_dialect_library(MLIRMeshDialect
+ MeshOps.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Mesh
+
+ DEPENDS
+ MLIRMeshOpsAttrIncGen
+ MLIRMeshOpsIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRArithDialect
+ MLIRIR
+ MLIRSupport
+)
diff --git a/mlir/lib/Dialect/Mesh/MeshOps.cpp b/mlir/lib/Dialect/Mesh/MeshOps.cpp
new file mode 100644
index 000000000000000..fd9f332e204de45
--- /dev/null
+++ b/mlir/lib/Dialect/Mesh/MeshOps.cpp
@@ -0,0 +1,122 @@
+//===- MeshOps.cpp - Mesh Dialect Operations ------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Mesh/IR/MeshOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+#define DEBUG_TYPE "mesh-ops"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
+
+using namespace mlir;
+using namespace mlir::mesh;
+
+#include "mlir/Dialect/Mesh/IR/MeshOpsDialect.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// Mesh dialect
+//===----------------------------------------------------------------------===//
+
+void MeshDialect::initialize() {
+ addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/Mesh/IR/MeshOps.cpp.inc"
+ >();
+ addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/Dialect/Mesh/IR/MeshOpsAttributes.cpp.inc"
+ >();
+}
+
+Operation *MeshDialect::materializeConstant(OpBuilder &builder, Attribute value,
+ Type type, Location loc) {
+ return arith::ConstantOp::materialize(builder, value, type, loc);
+}
+
+//===----------------------------------------------------------------------===//
+// mesh.cluster op
+//===----------------------------------------------------------------------===//
+
+LogicalResult ClusterOp::verify() {
+ ArrayRef<int64_t> dimSizes = getDimSizes();
+ size_t rank = getRank();
+
+ if (rank == 0)
+ return emitOpError("rank of cluster is expected to be a positive integer");
+
+ if (dimSizes.size() > rank)
+ return emitOpError(
+ "rank of dim_sizes is not expected to be larger than rank of cluster");
+
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// mesh.shard op
+//===----------------------------------------------------------------------===//
+
+LogicalResult ShardOp::verify() {
+ bool asResult = getAsResult();
+ if (asResult) {
+ Value src = getSrc();
+ Operation *defOp = src.getDefiningOp();
+ if (llvm::isa_and_nonnull<ShardOp>(defOp))
+ return emitOpError("two mesh.shard ops with as_result = true are not "
+ "expected to be stacked together");
+
+ unsigned numShard = llvm::count_if(src.getUsers(), [](Operation *user) {
+ return llvm::isa<ShardOp>(user);
+ });
+ if (numShard > 1)
+ return emitOpError(
+ "when than one mesh.shard ops operate on the same tensor, all of "
+ "their as_result attributes are expected to be false");
+
+ } else {
+ ShardOp defShardOp = getSrc().getDefiningOp<ShardOp>();
+ if (defShardOp && !defShardOp.getAsResult())
+ return emitOpError("two mesh.shard ops with as_result = false are not "
+ "expected to be stacked together");
+ }
+
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// mesh.shard op
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+MeshShardingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+ SymbolRefAttr, ArrayRef<DenseI64ArrayAttr> axes) {
+ // TODO: At present cluster symbol ref is not verified. This is due to the
+ // difficulty in fetching the corresponding symbol op based on an attribute.
+
+ DenseSet<int64_t> visitedAxes;
+ for (DenseI64ArrayAttr subAxes : axes) {
+ ArrayRef<int64_t> subAxesArray = subAxes.asArrayRef();
+ for (int64_t axis : subAxesArray) {
+ if (!visitedAxes.insert(axis).second)
+ return emitError() << "mesh axis duplicated";
+ }
+ }
+
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Mesh/IR/MeshOps.cpp.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Mesh/IR/MeshOpsAttributes.cpp.inc"
diff --git a/mlir/test/Dialect/Mesh/invalid.mlir b/mlir/test/Dialect/Mesh/invalid.mlir
new file mode 100644
index 000000000000000..8898ab000e6b6ed
--- /dev/null
+++ b/mlir/test/Dialect/Mesh/invalid.mlir
@@ -0,0 +1,65 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s
+
+// expected-error at +1 {{rank of cluster is expected to be a positive integer}}
+mesh.cluster @mesh0(rank = 0)
+
+// -----
+
+// expected-error at +1 {{rank of dim_sizes is not expected to be larger than rank of cluster}}
+mesh.cluster @mesh0(rank = 2, dim_sizes = [2, 3, 4])
+
+// -----
+
+mesh.cluster @mesh0(rank = 2, dim_sizes = [2, 4])
+
+func.func @mesh_shard_op_stacked_true_true(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %0 = mesh.shard %arg0 to <@mesh0, [[], [0]]> : tensor<4x8xf32>
+ // expected-error at +1 {{two mesh.shard ops with as_result = true are not expected to be stacked together}}
+ %1 = mesh.shard %0 to <@mesh0, [[], [0, 1]]> : tensor<4x8xf32>
+ return %1 : tensor<4x8xf32>
+}
+
+// -----
+
+mesh.cluster @mesh0(rank = 2, dim_sizes = [2, 4])
+
+func.func @mesh_shard_op_stacked_false_false(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %0 = mesh.shard %arg0 to <@mesh0, [[], [0]]> {as_result = false} : tensor<4x8xf32>
+ // expected-error at +1 {{two mesh.shard ops with as_result = false are not expected to be stacked together}}
+ %1 = mesh.shard %0 to <@mesh0, [[], [0, 1]]> {as_result = false} : tensor<4x8xf32>
+ return %1 : tensor<4x8xf32>
+}
+
+// -----
+
+mesh.cluster @mesh0(rank = 2, dim_sizes = [2, 4])
+
+func.func @mesh_shard_ops_on_same_tensor(%arg0 : tensor<4x8xf32>) ->
+ (tensor<4x8xf32>, tensor<4x8xf32>) {
+ // expected-error at +1 {{when than one mesh.shard ops operate on the same tensor, all of their as_result attributes are expected to be false}}
+ %0 = mesh.shard %arg0 to <@mesh0, [[], [0]]> : tensor<4x8xf32>
+ %1 = mesh.shard %arg0 to <@mesh0, [[], [0, 1]]> : tensor<4x8xf32>
+ return %0, %1 : tensor<4x8xf32>, tensor<4x8xf32>
+}
+
+// -----
+
+mesh.cluster @mesh0(rank = 2, dim_sizes = [2, 4])
+
+func.func @mesh_axis_duplicated(
+ // expected-error at +1 {{mesh axis duplicated}}
+ %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0], [0]]>>) ->
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[0], [0]]>> {
+ return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0], [0]]>>
+}
+
+// -----
+
+mesh.cluster @mesh0(rank = 2, dim_sizes = [2, 4])
+
+func.func @mesh_axis_duplicated_2(
+ // expected-error at +1 {{mesh axis duplicated}}
+ %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0, 0]]>>) ->
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[0, 0]]>> {
+ return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0, 0]]>>
+}
diff --git a/mlir/test/Dialect/Mesh/ops.mlir b/mlir/test/Dialect/Mesh/ops.mlir
new file mode 100644
index 000000000000000..c87f07257d66096
--- /dev/null
+++ b/mlir/test/Dialect/Mesh/ops.mlir
@@ -0,0 +1,69 @@
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// CHECK: mesh.cluster @mesh0
+mesh.cluster @mesh0(rank = 3, dim_sizes = [2, 2, 4])
+
+// CHECK: mesh.cluster @mesh1
+mesh.cluster @mesh1(rank = 2, dim_sizes = [4])
+
+// CHECK: mesh.cluster @mesh2
+mesh.cluster @mesh2(rank = 2, dim_sizes = [0, 4])
+
+// CHECK: mesh.cluster @mesh3
+mesh.cluster @mesh3(rank = 2)
+
+// CHECK-LABEL: func @mesh_shard_encoding_fully_replicated
+func.func @mesh_shard_encoding_fully_replicated(
+ %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[]]>>) ->
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[]]>> {
+ return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[]]>>
+}
+
+// CHECK-LABEL: func @mesh_shard_encoding_1st_dim
+func.func @mesh_shard_encoding_1st_dim(
+ %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>>) ->
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>> {
+ return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>>
+}
+
+// CHECK-LABEL: func @mesh_shard_encoding_2nd_dim
+func.func @mesh_shard_encoding_2nd_dim(
+ %arg0 : tensor<4x8xf32, #mesh.shard<@mesh1, [[], [0]]>>) ->
+ tensor<4x8xf32, #mesh.shard<@mesh1, [[], [0]]>> {
+ return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh1, [[], [0]]>>
+}
+
+// CHECK-LABEL: func @mesh_shard_encoding_1st_and_3rd_dim
+func.func @mesh_shard_encoding_1st_and_3rd_dim(
+ %arg0 : tensor<4x8x16xf32, #mesh.shard<@mesh3, [[0], [], [1]]>>) ->
+ tensor<4x8x16xf32, #mesh.shard<@mesh3, [[0], [], [1]]>> {
+ return %arg0 : tensor<4x8x16xf32, #mesh.shard<@mesh3, [[0], [], [1]]>>
+}
+
+// CHECK-LABEL: func @mesh_shard_op_1st_dim
+func.func @mesh_shard_op_1st_dim(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
+ return %0 : tensor<4x8xf32>
+}
+
+// CHECK-LABEL: func @mesh_shard_op_2nd_dim
+func.func @mesh_shard_op_2nd_dim(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %0 = mesh.shard %arg0 to <@mesh1, [[], [0]]> : tensor<4x8xf32>
+ return %0 : tensor<4x8xf32>
+}
+
+// CHECK-LABEL: func @mesh_shard_op_1st_and_3rd_dim
+func.func @mesh_shard_op_1st_and_3rd_dim(
+ %arg0 : tensor<4x8x16xf32>) -> tensor<4x8x16xf32> {
+ %0 = mesh.shard %arg0 to <@mesh3, [[0], [], [1]]> : tensor<4x8x16xf32>
+ return %0 : tensor<4x8x16xf32>
+}
+
+// CHECK-LABEL: func @mesh_shard_op_two_users
+func.func @mesh_shard_op_two_users(%arg0 : tensor<4x8xf32>) ->
+ (tensor<4x8xf32>, tensor<4x8xf32>) {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
+ %1 = mesh.shard %0 to <@mesh0, [[1]]> {as_result = false} : tensor<4x8xf32>
+ %2 = mesh.shard %0 to <@mesh0, [[2]]> {as_result = false} : tensor<4x8xf32>
+ return %1, %2 : tensor<4x8xf32>, tensor<4x8xf32>
+}
More information about the Mlir-commits
mailing list