[Mlir-commits] [mlir] [Mesh] initialize mesh dialect (PR #68007)
Mehdi Amini
llvmlistbot at llvm.org
Mon Oct 2 11:41:19 PDT 2023
================
@@ -0,0 +1,91 @@
+//===- 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.
+// 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.
----------------
joker-eph wrote:
What about using a separate extra field for the partial instead?
https://github.com/llvm/llvm-project/pull/68007
More information about the Mlir-commits
mailing list