[Mlir-commits] [mlir] [Mesh] initialize mesh dialect (PR #68007)
Chengji Yao
llvmlistbot at llvm.org
Tue Oct 3 08:29:56 PDT 2023
https://github.com/yaochengji updated https://github.com/llvm/llvm-project/pull/68007
>From 655ca9e9775fdded3c0e479a6bf6d3b82f24c5d2 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 1/3] [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>
+}
>From ded65727a402f7a148b1be9aa0801c25a879f572 Mon Sep 17 00:00:00 2001
From: Chengji Yao <yaochengji at hotmial.com>
Date: Mon, 2 Oct 2023 17:55:54 +0000
Subject: [PATCH 2/3] small fix
---
mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td | 2 +-
mlir/lib/Dialect/Mesh/CMakeLists.txt | 16 +---------------
mlir/lib/Dialect/Mesh/IR/CMakeLists.txt | 15 +++++++++++++++
mlir/lib/Dialect/Mesh/{ => IR}/MeshOps.cpp | 0
4 files changed, 17 insertions(+), 16 deletions(-)
create mode 100644 mlir/lib/Dialect/Mesh/IR/CMakeLists.txt
rename mlir/lib/Dialect/Mesh/{ => IR}/MeshOps.cpp (100%)
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td b/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
index 05b5eb75c66d3f2..242d526e42c02a5 100644
--- a/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
+++ b/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
@@ -1,4 +1,4 @@
-//===- MeshBae.td - Mesh Dialect Dialect -------------------*- tablegen -*-===//
+//===- MeshBase.td - Mesh 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.
diff --git a/mlir/lib/Dialect/Mesh/CMakeLists.txt b/mlir/lib/Dialect/Mesh/CMakeLists.txt
index fc3f7c96f05c77f..f33061b2d87cffc 100644
--- a/mlir/lib/Dialect/Mesh/CMakeLists.txt
+++ b/mlir/lib/Dialect/Mesh/CMakeLists.txt
@@ -1,15 +1 @@
-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
-)
+add_subdirectory(IR)
diff --git a/mlir/lib/Dialect/Mesh/IR/CMakeLists.txt b/mlir/lib/Dialect/Mesh/IR/CMakeLists.txt
new file mode 100644
index 000000000000000..fc3f7c96f05c77f
--- /dev/null
+++ b/mlir/lib/Dialect/Mesh/IR/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/IR/MeshOps.cpp
similarity index 100%
rename from mlir/lib/Dialect/Mesh/MeshOps.cpp
rename to mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
>From ee0179330df02f33310c441d3c40dae615bcc322 Mon Sep 17 00:00:00 2001
From: Chengji Yao <yaochengji at hotmial.com>
Date: Tue, 3 Oct 2023 15:23:31 +0000
Subject: [PATCH 3/3] fix comments
---
.../mlir/Dialect/Mesh/IR/CMakeLists.txt | 5 ++
mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td | 54 +++++++++---
mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h | 12 +--
mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td | 88 ++++++++++++++-----
mlir/lib/Dialect/Mesh/IR/CMakeLists.txt | 1 +
mlir/lib/Dialect/Mesh/IR/MeshOps.cpp | 49 ++++-------
mlir/test/Dialect/Mesh/invalid.mlir | 44 +++-------
mlir/test/Dialect/Mesh/ops.mlir | 54 +++++++++++-
8 files changed, 196 insertions(+), 111 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt
index f1d15891d029e1a..cfc948e305638fa 100644
--- a/mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt
@@ -6,3 +6,8 @@ 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)
+
+set(LLVM_TARGET_DEFINITIONS MeshBase.td)
+mlir_tablegen(MeshOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(MeshOpsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRMeshOpsEnumsIncGen)
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td b/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
index 242d526e42c02a5..4a5feacd110a80e 100644
--- a/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
+++ b/mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
@@ -12,6 +12,7 @@
include "mlir/IR/OpBase.td"
include "mlir/IR/AttrTypeBase.td"
include "mlir/IR/BuiltinTypeInterfaces.td"
+include "mlir/IR/EnumAttr.td"
//===----------------------------------------------------------------------===//
// Mesh Dialect
@@ -28,12 +29,25 @@ def Mesh_Dialect : Dialect {
}];
let dependentDialects = [
- "arith::ArithDialect"
+ "arith::ArithDialect" // For materializeConstant()
];
let useDefaultAttributePrinterParser = 1;
let hasConstantMaterializer = 1;
}
+//===----------------------------------------------------------------------===//
+// Mesh Enums.
+//===----------------------------------------------------------------------===//
+
+def Mesh_Partial : I32EnumAttr<"Partial", "partial type of a distributed tensor", [
+ I32EnumAttrCase<"Sum", 0, "partial_sum">,
+ I32EnumAttrCase<"Max", 1, "partial_max">,
+ I32EnumAttrCase<"Min", 2, "partial_min">,
+ I32EnumAttrCase<"Generic", 100, "partial_generic">
+]> {
+ let genSpecializedAttr = 0;
+ let cppNamespace = "::mlir::mesh";
+}
//===----------------------------------------------------------------------===//
// Mesh Attribute
@@ -43,25 +57,39 @@ def MeshSharding : AttrDef<Mesh_Dialect, "MeshSharding"> {
let mnemonic = "shard";
let parameters = (ins
- OptionalParameter<"::mlir::SymbolRefAttr">:$cluster,
- ArrayRefParameter<"::mlir::DenseI64ArrayAttr">:$axes
+ AttrParameter<"::mlir::SymbolRefAttr", "cluster placed">:$cluster,
+ ArrayRefParameter<"::mlir::DenseI8ArrayAttr">:$axes,
+ DefaultValuedParameter<"::mlir::mesh::Partial",
+ "::mlir::mesh::Partial::Sum">:$partial_type
);
let summary = "Attribute that extends tensor type to distributed tensor type.";
let description = [{
- The mesh.shard attribute contains two attribute in it:
+ The MeshSharding attribute could be used in the encoding of a
+ `RankedTensorType` or the mesh.shard op. it contains three sub-attributes:
1. `cluster`: this attribute is a SymbolRefAttr that refers to the mesh
- cluster where the distributed tensor is placed.
+ cluster where the distributed tensor is placed. The symbol must resolve to a
+ `mesh.cluster` operation.
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.
+ - When i == `rank`, it signifies that the tensor is partial one along the x
+ and y axes. An all-reduce should be applied to get the whole tensor, and the
+ reduction type is specified by `partial_type`.
+
+ 3. `partial_type`: indicates the reduction type of the possible all-reduce
+ op. It has 4 possible values:
+ - `partial_sum`: denotes it's an all-reduce-sum
+ - `partial_max`: denotes it's an all-reduce-max
+ - `partial_min`: denotes it's an all-reduce-min
+ - `partial_generic`: denotes that the all-reduce type is complex and cannot
+ be represented merely by a simple sum, max, or min. The exact reduction
+ computation may be derived from the semantics of the corresponding operation
+ or from the reduction computation IR
Example:
@@ -77,14 +105,18 @@ def MeshSharding : AttrDef<Mesh_Dialect, "MeshSharding"> {
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.
+ // it is also a partial_sum along mesh axis 1.
tensor<4x8xf32, #mesh.shard<@mesh0, [[0], [], [1]]>
- // Could also be used in the attribute of mesh.shard op
+ // The tensor is sharded on the first dimension along axis 0 of @mesh0 and
+ // it is also a partial_max along mesh axis 1.
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[0], [], [1]], partial_max>
+
+ // Could be used in the attribute of mesh.shard op
%0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
```
}];
- let assemblyFormat = "`<` ($cluster^ `,`)? `[` $axes `]` `>`";
+ let assemblyFormat = "`<` $cluster `,` `[` $axes `]` (`,` $partial_type^)?`>`";
let genVerifyDecl = 1;
}
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h
index 69159fc0d8cca89..9dfeca84d012165 100644
--- a/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h
+++ b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h
@@ -14,18 +14,10 @@
#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"
+#include "mlir/Dialect/Mesh/IR/MeshOpsEnums.h.inc"
+
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/Mesh/IR/MeshOpsAttributes.h.inc"
diff --git a/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td
index bbd4891a2947517..3acb7cabdc634f6 100644
--- a/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td
+++ b/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td
@@ -40,9 +40,9 @@ def Mesh_ClusterOp : Mesh_Op<"cluster", [Symbol]> {
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 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:
@@ -63,22 +63,20 @@ def Mesh_ClusterOp : Mesh_Op<"cluster", [Symbol]> {
// 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>
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>>
```
}];
let arguments = (ins
SymbolNameAttr:$sym_name,
- I64Attr:$rank,
+ I8Attr:$rank,
DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$dim_sizes
);
- let assemblyFormat = "$sym_name `(` `rank` `=` $rank (`,` `dim_sizes` `=` $dim_sizes^)? `)` attr-dict";
+ let assemblyFormat = [{
+ $sym_name `(` `rank` `=` $rank (`,` `dim_sizes` `=` $dim_sizes^)? `)`
+ attr-dict
+ }];
let hasVerifier = 1;
}
@@ -95,38 +93,82 @@ def Mesh_ShardOp : Mesh_Op<"shard", [Pure, SameOperandsAndResultType]> {
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.
+ 3. `annotate_for_users`: A unit 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 specified, the sharding pertains to specific
+ users of the tensor value, indicating how it should be considered when used
+ as an operand in subsequent operations. If not, the sharding applies to the
+ operation that defines the tensor value.
Example:
```
+ func.func @only_result_annotated(%arg0 : tensor<4x8xf32>) -> () {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
+ ...
+ }
+
+ func.func @only_operand_annotated(%arg0 : tensor<4x8xf32>) -> () {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> annotate_for_users : tensor<4x8xf32>
+ ...
+ }
+
// 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>) -> () {
+ func.func @both_result_and_multi_operands_annotated(
+ %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>
+ %1 = mesh.shard %0 to <@mesh0, [[1]]> annotate_for_users : tensor<4x8xf32>
+ %2 = mesh.shard %0 to <@mesh0, [[2]]> annotate_for_users : tensor<4x8xf32>
"op0"(%1) : ...
"op1"(%2) : ...
...
}
```
+
+ The following usages are undefined:
+ ```
+ func.func @annotate_on_same_result_with_different_sharding(
+ %arg0 : tensor<4x8xf32>) -> () {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
+ %1 = mesh.shard %0 to <@mesh0, [[1]]> : tensor<4x8xf32>
+ ...
+ }
+
+ func.func @annotate_on_same_result_same_value_with_different_sharding(
+ %arg0 : tensor<4x8xf32>) -> () {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
+ %1 = mesh.shard %arg0 to <@mesh0, [[1]]> : tensor<4x8xf32>
+ ...
+ }
+
+ func.func @annotate_on_same_operand_with_different_sharding(
+ %arg0 : tensor<4x8xf32>) -> () {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> annotate_for_users : tensor<4x8xf32>
+ %1 = mesh.shard %0 to <@mesh0, [[1]]> annotate_for_users : tensor<4x8xf32>
+ ...
+ }
+
+ func.func @result_annotated_after_operand(
+ %arg0 : tensor<4x8xf32>) -> () {
+ %0 = mesh.shard %arg0 to <@mesh0, [[0]]> annotate_for_users : tensor<4x8xf32>
+ %1 = mesh.shard %0 to <@mesh0, [[1]]> : tensor<4x8xf32>
+ ...
+ }
+ ```
}];
let arguments = (ins
Builtin_RankedTensor:$src,
MeshSharding:$shard,
- DefaultValuedAttr<BoolAttr, "true">:$as_result
+ UnitAttr:$annotate_for_users
);
let results = (outs
Builtin_RankedTensor:$result
);
- let assemblyFormat = "$src `to` $shard attr-dict `:` type($result)";
- let hasVerifier = 1;
+ let assemblyFormat = [{
+ $src `to` $shard (`annotate_for_users` $annotate_for_users^)? attr-dict `:`
+ type($result)
+ }];
}
#endif // MLIR_DIALECT_MESH_IR_MESHOPS_TD
diff --git a/mlir/lib/Dialect/Mesh/IR/CMakeLists.txt b/mlir/lib/Dialect/Mesh/IR/CMakeLists.txt
index fc3f7c96f05c77f..700e6e21f36b677 100644
--- a/mlir/lib/Dialect/Mesh/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Mesh/IR/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_dialect_library(MLIRMeshDialect
DEPENDS
MLIRMeshOpsAttrIncGen
+ MLIRMeshOpsEnumsIncGen
MLIRMeshOpsIncGen
LINK_LIBS PUBLIC
diff --git a/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp b/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
index fd9f332e204de45..5781bc8130aef8a 100644
--- a/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
+++ b/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
@@ -10,6 +10,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/TypeSwitch.h"
#define DEBUG_TYPE "mesh-ops"
@@ -46,7 +47,7 @@ Operation *MeshDialect::materializeConstant(OpBuilder &builder, Attribute value,
LogicalResult ClusterOp::verify() {
ArrayRef<int64_t> dimSizes = getDimSizes();
- size_t rank = getRank();
+ uint8_t rank = getRank();
if (rank == 0)
return emitOpError("rank of cluster is expected to be a positive integer");
@@ -55,35 +56,10 @@ LogicalResult ClusterOp::verify() {
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)
+ for (int64_t dimSize : dimSizes) {
+ if (dimSize < 0)
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");
+ "dimension size of a mesh cluster is expected to be non-negative");
}
return success();
@@ -95,14 +71,17 @@ LogicalResult ShardOp::verify() {
LogicalResult
MeshShardingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
- SymbolRefAttr, ArrayRef<DenseI64ArrayAttr> axes) {
+ SymbolRefAttr, ArrayRef<DenseI8ArrayAttr> axes,
+ Partial) {
// 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) {
+ llvm::SmallSet<int8_t, 4> visitedAxes;
+ for (DenseI8ArrayAttr subAxes : axes) {
+ ArrayRef<int8_t> subAxesArray = subAxes.asArrayRef();
+ for (int8_t axis : subAxesArray) {
+ if (axis < 0)
+ return emitError() << "mesh axis is expected to be non-negative";
if (!visitedAxes.insert(axis).second)
return emitError() << "mesh axis duplicated";
}
@@ -120,3 +99,5 @@ MeshShardingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/Mesh/IR/MeshOpsAttributes.cpp.inc"
+
+#include "mlir/Dialect/Mesh/IR/MeshOpsEnums.cpp.inc"
diff --git a/mlir/test/Dialect/Mesh/invalid.mlir b/mlir/test/Dialect/Mesh/invalid.mlir
index 8898ab000e6b6ed..39a4645176459f9 100644
--- a/mlir/test/Dialect/Mesh/invalid.mlir
+++ b/mlir/test/Dialect/Mesh/invalid.mlir
@@ -10,37 +10,8 @@ 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>
-}
+// expected-error at +1 {{dimension size of a mesh cluster is expected to be non-negative}}
+mesh.cluster @mesh0(rank = 2, dim_sizes = [-1])
// -----
@@ -63,3 +34,14 @@ func.func @mesh_axis_duplicated_2(
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_negtive(
+ // expected-error at +1 {{mesh axis is expected to be non-negative}}
+ %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[-1]]>>) ->
+ tensor<4x8xf32, #mesh.shard<@mesh0, [[-1]]>> {
+ return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[-1]]>>
+}
diff --git a/mlir/test/Dialect/Mesh/ops.mlir b/mlir/test/Dialect/Mesh/ops.mlir
index c87f07257d66096..a2c18347a1313f2 100644
--- a/mlir/test/Dialect/Mesh/ops.mlir
+++ b/mlir/test/Dialect/Mesh/ops.mlir
@@ -14,6 +14,7 @@ mesh.cluster @mesh3(rank = 2)
// CHECK-LABEL: func @mesh_shard_encoding_fully_replicated
func.func @mesh_shard_encoding_fully_replicated(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32, #mesh.shard<@mesh0, {{\[\[}}]]>>
%arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[]]>>) ->
tensor<4x8xf32, #mesh.shard<@mesh0, [[]]>> {
return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[]]>>
@@ -21,6 +22,7 @@ func.func @mesh_shard_encoding_fully_replicated(
// CHECK-LABEL: func @mesh_shard_encoding_1st_dim
func.func @mesh_shard_encoding_1st_dim(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32, #mesh.shard<@mesh0, {{\[\[}}0]]>>
%arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>>) ->
tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>> {
return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh0, [[0]]>>
@@ -28,6 +30,7 @@ func.func @mesh_shard_encoding_1st_dim(
// CHECK-LABEL: func @mesh_shard_encoding_2nd_dim
func.func @mesh_shard_encoding_2nd_dim(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32, #mesh.shard<@mesh1, {{\[\[}}], [0]]>>
%arg0 : tensor<4x8xf32, #mesh.shard<@mesh1, [[], [0]]>>) ->
tensor<4x8xf32, #mesh.shard<@mesh1, [[], [0]]>> {
return %arg0 : tensor<4x8xf32, #mesh.shard<@mesh1, [[], [0]]>>
@@ -35,35 +38,82 @@ func.func @mesh_shard_encoding_2nd_dim(
// CHECK-LABEL: func @mesh_shard_encoding_1st_and_3rd_dim
func.func @mesh_shard_encoding_1st_and_3rd_dim(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8x16xf32, #mesh.shard<@mesh3, {{\[\[}}0], [], [1]]>>
%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
+// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>
func.func @mesh_shard_op_1st_dim(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ // CHECK-NEXT: mesh.shard %[[ARG]] to <@mesh0, {{\[\[}}0]]> : tensor<4x8xf32>
%0 = mesh.shard %arg0 to <@mesh0, [[0]]> : tensor<4x8xf32>
return %0 : tensor<4x8xf32>
}
// CHECK-LABEL: func @mesh_shard_op_2nd_dim
+// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>
func.func @mesh_shard_op_2nd_dim(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ // CHECK-NEXT: mesh.shard %[[ARG]] to <@mesh1, {{\[\[}}], [0]]> : 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(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8x16xf32>
%arg0 : tensor<4x8x16xf32>) -> tensor<4x8x16xf32> {
+ // CHECK-NEXT: mesh.shard %[[ARG]] to <@mesh3, {{\[\[}}0], [], [1]]> : tensor<4x8x16xf32>
%0 = mesh.shard %arg0 to <@mesh3, [[0], [], [1]]> : tensor<4x8x16xf32>
return %0 : tensor<4x8x16xf32>
}
+// CHECK-LABEL: func @mesh_shard_op_partial_max
+func.func @mesh_shard_op_partial_max(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>
+ %arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ // CHECK-NEXT: mesh.shard %[[ARG]] to <@mesh3, {{\[\[}}0], [], [1]], partial_max> : tensor<4x8xf32>
+ %0 = mesh.shard %arg0 to <@mesh3, [[0], [], [1]], partial_max> : tensor<4x8xf32>
+ return %0 : tensor<4x8xf32>
+}
+
+// CHECK-LABEL: func @mesh_shard_op_partial_min
+func.func @mesh_shard_op_partial_min(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>
+ %arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ // CHECK-NEXT: mesh.shard %[[ARG]] to <@mesh3, {{\[\[}}0], [], [1]], partial_min> : tensor<4x8xf32>
+ %0 = mesh.shard %arg0 to <@mesh3, [[0], [], [1]], partial_min> : tensor<4x8xf32>
+ return %0 : tensor<4x8xf32>
+}
+
+// CHECK-LABEL: func @mesh_shard_op_partial_generic
+func.func @mesh_shard_op_partial_generic(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>
+ %arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ // CHECK-NEXT: mesh.shard %[[ARG]] to <@mesh3, {{\[\[}}0], [], [1]], partial_generic> : tensor<4x8xf32>
+ %0 = mesh.shard %arg0 to <@mesh3, [[0], [], [1]], partial_generic> : tensor<4x8xf32>
+ return %0 : tensor<4x8xf32>
+}
+
+// CHECK-LABEL: func @mesh_shard_op_partial_sum_explict
+func.func @mesh_shard_op_partial_sum_explict(
+ // CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>
+ %arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {
+ // CHECK-NEXT: mesh.shard %[[ARG]] to <@mesh3, {{\[\[}}0], [], [1]]> : tensor<4x8xf32>
+ %0 = mesh.shard %arg0 to <@mesh3, [[0], [], [1]], partial_sum> : tensor<4x8xf32>
+ return %0 : tensor<4x8xf32>
+}
+
// CHECK-LABEL: func @mesh_shard_op_two_users
+// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>
func.func @mesh_shard_op_two_users(%arg0 : tensor<4x8xf32>) ->
(tensor<4x8xf32>, tensor<4x8xf32>) {
+ // CHECK-NEXT: %[[V0:.*]] = mesh.shard %[[ARG]] to <@mesh0, {{\[\[}}0]]> : 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>
+ // CHECK-DAG: mesh.shard %[[V0]] to <@mesh0, {{\[\[}}1]]> annotate_for_users : tensor<4x8xf32>
+ %1 = mesh.shard %0 to <@mesh0, [[1]]> annotate_for_users : tensor<4x8xf32>
+ // CHECK-DAG: mesh.shard %[[V0]] to <@mesh0, {{\[\[}}2]]> annotate_for_users : tensor<4x8xf32>
+ %2 = mesh.shard %0 to <@mesh0, [[2]]> annotate_for_users : tensor<4x8xf32>
return %1, %2 : tensor<4x8xf32>, tensor<4x8xf32>
}
More information about the Mlir-commits
mailing list