From 5c8cf91ca3dc97a0e1a2d1871bb7bcd19f41482c Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 01/11] Mesh Logical Axis Ops --- src/enzyme_ad/jax/BUILD | 2 + .../jax/Dialect/Distributed/Dialect.h | 2 + src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp | 48 +++++++++++++++++++ src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 48 +++++++++++++++++++ .../jax/Dialect/Distributed/Types.td | 8 ++++ 5 files changed, 108 insertions(+) diff --git a/src/enzyme_ad/jax/BUILD b/src/enzyme_ad/jax/BUILD index 56fab16aa5..98dcda1986 100644 --- a/src/enzyme_ad/jax/BUILD +++ b/src/enzyme_ad/jax/BUILD @@ -587,6 +587,7 @@ td_library( ], deps = [ "@llvm-project//mlir:BuiltinDialectTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", "@llvm-project//mlir:OpBaseTdFiles", ], ) @@ -632,6 +633,7 @@ gentbl_cc_library( td_file = "Dialect/Distributed/Ops.td", deps = [ ":DistributedDialectFiles", + "@llvm-project//mlir:InferTypeOpInterface", ], ) diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h index e8277364c6..009400ea57 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h @@ -9,6 +9,8 @@ #include "mlir/IR/Region.h" #include "mlir/IR/SymbolTable.h" #include "mlir/IR/Types.h" +#include "mlir/Interfaces/InferTypeOpInterface.h" + // Include the dialect #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedDialect.h.inc" diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp index 21d62cecf6..c4ce290d95 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp @@ -1,4 +1,5 @@ #include "mlir/IR/Builders.h" +#include "mlir/Interfaces/InferTypeOpInterface.h" #include "llvm/ADT/TypeSwitch.h" #include "Dialect.h" @@ -125,6 +126,53 @@ SplitBranchOp::verifySymbolUses(::mlir::SymbolTableCollection &symbol_table) { return mlir::success(); } +LogicalResult AxisFactorOp::verify() { + auto factors = getFactors(); + for (auto factor_attr : factors) { + if (auto factor = dyn_cast(factor_attr)) { + if (factor.getValue().getSExtValue() <= 0) { + return emitOpError() << "requires all factors to be > 0"; + } + } else { + return emitOpError() << "requires all factors to be integer attributes"; + } + } + + if (getLogicalAxes().size() != factors.size()) { + return emitOpError() + << "requires one logical axis result per factor (got " + << getLogicalAxes().size() << " results for " << factors.size() + << " factors)"; + } + + // TODO: Verify that product(factors) equals physical_axis size. + return mlir::success(); +} + +LogicalResult AxisFactorOp::inferReturnTypes( + MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, OpaqueProperties properties, + RegionRange regions, SmallVectorImpl &inferredReturnTypes) { + AxisFactorOpAdaptor adaptor(operands, attributes, properties, regions); + if (!adaptor.getFactors()) { + if (location) + mlir::emitError(*location) << "missing factors attribute"; + return mlir::failure(); + } + + inferredReturnTypes.reserve(adaptor.getFactors().size()); + for (auto _ : adaptor.getFactors()) { + inferredReturnTypes.push_back(LogicalCommAxisType::get(context)); + } + return mlir::success(); +} + +LogicalResult +AxisProductOp::verify() { + // TODO disjointness + return mlir::success(); +} + LogicalResult DefineTokenOp::verifySymbolUses(::mlir::SymbolTableCollection &symbol_table) { // Tokens need to indicate which channel they communicate over diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td index 0857ce2f2c..f59cbbcf54 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -3,11 +3,59 @@ include "mlir/IR/BuiltinAttributes.td" include "mlir/IR/SymbolInterfaces.td" +include "mlir/Interfaces/InferTypeOpInterface.td" include "Dialect.td" include "Types.td" include "Interfaces.td" +// Mesh definition ops +def PhysicalMeshOp : DistributedOp<"PhysicalMesh"> { + let description = [{Defines a physical mesh of devices based on the hardware topology. A physical mesh is a Cartesian product + of zero or more physical axes. Axes should be distinct! Two copies of the same axis is permitted (but may better be represented + as a larger axis), but passing the same axis twice is an error.}]; + let arguments = (ins + Variadic:$axes + ); + let results = (outs PhysicalMeshType:$mesh); + let assemblyFormat = "$axes attr-dict"; +} + +def AxisFactorOp : DistributedOp<"AxisFactor", [DeclareOpInterfaceMethods]> { + let description = [{ + Factors a physical axis into multiple logical axes, according to a static array of factor sizes. + The lowest index of factors corresponds to the lowest-significance / most contiguous new axis. + Factors must multiply into the original physical axis size. + Returns a logical axis per factor. + }]; + let arguments = (ins + PhysicalCommAxisType:$physical_axis, + I32ArrayAttr:$factors + ); + let results = (outs Variadic:$logical_axes); + let assemblyFormat = "$physical_axis $factors attr-dict"; + let hasVerifier = 1; // TODO: verify factors multiply to physical axis size, and correct number of outputs. +} + + +def AxisProductOp : DistributedOp<"AxisProduct"> { + let description = [{ + Combines multiple logical axes into a single logical axis. Logical axis must be disjoint, + but may come from any mix of physical axes, including different strides over the same physical + axis. The order encodes the significance, with the lowest-index/first logical axis being + index as the least-significant/most contiguous dimension. + For simplicity for implementation, we require any factors of the same physical axis to be + defined from the same factorization op. (For instance, this makes disjointness easier) + }]; + let arguments = (ins + Variadic:$logical_axes + ); + let results = (outs LogicalCommAxisType:$logical_axis); + let assemblyFormat = "$logical_axes attr-dict"; + let hasVerifier = 1; // TODO: axis disjointness. +} + + // Device definition ops def ChannelOp : DistributedOp<"Channel", [Symbol, ChannelDefTrait, DeclareOpInterfaceMethods]>{ diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Types.td b/src/enzyme_ad/jax/Dialect/Distributed/Types.td index 517c63753c..5730207dd5 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Types.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Types.td @@ -3,6 +3,14 @@ include "Dialect.td" +// A natural axis defined by the hardware topology. +def PhysicalCommAxisType : DistributedType<"PhysicalCommAxis", "physical_comm_axis">; +// Includes physical axis factors and products of other logical axis. +def LogicalCommAxisType : DistributedType<"LogicalCommAxis", "logical_comm_axis">; + +def PhysicalMeshType : DistributedType<"PhysicalMesh", "physical_mesh">; +def LogicalMeshType : DistributedType<"LogicalMesh", "logical_mesh">; + def TokenType : DistributedType<"Token", "token">; #endif // ENZYME_DISTRIBUTED_DIALECT_TYPES_H \ No newline at end of file From e787ec321a63a9d13e572fa743b2727266c9e601 Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 02/11] refactor dialect: axes --- src/enzyme_ad/jax/BUILD | 58 ++++ .../jax/Dialect/Distributed/AxisOps.cpp | 136 ++++++++ .../jax/Dialect/Distributed/Dialect.cpp | 4 + .../jax/Dialect/Distributed/Dialect.h | 27 +- .../jax/Dialect/Distributed/Interfaces.td | 95 +++++- .../Distributed/MeshComputationOps.cpp | 317 ++++++++++++++++++ src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp | 185 ---------- src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 250 ++++++++------ .../jax/Dialect/Distributed/Traits.h | 8 +- .../jax/Dialect/Distributed/Types.cpp | 41 +++ .../jax/Dialect/Distributed/Types.td | 64 +++- .../jax/Dialect/Distributed/Utilities.cpp | 161 +++++++++ .../jax/Dialect/Distributed/Utilities.h | 45 +++ .../jax/Passes/Distributed/Passes.cpp | 7 + src/enzyme_ad/jax/Passes/Distributed/Passes.h | 16 + .../jax/Passes/Distributed/Passes.td | 10 + src/enzyme_ad/jax/RegistryUtils.cpp | 2 + .../distributed/abcda_batch_matmul.mlir | 43 +++ .../branched_pipeline_friendly.mlir | 33 ++ test/lit_tests/distributed/concat.mlir | 32 ++ .../diamond_pipeline_friendly.mlir | 31 ++ test/lit_tests/distributed/dotadd.mlir | 23 ++ .../innerouterinner_dataparallel.mlir | 24 ++ .../innerouterinner_operatorparallel.mlir | 23 ++ .../innerouterinner_pipeline_friendly.mlir | 19 ++ test/lit_tests/distributed/roundtrip.mlir | 42 --- test/lit_tests/distributed/test_timing.mlir | 47 +++ test/lit_tests/distributed/two_matmul.mlir | 18 + 28 files changed, 1417 insertions(+), 344 deletions(-) create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp delete mode 100644 src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/Types.cpp create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/Utilities.h create mode 100644 src/enzyme_ad/jax/Passes/Distributed/Passes.cpp create mode 100644 src/enzyme_ad/jax/Passes/Distributed/Passes.h create mode 100644 src/enzyme_ad/jax/Passes/Distributed/Passes.td create mode 100644 test/lit_tests/distributed/abcda_batch_matmul.mlir create mode 100644 test/lit_tests/distributed/branched_pipeline_friendly.mlir create mode 100644 test/lit_tests/distributed/concat.mlir create mode 100644 test/lit_tests/distributed/diamond_pipeline_friendly.mlir create mode 100644 test/lit_tests/distributed/dotadd.mlir create mode 100644 test/lit_tests/distributed/innerouterinner_dataparallel.mlir create mode 100644 test/lit_tests/distributed/innerouterinner_operatorparallel.mlir create mode 100644 test/lit_tests/distributed/innerouterinner_pipeline_friendly.mlir delete mode 100644 test/lit_tests/distributed/roundtrip.mlir create mode 100644 test/lit_tests/distributed/test_timing.mlir create mode 100644 test/lit_tests/distributed/two_matmul.mlir diff --git a/src/enzyme_ad/jax/BUILD b/src/enzyme_ad/jax/BUILD index 98dcda1986..a90ad8ad90 100644 --- a/src/enzyme_ad/jax/BUILD +++ b/src/enzyme_ad/jax/BUILD @@ -585,10 +585,16 @@ td_library( "Dialect/Distributed/Ops.td", "Dialect/Distributed/Types.td", ], + includes = [ + ".", + "../../../external/shardy", + ], deps = [ "@llvm-project//mlir:BuiltinDialectTdFiles", "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", "@llvm-project//mlir:OpBaseTdFiles", + "@stablehlo//:stablehlo_ops_td_files", + "@shardy//shardy/dialect/sdy/ir:sdy_td_files", ], ) @@ -634,6 +640,7 @@ gentbl_cc_library( deps = [ ":DistributedDialectFiles", "@llvm-project//mlir:InferTypeOpInterface", + "@shardy//shardy/dialect/sdy/ir:sdy_td_files", ], ) @@ -675,6 +682,25 @@ gentbl_cc_library( ], ) +gentbl_cc_library( + name = "DistributedTypeInterfacesIncGen", + tbl_outs = [ + ( + ["--gen-type-interface-decls"], + "Dialect/Distributed/DistributedTypeInterfaces.h.inc", + ), + ( + ["--gen-type-interface-defs"], + "Dialect/Distributed/DistributedTypeInterfaces.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/Distributed/Interfaces.td", + deps = [ + ":DistributedDialectFiles", + ], +) + td_library( name = "TesseraDialectTdFiles", srcs = [ @@ -826,6 +852,33 @@ gentbl_cc_library( ], ) +td_library( + name = "DistributedPassesTdFiles", + srcs = [], + deps = [ + "@llvm-project//mlir:PassBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "DistributedPassesIncGen", + tbl_outs = [ + ( + [ + "-gen-pass-decls", + "-name=distributed", + ], + "Passes/Distributed/Passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Passes/Distributed/Passes.td", + deps = [ + ":DistributedPassesTdFiles", + "@llvm-project//mlir:FunctionInterfacesTdFiles", + ], +) + cc_library( name = "CheckedRewrite", hdrs = ["CheckedRewrite.h"], @@ -1197,6 +1250,7 @@ cc_library( "Analysis/*.cpp", "Implementations/*.cpp", "Passes/*.cpp", + "Passes/Distributed/*.cpp", "Passes/Tessera/*.cpp", "Dialect/*.cpp", "Dialect/Distributed/*.cpp", @@ -1215,6 +1269,7 @@ cc_library( "Analysis/*.h", "Implementations/*.h", "Passes/*.h", + "Passes/Distributed/*.h", "Passes/Tessera/*.h", "Dialect/*.h", "Dialect/Distributed/*.h", @@ -1236,7 +1291,10 @@ cc_library( deps = [ ":CheckedRewrite", ":DistributedDialectIncGen", + ":DistributedInterfacesIncGen", + ":DistributedPassesIncGen", ":DistributedOpsIncGen", + ":DistributedTypeInterfacesIncGen", ":DistributedTypesIncGen", ":EnzymeHLOPatternsIncGen", ":EnzymeXLAAttrsIncGen", diff --git a/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp new file mode 100644 index 0000000000..9d491da843 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp @@ -0,0 +1,136 @@ +#include "Dialect.h" +#include "Utilities.h" + +namespace mlir::enzyme::distributed { + +namespace { + +static LogicalResult verifyFactorExtents(ArrayRef extents, + unsigned sourceExtent, + Operation *op) { + uint64_t product = 1; + for (int32_t extent : extents) { + if (extent <= 0) { + return op->emitOpError() << "requires all factor extents to be > 0"; + } + product *= static_cast(extent); + } + if (product != sourceExtent) { + return op->emitOpError() << "requires product(factor_extents) == axis " + "extent (" + << product << " != " << sourceExtent << ")"; + } + return success(); +} + +static void computeMajorToMinorStrides(ArrayRef extents, + SmallVectorImpl &strides) { + // Leftmost factors are most major. + strides.resize(extents.size()); + unsigned running = 1; + for (int i = static_cast(extents.size()) - 1; i >= 0; --i) { + strides[i] = running; + running *= static_cast(extents[i]); + } +} + +} // namespace + +LogicalResult GetPhysicalAxisOp::verify() { + FailureOr physicalAxis = + resolveSymbolOpFromAttr(*this, + getPhysicalAxisAttr()); + if (failed(physicalAxis)) { + return emitOpError() << "references unknown physical axis symbol " + << getPhysicalAxisAttr(); + } + + auto axisType = dyn_cast(getAxis().getType()); + if (!axisType) { + return emitOpError() << "requires result type to be PhysicalCommAxisType"; + } + + unsigned expectedExtent = + static_cast(physicalAxis->getPhysicalAxisSize()); + if (axisType.getExtent() != expectedExtent) { + return emitOpError() << "requires result type extent to match referenced " + "physical axis size (" + << axisType.getExtent() << " != " << expectedExtent + << ")"; + } + + return success(); +} + +LogicalResult GenericAxisOp::verify() { + if (!isa(getAxis().getType())) { + return emitOpError() << "requires result type to implement AxisTypeInterface"; + } + + return success(); +} + +LogicalResult AxisFactorOp::verify() { + auto axisIface = dyn_cast(getAxis().getType()); + if (!axisIface) { + return emitOpError() << "requires axis operand type to implement " + "AxisTypeInterface"; + } + + ArrayRef factorExtents = getFactorExtents(); + + if (failed(verifyFactorExtents(factorExtents, axisIface.extent(), + getOperation()))) { + return failure(); + } + + if (getAxisFactors().size() != factorExtents.size()) { + return emitOpError() << "requires number of results to match number of " + "factor extents"; + } + + SmallVector expectedStrides; + computeMajorToMinorStrides(factorExtents, expectedStrides); + + for (auto [idx, axisFactorVal] : llvm::enumerate(getAxisFactors())) { + auto axisFactorType = dyn_cast(axisFactorVal.getType()); + if (!axisFactorType) { + return emitOpError() << "requires all results to have AxisFactorType"; + } + if (axisFactorType.getAxisType() != getAxis().getType()) { + return emitOpError() << "requires result #" << idx + << " to have base axis type equal to operand type"; + } + if (axisFactorType.getExtent() != static_cast(factorExtents[idx])) { + return emitOpError() << "requires result #" << idx + << " extent to match factor extent"; + } + if (axisFactorType.getStride() != expectedStrides[idx]) { + return emitOpError() << "requires result #" << idx + << " stride to follow leftmost-major convention"; + } + } + + return success(); +} + +LogicalResult AxisFactorOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + AxisFactorOpAdaptor adaptor(operands, attributes, properties, regions); + ArrayRef factorExtents = adaptor.getFactorExtents(); + + SmallVector strides; + computeMajorToMinorStrides(factorExtents, strides); + + inferredReturnTypes.reserve(factorExtents.size()); + Type axisType = adaptor.getAxis().getType(); + for (auto [idx, factorExtent] : llvm::enumerate(factorExtents)) { + inferredReturnTypes.push_back(AxisFactorType::get( + context, axisType, static_cast(factorExtent), strides[idx])); + } + return success(); +} + +} // namespace mlir::enzyme::distributed diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp index 343d327678..2715e256b3 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp @@ -8,6 +8,10 @@ #define GET_TYPEDEF_CLASSES #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedTypes.cpp.inc" +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedTypeInterfaces.cpp.inc" + +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedInterfaces.cpp.inc" + // Initialize the dialect void mlir::enzyme::distributed::DistributedDialect::initialize() { addTypes< diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h index 009400ea57..14f6f957fe 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h @@ -2,6 +2,7 @@ #define ENZYME_AD_JAX_DIALECT_DISTRIBUTED_DIALECT_H #include "mlir/Bytecode/BytecodeOpInterface.h" +#include "mlir/Dialect/Quant/IR/QuantTypes.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Dialect.h" #include "mlir/IR/DialectImplementation.h" @@ -10,15 +11,39 @@ #include "mlir/IR/SymbolTable.h" #include "mlir/IR/Types.h" #include "mlir/Interfaces/InferTypeOpInterface.h" +#include "mlir/Support/LLVM.h" +#include "stablehlo/dialect/Base.h" +#include "shardy/dialect/sdy/ir/dialect.h" + +// Need to declare the TypedOpResult struct before including the generated +// interface declarations +namespace mlir::enzyme::distributed { +template struct TypedOpResult { + ::mlir::OpResult value; + TypedOpResult(mlir::Value value) : value(llvm::cast(value)) { + assert(isa(value.getType()) && "value must have the correct type"); + } + TypedOpResult(::mlir::OpResult value) : value(value) { + assert(isa(value.getType()) && "value must have the correct type"); + } + operator mlir::OpResult() const { return value; } + + mlir::OpResult asOpResult() const { return value; } +}; +} // namespace mlir::enzyme::distributed // Include the dialect #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedDialect.h.inc" -// Traits and interfaces +// Traits #include "Traits.h" +// Type interfaces +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedTypeInterfaces.h.inc" // Types #define GET_TYPEDEF_CLASSES #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedTypes.h.inc" +// Interfaces +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedInterfaces.h.inc" // Operations #define GET_OP_CLASSES #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedOps.h.inc" diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td b/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td index 8fd74fd4e2..1fe5f334ba 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td @@ -3,7 +3,98 @@ include "mlir/IR/OpBase.td" -def DeviceDefTrait : NativeOpTrait<"enzyme::distributed::DeviceDefTrait">; -def ChannelDefTrait : NativeOpTrait<"enzyme::distributed::ChannelDefTrait">; +def AxisTypeInterface : TypeInterface<"AxisTypeInterface"> { + let cppNamespace = "::mlir::enzyme::distributed"; + let description = [{ + Interface for axis types with a statically encoded extent. + }]; + + let methods = [ + InterfaceMethod< + [{Returns the compile-time extent of this axis.}], + "unsigned", "extent", (ins), "", + [{ + return $_type.getExtent(); + }]>, + InterfaceMethod< + [{Returns whether two axis SSA values refer to the same underlying axis under this axis type's semantics. + Axis are allowed to either be fully equivalent or disjoint: canonical axes may not partially alias.}], + "bool", "aliases", + (ins "::mlir::Value":$selfAxis, "::mlir::Value":$otherAxis)> + ]; +} + +def PhysicalCommAxisOpInterface : + OpInterface<"PhysicalCommAxisOpInterface"> { + let description = [{ + Interface for ops that define a physical communication axis. + }]; + + let methods = [ + InterfaceMethod< + [{Returns the extent/size of this physical axis in number of devices.}], + "int64_t", "getPhysicalAxisSize", (ins)> + ]; +} + +// Interfaces for timing analysis +def LaneOpInterface : + OpInterface<"LaneOpInterface"> { + let methods = [ + InterfaceMethod< + [{Returns the regions corresponding to the different lanes of this operation}], + "llvm::SmallVector", "getLanes", (ins) + > + ]; +} + +def HandleDeclarationOpInterface : + OpInterface<"HandleDeclarationOpInterface"> { + let methods = [ + InterfaceMethod< + [{Returns the SSA value representing an event handle}], + "mlir::Value", "getHandle", (ins) + > + ]; +} + +def EventOpInterface : + OpInterface<"EventOpInterface"> { + let description = [{ + An op representing part or all of an event. Not all ops in a lane + need to be events, but any op spanning multiple lanes must be. + Events are ordered by a happens-before relationship, determined by: + - Explicit happensAfter(): a list of handles whose events must complete before this event can start. + - Implicit happens after: No interface needed, this is determined by the order of operations within a + basic block and by SSA dependencies. + - simultaneousWith(): a handle that this event happens simultaneously with. Multiple events + (including those on the same lane) may be simultaneous with the same handle, forming an + equivalence relation. From a happens-before view, these events are merged into the same + event and share predecessors and successors. Simultaneous events must not have a happens-before + relationship between each other, with one exception: concurrent ops in the same lane whose only + dependency is from implicit op order, with no intervening non-simultaneous ops. + - concurrentWith(Op other): A "must be concurrent" check allowing an underapproximation of + concurrency. Currently this is only used when checking the legality of simultaneousWith(). + }]; + let methods = [ + InterfaceMethod< + [{Returns the SSA value handles whose events must complete before this event can start}], + "::llvm::SmallVector<::mlir::Value>", "happensAfter", (ins) + >, + InterfaceMethod< + [{Returns the SSA value handles whose events must happen after this event.}], + "::llvm::SmallVector<::mlir::Value>", "happensBefore", (ins) + >, + InterfaceMethod< + [{Returns the SSA value handles whose events this event happens simultaneously with}], + "::llvm::SmallVector<::mlir::Value>", "simultaneousWith", (ins) + >, + InterfaceMethod< + [{An underapproximating relation allowing concurrency between two ops. }], + "bool", "concurrentWith", (ins "mlir::Operation *":$other) + > + ]; +} + #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_INTERFACES \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp new file mode 100644 index 0000000000..79f91e1124 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp @@ -0,0 +1,317 @@ +#include "Dialect.h" +#include "Utilities.h" + +#include "shardy/dialect/sdy/ir/utils.h" + +namespace mlir::enzyme::distributed { + +llvm::SmallVector MeshComputationOp::getLanes() { + llvm::SmallVector lanes; + lanes.reserve((*this)->getNumRegions()); + for (Region ®ion : (*this)->getRegions()) { + lanes.push_back(®ion); + } + return lanes; +} + +// Resolves a tensor SSA value back to its input-operand slot. +FailureOr MeshComputationOp::findInputTensorIndex(Value tensor) const { + auto inputTensors = const_cast(this)->getInputTensors(); + for (auto [index, inputTensor] : llvm::enumerate(inputTensors)) { + if (inputTensor == tensor) { + return static_cast(index); + } + } + return failure(); +} + +LogicalResult MeshComputationOp::setInputTensorSharding( + Value tensor, sdy::TensorShardingAttr sharding) { + if (failed(findInputTensorIndex(tensor))) { + return emitOpError() + << "expected tensor to be one of the mesh computation input tensors"; + } + sdy::setSharding(tensor, sharding); + return success(); +} + +// Looks up the distributed logical axis corresponding to a recorded Shardy +// axis name. +FailureOr +MeshComputationOp::findLogicalAxisForShardyAxisName(StringRef shardyAxisName) const { + ArrayAttr shardyAxisNames = const_cast(this)->getShardyAxisNames(); + if (!shardyAxisNames) { + return failure(); + } + + llvm::SmallVector logicalAxes; + logicalAxes.append(getSpmdAxes().begin(), getSpmdAxes().end()); + logicalAxes.append(getMpmdAxes().begin(), getMpmdAxes().end()); + if (logicalAxes.size() != shardyAxisNames.size()) { + return failure(); + } + + for (auto [index, nameAttr] : llvm::enumerate(shardyAxisNames)) { + auto axisName = dyn_cast(nameAttr); + if (!axisName) { + return failure(); + } + if (axisName.getValue() == shardyAxisName) { + return logicalAxes[index]; + } + } + + return failure(); +} + +// Looks up the recorded Shardy axis name corresponding to a logical axis SSA +// value. +FailureOr +MeshComputationOp::findShardyAxisNameForLogicalAxis(Value logicalAxis) const { + ArrayAttr shardyAxisNames = const_cast(this)->getShardyAxisNames(); + if (!shardyAxisNames) { + return failure(); + } + + llvm::SmallVector logicalAxes; + logicalAxes.append(getSpmdAxes().begin(), getSpmdAxes().end()); + logicalAxes.append(getMpmdAxes().begin(), getMpmdAxes().end()); + if (logicalAxes.size() != shardyAxisNames.size()) { + return failure(); + } + + for (auto [index, axis] : llvm::enumerate(logicalAxes)) { + if (axis == logicalAxis) { + auto axisName = dyn_cast(shardyAxisNames[index]); + if (!axisName) { + return failure(); + } + return axisName; + } + } + + return failure(); +} + +ValueRange MeshComputationOp::getSpmdAxes() const { + auto segmentSizesAttr = + (*this)->getAttrOfType("operandSegmentSizes"); + if (!segmentSizesAttr || segmentSizesAttr.size() != 3) { + return {}; + } + return ValueRange( + (*this)->getOperands().slice(/*start=*/0, segmentSizesAttr[0])); +} + +ValueRange MeshComputationOp::getMpmdAxes() const { + auto segmentSizesAttr = + (*this)->getAttrOfType("operandSegmentSizes"); + if (!segmentSizesAttr || segmentSizesAttr.size() != 3) { + return {}; + } + return ValueRange( + (*this)->getOperands().slice(/*start=*/segmentSizesAttr[0], + segmentSizesAttr[1])); +} + +uint32_t MeshComputationOp::getNumDeviceBodies() const { + auto attr = (*this)->getAttrOfType("num_device_bodies"); + return attr ? static_cast(attr.getInt()) : 0; +} + +uint32_t MeshComputationOp::getNumCommunicationBodies() const { + auto attr = + (*this)->getAttrOfType("num_communication_bodies"); + return attr ? static_cast(attr.getInt()) : 0; +} + +Region &MeshComputationOp::getDeviceBody(unsigned idx) { + return getOperation()->getRegion(idx); +} + +const Region &MeshComputationOp::getDeviceBody(unsigned idx) const { + return (*this)->getRegion(idx); +} + +// Maps a multi-axis MPMD device coordinate to a flat index in the device-body +// partition. Lower-index (leftmost) MPMD axes are treated as higher +// significance. +FailureOr MeshComputationOp::findComputationBodyIndexByDeviceIndex( + const DeviceIndex &deviceIndex) const { + auto mpmdAxes = getMpmdAxes(); + if (deviceIndex.size() != mpmdAxes.size()) { + return failure(); + } + + unsigned flatIndex = 0; + unsigned stride = 1; + for (size_t axisPos = mpmdAxes.size(); axisPos-- > 0;) { + Value axis = mpmdAxes[axisPos]; + int64_t axisSize = static_cast(getAxisSize(axis)); + if (axisSize <= 0) { + return failure(); + } + + unsigned coordinate = deviceIndex[axisPos]; + if (coordinate >= static_cast(axisSize)) { + return failure(); + } + + flatIndex += coordinate * stride; + stride *= static_cast(axisSize); + } + + uint32_t numDeviceBodies = getNumDeviceBodies(); + if (flatIndex >= static_cast(numDeviceBodies)) { + return failure(); + } + return flatIndex; +} + +Region &MeshComputationOp::getDeviceBodyByDeviceIndex( + const DeviceIndex &deviceIndex) { + auto bodyIndex = findComputationBodyIndexByDeviceIndex(deviceIndex); + assert(succeeded(bodyIndex) && + "expected a valid DeviceIndex in the MPMD submesh"); + return getDeviceBody(*bodyIndex); +} + +const Region &MeshComputationOp::getDeviceBodyByDeviceIndex( + const DeviceIndex &deviceIndex) const { + auto bodyIndex = findComputationBodyIndexByDeviceIndex(deviceIndex); + assert(succeeded(bodyIndex) && + "expected a valid DeviceIndex in the MPMD submesh"); + return getDeviceBody(*bodyIndex); +} + +Region &MeshComputationOp::getCommunicationBody(unsigned idx) { + unsigned communicationBodyStart = static_cast(getNumDeviceBodies()); + return getOperation()->getRegion(communicationBodyStart + idx); +} + +const Region &MeshComputationOp::getCommunicationBody(unsigned idx) const { + unsigned communicationBodyStart = getNumDeviceBodies(); + return (*this)->getRegion(communicationBodyStart + idx); +} + +// Returns the communication-body index for a communication axis in the +// combined communication partition. The ordering is SPMD axes first, then MPMD +// axes, and the returned index is relative to the start of that partition. +FailureOr +MeshComputationOp::findCommunicationBodyIndexForAxis(Value axis) const { + unsigned communicationBodyIndex = 0; + auto spmdAxes = getSpmdAxes(); + auto mpmdAxes = getMpmdAxes(); + for (bool spmdCase : {true, false}) { + auto axisRange = spmdCase ? spmdAxes : mpmdAxes; + for (Value candidate : axisRange) { + if (candidate == axis) { + return communicationBodyIndex; + } + ++communicationBodyIndex; + } + } + return failure(); +} + +Region &MeshComputationOp::getCommunicationBodyForAxis(Value axis) { + auto bodyIndex = findCommunicationBodyIndexForAxis(axis); + assert(succeeded(bodyIndex) && + "communication body lookup expects a logical axis from either the " + "SPMD or MPMD partitions"); + return getCommunicationBody(*bodyIndex); +} + +const Region &MeshComputationOp::getCommunicationBodyForAxis(Value axis) const { + auto bodyIndex = findCommunicationBodyIndexForAxis(axis); + assert(succeeded(bodyIndex) && + "communication body lookup expects a logical axis from either the " + "SPMD or MPMD partitions"); + return getCommunicationBody(*bodyIndex); +} + +// Verifies region partition sizes and axis/body cardinality invariants. +LogicalResult MeshComputationOp::verify() { + auto segmentSizesAttr = + (*this)->getAttrOfType("operandSegmentSizes"); + if (!segmentSizesAttr) { + return emitOpError() << "requires operandSegmentSizes attribute"; + } + if (segmentSizesAttr.size() != 3) { + return emitOpError() + << "requires operandSegmentSizes to have 3 entries " + "(spmd_axes, mpmd_axes, input_tensors)"; + } + + auto numDeviceBodiesAttr = + (*this)->getAttrOfType("num_device_bodies"); + if (!numDeviceBodiesAttr) { + return emitOpError() << "requires num_device_bodies attribute"; + } + auto numCommunicationBodiesAttr = + (*this)->getAttrOfType("num_communication_bodies"); + if (!numCommunicationBodiesAttr) { + return emitOpError() << "requires num_communication_bodies attribute"; + } + + if (numDeviceBodiesAttr.getInt() < 0) { + return emitOpError() << "requires num_device_bodies to be non-negative"; + } + if (numCommunicationBodiesAttr.getInt() < 0) { + return emitOpError() + << "requires num_communication_bodies to be non-negative"; + } + + // Check overall region count + unsigned expectedBodyCount = + static_cast(getNumDeviceBodies()) + + static_cast(getNumCommunicationBodies()); + if (getOperation()->getNumRegions() != expectedBodyCount) { + return emitOpError() + << "requires the number of regions to equal device bodies + " + "communication bodies"; + } + + // Check communication region count matches attributes + size_t expectedCommunicationBodyCount = + getSpmdAxes().size() + getMpmdAxes().size(); + if (expectedCommunicationBodyCount != + static_cast(getNumCommunicationBodies())) { + return emitOpError() + << "requires the number of communication bodies to equal the " + "number of SPMD + MPMD axes"; + } + + // Check device region count matches the product of MPMD axis sizes + int64_t expectedDeviceBodyCount = 1; + for (Value axis : getMpmdAxes()) { + expectedDeviceBodyCount *= static_cast(getAxisSize(axis)); + } + if (expectedDeviceBodyCount != getNumDeviceBodies()) { + return emitOpError() + << "requires the number of device bodies to equal the product of " + "the MPMD axis sizes"; + } + + // Check there are no duplicated canonical logical axes. + llvm::SmallDenseSet seenAxes; + for (Value axis : getSpmdAxes()) { + if (!seenAxes.insert(axis).second) { + return emitOpError() + << "requires SPMD and MPMD axes to be pairwise disjoint"; + } + } + for (Value axis : getMpmdAxes()) { + if (!seenAxes.insert(axis).second) { + return emitOpError() + << "requires SPMD and MPMD axes to be pairwise disjoint"; + } + } + + return success(); +} + +} // namespace mlir::enzyme::distributed + +#define GET_OP_CLASSES +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedOps.cpp.inc" diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp deleted file mode 100644 index c4ce290d95..0000000000 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp +++ /dev/null @@ -1,185 +0,0 @@ -#include "mlir/IR/Builders.h" -#include "mlir/Interfaces/InferTypeOpInterface.h" -#include "llvm/ADT/TypeSwitch.h" - -#include "Dialect.h" - -using mlir::OpTrait::enzyme::distributed::ChannelDefTrait; -using mlir::OpTrait::enzyme::distributed::DeviceDefTrait; -namespace mlir::enzyme::distributed { - -/** - * Returns success if the symbol reference attribute refers to an operation of - * type T in the given symbol table. - */ -template -llvm::LogicalResult checkSymbolIsA(::mlir::SymbolTableCollection &symbol_table, - mlir::Operation *owning_op, - mlir::SymbolRefAttr attr) { - auto loc = owning_op->getLoc(); - Operation *symOp = symbol_table.lookupNearestSymbolFrom(owning_op, attr); - if (dyn_cast_or_null(symOp) == nullptr) { - mlir::emitError(loc) << "invalid symbol reference or symbol type: " << attr; - return mlir::failure(); - } - return mlir::success(); -} - -/** - * Returns success if the attribute is a symbol reference to an operation with - * the given trait in the provided symbol table. - */ -template