From fa98bf77312338a2bd0996bc3103114dd56eff8d Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 01/11] Axis segments --- src/enzyme_ad/jax/Dialect/Axis/Ops.cpp | 96 +++++++- src/enzyme_ad/jax/Dialect/Axis/Ops.td | 22 +- src/enzyme_ad/jax/Dialect/Axis/Types.td | 7 + src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp | 222 ++++++++++++++++-- src/enzyme_ad/jax/Dialect/Axis/Utilities.h | 29 ++- .../jax/Dialect/Distributed/Utilities.h | 4 + test/lit_tests/axis/roundtrip.mlir | 7 +- test/lit_tests/axis/verify.mlir | 50 +++- 8 files changed, 400 insertions(+), 37 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Axis/Ops.cpp b/src/enzyme_ad/jax/Dialect/Axis/Ops.cpp index 0df959618f..d1e27e02ea 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Ops.cpp +++ b/src/enzyme_ad/jax/Dialect/Axis/Ops.cpp @@ -22,6 +22,24 @@ static LogicalResult verifyFactorExtents(ArrayRef extents, return success(); } +static LogicalResult verifySegmentExtents(ArrayRef extents, + unsigned sourceExtent, + Operation *op) { + uint64_t sum = 0; + for (int32_t extent : extents) { + if (extent <= 0) { + return op->emitOpError() << "requires all segment extents to be > 0"; + } + sum += static_cast(extent); + } + if (sum != sourceExtent) { + return op->emitOpError() << "requires sum(segment_extents) == axis extent " + "(" + << sum << " != " << sourceExtent << ")"; + } + return success(); +} + static void computeMajorToMinorStrides(ArrayRef extents, SmallVectorImpl &strides) { // Leftmost factors are most major. @@ -151,7 +169,75 @@ LogicalResult AxisFactorOp::inferReturnTypes( return success(); } -LogicalResult AxisGroupOp::verify() { +LogicalResult AxisSegmentOp::verify() { + if (failed(getAxisProvenanceOp(getAxis()))) { + return emitOpError() + << "requires axis operand to be traceable to an op result"; + } + + auto axisIface = dyn_cast(getAxis().getType()); + if (!axisIface) { + return emitOpError() << "requires axis operand type to implement " + "AxisTypeInterface"; + } + + ArrayRef segmentExtents = getSegmentExtents(); + if (failed(verifySegmentExtents(segmentExtents, axisIface.extent(), + getOperation()))) { + return failure(); + } + + if (getAxisSegments().size() != segmentExtents.size()) { + return emitOpError() << "requires number of results to match number of " + "segment extents"; + } + + unsigned runningOffset = 0; + for (auto [idx, axisSegmentVal] : llvm::enumerate(getAxisSegments())) { + auto axisSegmentType = dyn_cast(axisSegmentVal.getType()); + if (!axisSegmentType) { + return emitOpError() << "requires all results to have AxisSegmentType"; + } + if (axisSegmentType.getAxisType() != getAxis().getType()) { + return emitOpError() << "requires result #" << idx + << " to have base axis type equal to operand type"; + } + if (axisSegmentType.getExtent() != + static_cast(segmentExtents[idx])) { + return emitOpError() << "requires result #" << idx + << " extent to match segment extent"; + } + if (axisSegmentType.getOffset() != runningOffset) { + return emitOpError() << "requires result #" << idx + << " offset to match cumulative segment layout " + "(low result index maps to low axis values)"; + } + runningOffset += static_cast(segmentExtents[idx]); + } + + return success(); +} + +LogicalResult AxisSegmentOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + AxisSegmentOpAdaptor adaptor(operands, attributes, properties, regions); + ArrayRef segmentExtents = adaptor.getSegmentExtents(); + + inferredReturnTypes.reserve(segmentExtents.size()); + Type axisType = adaptor.getAxis().getType(); + unsigned runningOffset = 0; + for (int32_t segmentExtent : segmentExtents) { + inferredReturnTypes.push_back(AxisSegmentType::get( + context, axisType, static_cast(segmentExtent), + runningOffset)); + runningOffset += static_cast(segmentExtent); + } + return success(); +} + +LogicalResult AxisProductOp::verify() { uint64_t extentProduct = 1; for (Value factor : getFactors()) { if (!isa(factor)) { @@ -169,18 +255,18 @@ LogicalResult AxisGroupOp::verify() { extentProduct *= static_cast(factorType.getExtent()); } - if (getGroup().getType().getExtent() != extentProduct) { + if (getProduct().getType().getExtent() != extentProduct) { return emitOpError() - << "requires group extent to equal product of factor extents"; + << "requires product extent to equal product of factor extents"; } return success(); } -LogicalResult AxisGroupOp::inferReturnTypes( +LogicalResult AxisProductOp::inferReturnTypes( MLIRContext *context, std::optional location, ValueRange operands, DictionaryAttr attributes, PropertyRef properties, RegionRange regions, SmallVectorImpl &inferredReturnTypes) { - AxisGroupOpAdaptor adaptor(operands, attributes, properties, regions); + AxisProductOpAdaptor adaptor(operands, attributes, properties, regions); uint64_t extentProduct = 1; for (Value factor : adaptor.getFactors()) { auto factorType = dyn_cast(factor.getType()); diff --git a/src/enzyme_ad/jax/Dialect/Axis/Ops.td b/src/enzyme_ad/jax/Dialect/Axis/Ops.td index fcb6712eeb..590cbf7314 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Ops.td @@ -35,12 +35,28 @@ def AxisFactorOp : AxisOp<"factor", [Pure, DeclareOpInterfaceMethods]> { +def AxisSegmentOp : AxisOp<"segment", [Pure, DeclareOpInterfaceMethods]> { let description = [{ - Recombines a leftmost-major factor list into one virtual factor group. + Splits one canonical axis into left-to-right contiguous linear segments. + Segment result index is value-ordered: lower result indices denote lower + axis values. Concretely, result #0 starts at offset 0 and each following + result starts at the cumulative extent of all prior segments. + }]; + let arguments = (ins + AnyType:$axis, + DenseI32ArrayAttr:$segment_extents + ); + let results = (outs Variadic:$axis_segments); + let assemblyFormat = "$axis $segment_extents attr-dict `:` type($axis)"; + let hasVerifier = 1; +} + +def AxisProductOp : AxisOp<"product", [Pure, DeclareOpInterfaceMethods]> { + let description = [{ + Recombines a leftmost-major factor list into one virtual factor product. }]; let arguments = (ins Variadic:$factors); - let results = (outs FactorGroupType:$group); + let results = (outs FactorGroupType:$product); let assemblyFormat = "$factors attr-dict `:` type($factors)"; let hasVerifier = 1; } diff --git a/src/enzyme_ad/jax/Dialect/Axis/Types.td b/src/enzyme_ad/jax/Dialect/Axis/Types.td index 723b114681..3bb87c9efd 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Types.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Types.td @@ -25,6 +25,13 @@ def AxisFactorType : AxisDialectType<"AxisFactor", "axis_factor"> { let assemblyFormat = "`<` $axisType `,` $extent `,` $stride `>`"; } +// Contiguous linear segment of a canonical axis with static extent and offset. +def AxisSegmentType : AxisDialectType<"AxisSegment", "axis_segment"> { + let parameters = + (ins "::mlir::Type":$axisType, "unsigned":$extent, "unsigned":$offset); + let assemblyFormat = "`<` $axisType `,` $extent `,` $offset `>`"; +} + // Virtual axis reconstructed from a list of factors. def FactorGroupType : AxisDialectType<"FactorGroup", "factor_group"> { let parameters = (ins "unsigned":$extent); diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp index 4a98fd7e7e..9b1282a2c0 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp @@ -1,5 +1,7 @@ #include "Utilities.h" +#include + namespace mlir::enzyme::axis { // Dispatches alias checks for canonical axes. Canonical axes are @@ -21,14 +23,49 @@ static bool areAxesEquivalent(Value lhs, Value rhs) { } // Tests if two axis factors are disjoint members of some valid factorization -// of a shared source axis. Assumes both factors are derived from the same -// source axis. -static bool arePairwiseFactorsDisjoint(const AxisFactorType &f1, - const AxisFactorType &f2) { - unsigned majorStride = f1.getStride(); - unsigned majorExtent = f1.getExtent(); - unsigned minorStride = f2.getStride(); - unsigned minorExtent = f2.getExtent(); +// of a shared source axis. +bool arePairwiseFactorsDisjoint(Value lhsFactor, Value rhsFactor, + Value lhsProvenanceAxis, + Value rhsProvenanceAxis) { + auto lhsType = dyn_cast(lhsFactor.getType()); + auto rhsType = dyn_cast(rhsFactor.getType()); + assert(lhsType && "factor value must have AxisFactorType"); + assert(rhsType && "factor value must have AxisFactorType"); + if (!lhsType || !rhsType) { + return false; + } + + Value lhsAxis = lhsProvenanceAxis; + if (!lhsAxis) { + auto lhsProvenance = + getFactorProvenanceAxis(cast>(lhsFactor)); + assert(succeeded(lhsProvenance) && "factor must have a provenance axis"); + if (failed(lhsProvenance)) { + return false; + } + lhsAxis = *lhsProvenance; + } + + Value rhsAxis = rhsProvenanceAxis; + if (!rhsAxis) { + auto rhsProvenance = + getFactorProvenanceAxis(cast>(rhsFactor)); + assert(succeeded(rhsProvenance) && "factor must have a provenance axis"); + if (failed(rhsProvenance)) { + return false; + } + rhsAxis = *rhsProvenance; + } + + // Factors from different canonical axes are disjoint by definition. + if (!areAxesEquivalent(lhsAxis, rhsAxis)) { + return true; + } + + unsigned majorStride = lhsType.getStride(); + unsigned majorExtent = lhsType.getExtent(); + unsigned minorStride = rhsType.getStride(); + unsigned minorExtent = rhsType.getExtent(); if (majorStride < minorStride) { std::swap(majorStride, minorStride); std::swap(majorExtent, minorExtent); @@ -59,6 +96,13 @@ int getFactorExtent(Value factor) { return static_cast(factorType.getExtent()); } +// Asserts a segment type and gets the extent. +int getSegmentExtent(Value segment) { + auto segmentType = dyn_cast(segment.getType()); + assert(segmentType && "segment type must be AxisSegmentType"); + return static_cast(segmentType.getExtent()); +} + // Returns the defining op for a canonical axis SSA value. FailureOr getAxisProvenanceOp(Value axis) { auto result = dyn_cast(axis); @@ -77,14 +121,23 @@ FailureOr getFactorProvenanceAxis(TypedValue factor) { return failure(); } -// Returns the factor list used to build a factor-group SSA value. +// Returns the defining source axis for a segment value. +FailureOr getSegmentProvenanceAxis(TypedValue segment) { + if (auto axisSegment = segment.getDefiningOp()) { + return axisSegment.getAxis(); + } + + return failure(); +} + +// Returns the factor list used to build a factor-product SSA value. FailureOr -getGroupProvenanceFactors(TypedValue factorGroup) { - auto groupOp = factorGroup.getDefiningOp(); - if (!groupOp) { +getProductProvenanceFactors(TypedValue factorProduct) { + auto productOp = factorProduct.getDefiningOp(); + if (!productOp) { return failure(); } - return groupOp.getFactors(); + return productOp.getFactors(); } // Checks factor compatibility and pairwise non-overlap metadata. @@ -97,7 +150,6 @@ bool areFactorsDisjoint(ValueRange factors) { "factor disjointness uses quadratic pairwise checks"); struct FactorInfo { - AxisFactorType factorType; Value provenance; }; @@ -116,18 +168,108 @@ bool areFactorsDisjoint(ValueRange factors) { auto provenance = getFactorProvenanceAxis(cast>(factor)); assert(succeeded(provenance) && "factor must have a provenance axis"); - cachedFactors.push_back({factorType, *provenance}); + cachedFactors.push_back({*provenance}); } for (size_t i = 0; i < cachedFactors.size(); ++i) { for (size_t j = i + 1; j < cachedFactors.size(); ++j) { - const auto &[lhsType, lhsAxis] = cachedFactors[i]; - const auto &[rhsType, rhsAxis] = cachedFactors[j]; + Value lhsAxis = cachedFactors[i].provenance; + Value rhsAxis = cachedFactors[j].provenance; if (!areAxesEquivalent(lhsAxis, rhsAxis)) { continue; } - if (!arePairwiseFactorsDisjoint(lhsType, rhsType)) { + if (!arePairwiseFactorsDisjoint(factors[i], factors[j], lhsAxis, + rhsAxis)) { + return false; + } + } + } + + return true; +} + +// Checks segment pairwise non-overlap metadata. +bool arePairwiseSegmentsDisjoint(Value lhsSegment, Value rhsSegment, + Value lhsProvenanceAxis, + Value rhsProvenanceAxis) { + auto lhsType = dyn_cast(lhsSegment.getType()); + auto rhsType = dyn_cast(rhsSegment.getType()); + assert(lhsType && "segment value must have AxisSegmentType"); + assert(rhsType && "segment value must have AxisSegmentType"); + if (!lhsType || !rhsType) { + return false; + } + + Value lhsAxis = lhsProvenanceAxis; + if (!lhsAxis) { + auto lhsProvenance = + getSegmentProvenanceAxis(cast>(lhsSegment)); + assert(succeeded(lhsProvenance) && "segment must have a provenance axis"); + if (failed(lhsProvenance)) { + return false; + } + lhsAxis = *lhsProvenance; + } + + Value rhsAxis = rhsProvenanceAxis; + if (!rhsAxis) { + auto rhsProvenance = + getSegmentProvenanceAxis(cast>(rhsSegment)); + assert(succeeded(rhsProvenance) && "segment must have a provenance axis"); + if (failed(rhsProvenance)) { + return false; + } + rhsAxis = *rhsProvenance; + } + + // Segments from different canonical axes are disjoint by definition. + if (!areAxesEquivalent(lhsAxis, rhsAxis)) { + return true; + } + + uint64_t lhsStart = static_cast(lhsType.getOffset()); + uint64_t lhsEnd = lhsStart + static_cast(lhsType.getExtent()); + uint64_t rhsStart = static_cast(rhsType.getOffset()); + uint64_t rhsEnd = rhsStart + static_cast(rhsType.getExtent()); + return lhsEnd <= rhsStart || rhsEnd <= lhsStart; +} + +// Checks segment group pairwise non-overlap metadata. +bool areSegmentsDisjoint(ValueRange segments) { + if (segments.empty()) { + return true; + } + + assert(segments.size() < 100 && + "segment disjointness uses quadratic pairwise checks"); + + SmallVector provenanceAxes; + provenanceAxes.reserve(segments.size()); + for (Value segment : segments) { + auto segmentType = dyn_cast(segment.getType()); + assert(segmentType && "segment value must have AxisSegmentType"); + if (!segmentType) { + return false; + } + assert(segmentType.getExtent() > 0 && "segment extent must be positive"); + + auto provenance = + getSegmentProvenanceAxis(cast>(segment)); + assert(succeeded(provenance) && "segment must have a provenance axis"); + if (failed(provenance)) { + return false; + } + provenanceAxes.push_back(*provenance); + } + + for (size_t i = 0; i < segments.size(); ++i) { + for (size_t j = i + 1; j < segments.size(); ++j) { + if (!areAxesEquivalent(provenanceAxes[i], provenanceAxes[j])) { + continue; + } + if (!arePairwiseSegmentsDisjoint(segments[i], segments[j], + provenanceAxes[i], provenanceAxes[j])) { return false; } } @@ -165,4 +307,48 @@ bool areFactorsComplete(Value axis, ValueRange factors) { return product == static_cast(getAxisExtent(axis)); } +// Checks that segments reconstruct the full source axis interval [0, extent). +bool areSegmentsComplete(Value axis, ValueRange segments) { + if (segments.empty() || !areSegmentsDisjoint(segments)) { + return false; + } + + SmallVector> intervals; + intervals.reserve(segments.size()); + + for (Value segment : segments) { + auto segmentType = dyn_cast(segment.getType()); + assert(segmentType && "segment value must have AxisSegmentType"); + if (!segmentType) { + return false; + } + + auto provenance = + getSegmentProvenanceAxis(cast>(segment)); + assert(succeeded(provenance) && "segment must have a provenance axis"); + if (failed(provenance) || *provenance != axis) { + return false; + } + + uint64_t start = static_cast(segmentType.getOffset()); + uint64_t end = start + static_cast(segmentType.getExtent()); + intervals.emplace_back(start, end); + } + + std::sort(intervals.begin(), intervals.end()); + if (intervals.front().first != 0) { + return false; + } + + uint64_t cursor = 0; + for (auto [start, end] : intervals) { + if (start != cursor) { + return false; + } + cursor = end; + } + + return cursor == static_cast(getAxisExtent(axis)); +} + } // namespace mlir::enzyme::axis diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h index 6c1c96b764..d031426b05 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h @@ -11,6 +11,9 @@ int getAxisExtent(::mlir::Value axis); // Returns the static extent for any factor SSA value. int getFactorExtent(::mlir::Value factor); +// Returns the static extent for any segment SSA value. +int getSegmentExtent(::mlir::Value segment); + // Returns the defining op for a canonical axis SSA value. ::mlir::FailureOr<::mlir::Operation *> getAxisProvenanceOp(::mlir::Value axis); @@ -18,16 +21,38 @@ ::mlir::FailureOr<::mlir::Operation *> getAxisProvenanceOp(::mlir::Value axis); ::mlir::FailureOr<::mlir::Value> getFactorProvenanceAxis(::mlir::TypedValue factor); -// Returns the factor list used to build a factor-group SSA value. +// Resolves the source canonical axis used to produce a segment value. +::mlir::FailureOr<::mlir::Value> +getSegmentProvenanceAxis(::mlir::TypedValue segment); + +// Returns the factor list used to build a factor-product SSA value. ::mlir::FailureOr<::mlir::ValueRange> -getGroupProvenanceFactors(::mlir::TypedValue factorGroup); +getProductProvenanceFactors(::mlir::TypedValue factorProduct); + +// Checks that factors are pairwise non-overlapping for one source axis. +bool arePairwiseFactorsDisjoint(::mlir::Value lhsFactor, + ::mlir::Value rhsFactor, + ::mlir::Value lhsProvenanceAxis = nullptr, + ::mlir::Value rhsProvenanceAxis = nullptr); // Checks that factors are pairwise non-overlapping for one source axis. bool areFactorsDisjoint(::mlir::ValueRange factors); +// Checks that segment intervals are pairwise non-overlapping per source axis. +bool arePairwiseSegmentsDisjoint(::mlir::Value lhsSegment, + ::mlir::Value rhsSegment, + ::mlir::Value lhsProvenanceAxis = nullptr, + ::mlir::Value rhsProvenanceAxis = nullptr); + +// Checks that segment intervals are pairwise non-overlapping per source axis. +bool areSegmentsDisjoint(::mlir::ValueRange segments); + // Checks that factors cover an axis exactly and therefore are disjoint. bool areFactorsComplete(::mlir::Value axis, ::mlir::ValueRange factors); +// Checks that segments exactly cover [0, axis extent) with no overlap or gaps. +bool areSegmentsComplete(::mlir::Value axis, ::mlir::ValueRange segments); + } // namespace mlir::enzyme::axis #endif // ENZYME_AD_JAX_DIALECT_AXIS_UTILITIES_H diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h index 39673deca2..7abbde75d9 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h @@ -26,9 +26,13 @@ ::mlir::FailureOr resolveSymbolOpFromAttr(::mlir::Operation *from, using ::mlir::enzyme::axis::areFactorsComplete; using ::mlir::enzyme::axis::areFactorsDisjoint; +using ::mlir::enzyme::axis::areSegmentsComplete; +using ::mlir::enzyme::axis::areSegmentsDisjoint; using ::mlir::enzyme::axis::getAxisExtent; using ::mlir::enzyme::axis::getFactorExtent; using ::mlir::enzyme::axis::getFactorProvenanceAxis; +using ::mlir::enzyme::axis::getSegmentExtent; +using ::mlir::enzyme::axis::getSegmentProvenanceAxis; } // namespace mlir::enzyme::distributed diff --git a/test/lit_tests/axis/roundtrip.mlir b/test/lit_tests/axis/roundtrip.mlir index d550be83ae..435ba56632 100644 --- a/test/lit_tests/axis/roundtrip.mlir +++ b/test/lit_tests/axis/roundtrip.mlir @@ -5,7 +5,9 @@ func.func @roundtrip_axis_ops() -> (!axis.shape_axis, 1>, !axis. %axis1 = axis.getaxis tensor<6x4xf32> 1 %f0, %f1 = axis.factor %axis0 [2, 3] : !axis.shape_axis, 0> - %g = axis.group %f0, %f1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> + %g = axis.product %f0, %f1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> + + %s0, %s1 = axis.segment %axis1 [1, 3] : !axis.shape_axis, 1> return %axis1, %g : !axis.shape_axis, 1>, !axis.factor_group<6> } @@ -14,5 +16,6 @@ func.func @roundtrip_axis_ops() -> (!axis.shape_axis, 1>, !axis. // CHECK: %[[AX0:.*]] = axis.getaxis tensor<6x4xf32> 0 // CHECK: %[[AX1:.*]] = axis.getaxis tensor<6x4xf32> 1 // CHECK: %[[FPAIR:.*]]:2 = axis.factor %[[AX0]] [2, 3] : !axis.shape_axis, 0> -// CHECK: %[[G:.*]] = axis.group %[[FPAIR]]#0, %[[FPAIR]]#1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> +// CHECK: %[[G:.*]] = axis.product %[[FPAIR]]#0, %[[FPAIR]]#1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> +// CHECK: %[[SEGS:.*]]:2 = axis.segment %[[AX1]] [1, 3] : !axis.shape_axis, 1> // CHECK: return %[[AX1]], %[[G]] : !axis.shape_axis, 1>, !axis.factor_group<6> diff --git a/test/lit_tests/axis/verify.mlir b/test/lit_tests/axis/verify.mlir index ed40ef5d0a..0652041f2c 100644 --- a/test/lit_tests/axis/verify.mlir +++ b/test/lit_tests/axis/verify.mlir @@ -74,29 +74,65 @@ func.func @factor_requires_stride_convention() { // ----- -func.func @group_requires_factor_op_results(%arg0: !axis.axis_factor, 0>, 2, 3>) { +func.func @product_requires_factor_op_results(%arg0: !axis.axis_factor, 0>, 2, 3>) { // expected-error @+1 {{requires factor operands to be op results}} - %g = axis.group %arg0 : !axis.axis_factor, 0>, 2, 3> + %g = axis.product %arg0 : !axis.axis_factor, 0>, 2, 3> return } // ----- -func.func @group_requires_axis_factor_producer() { +func.func @product_requires_axis_factor_producer() { %axis = axis.getaxis tensor<6xf32> 0 %f0, %f1 = axis.factor %axis [2, 3] : !axis.shape_axis, 0> %fake = builtin.unrealized_conversion_cast %f0 : !axis.axis_factor, 0>, 2, 3> to !axis.axis_factor, 0>, 2, 3> // expected-error @+1 {{requires factor operands to be produced by axis.factor}} - %g = axis.group %fake, %f1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> + %g = axis.product %fake, %f1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> return } // ----- -func.func @group_requires_extent_product_match() { +func.func @product_requires_extent_product_match() { %axis = axis.getaxis tensor<6xf32> 0 %f0, %f1 = axis.factor %axis [2, 3] : !axis.shape_axis, 0> - // expected-error @+1 {{requires group extent to equal product of factor extents}} - %g = "axis.group"(%f0, %f1) : (!axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1>) -> !axis.factor_group<5> + // expected-error @+1 {{requires product extent to equal product of factor extents}} + %g = "axis.product"(%f0, %f1) : (!axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1>) -> !axis.factor_group<5> + return +} + +// ----- + +func.func @segment_requires_positive_extents() { + %axis = axis.getaxis tensor<6xf32> 0 + // expected-error @+1 {{requires all segment extents to be > 0}} + %s0, %s1 = axis.segment %axis [0, 6] : !axis.shape_axis, 0> + return +} + +// ----- + +func.func @segment_requires_extent_sum_match() { + %axis = axis.getaxis tensor<6xf32> 0 + // expected-error @+1 {{requires sum(segment_extents) == axis extent (4 != 6)}} + %s0, %s1 = axis.segment %axis [2, 2] : !axis.shape_axis, 0> + return +} + +// ----- + +func.func @segment_requires_offset_layout() { + %axis = axis.getaxis tensor<6xf32> 0 + // expected-error @+1 {{requires result #1 offset to match cumulative segment layout (low result index maps to low axis values)}} + %s0, %s1 = "axis.segment"(%axis) {segment_extents = array} : (!axis.shape_axis, 0>) -> (!axis.axis_segment, 0>, 2, 0>, !axis.axis_segment, 0>, 4, 1>) + return +} + +// ----- + +func.func @segment_requires_first_result_to_start_at_low_values() { + %axis = axis.getaxis tensor<6xf32> 0 + // expected-error @+1 {{requires result #0 offset to match cumulative segment layout (low result index maps to low axis values)}} + %s0, %s1 = "axis.segment"(%axis) {segment_extents = array} : (!axis.shape_axis, 0>) -> (!axis.axis_segment, 0>, 2, 1>, !axis.axis_segment, 0>, 4, 3>) return } From 08de24189900738e44971d197c2d4624af333ef3 Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 02/11] Distributed function contexts --- src/enzyme_ad/jax/Dialect/Axis/Ops.td | 8 +- src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp | 66 ++++++++ src/enzyme_ad/jax/Dialect/Axis/Utilities.h | 6 + .../jax/Dialect/Distributed/FunctionOps.cpp | 152 ++++++++++++++++++ .../Distributed/MeshComputationOps.cpp | 10 +- src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 44 ++++- .../jax/Dialect/Distributed/Types.td | 8 +- .../jax/Dialect/Distributed/Utilities.cpp | 35 +++- .../jax/Dialect/Distributed/Utilities.h | 18 +-- .../distributed/abcda_batch_matmul.mlir | 43 ----- .../branched_pipeline_friendly.mlir | 33 ---- .../call_complete_replication.mlir | 54 +++++++ .../distributed/call_null_replication.mlir | 79 +++++++++ .../distributed/call_permutation_equiv.mlir | 52 ++++++ .../distributed/call_permutation_stress.mlir | 26 +++ .../distributed/call_verify_extra.mlir | 56 +++++++ test/lit_tests/distributed/concat.mlir | 32 ---- .../diamond_pipeline_friendly.mlir | 31 ---- test/lit_tests/distributed/dotadd.mlir | 23 --- .../distributed/function_verify.mlir | 44 +++++ .../innerouterinner_dataparallel.mlir | 24 --- .../innerouterinner_operatorparallel.mlir | 23 --- .../innerouterinner_pipeline_friendly.mlir | 19 --- test/lit_tests/distributed/mesh_verify.mlir | 30 ++++ test/lit_tests/distributed/roundtrip.mlir | 40 +++++ test/lit_tests/distributed/test_timing.mlir | 47 ------ test/lit_tests/distributed/two_matmul.mlir | 18 --- 27 files changed, 703 insertions(+), 318 deletions(-) create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/FunctionOps.cpp delete mode 100644 test/lit_tests/distributed/abcda_batch_matmul.mlir delete mode 100644 test/lit_tests/distributed/branched_pipeline_friendly.mlir create mode 100644 test/lit_tests/distributed/call_complete_replication.mlir create mode 100644 test/lit_tests/distributed/call_null_replication.mlir create mode 100644 test/lit_tests/distributed/call_permutation_equiv.mlir create mode 100644 test/lit_tests/distributed/call_permutation_stress.mlir create mode 100644 test/lit_tests/distributed/call_verify_extra.mlir delete mode 100644 test/lit_tests/distributed/concat.mlir delete mode 100644 test/lit_tests/distributed/diamond_pipeline_friendly.mlir delete mode 100644 test/lit_tests/distributed/dotadd.mlir create mode 100644 test/lit_tests/distributed/function_verify.mlir delete mode 100644 test/lit_tests/distributed/innerouterinner_dataparallel.mlir delete mode 100644 test/lit_tests/distributed/innerouterinner_operatorparallel.mlir delete mode 100644 test/lit_tests/distributed/innerouterinner_pipeline_friendly.mlir create mode 100644 test/lit_tests/distributed/mesh_verify.mlir create mode 100644 test/lit_tests/distributed/roundtrip.mlir delete mode 100644 test/lit_tests/distributed/test_timing.mlir delete mode 100644 test/lit_tests/distributed/two_matmul.mlir diff --git a/src/enzyme_ad/jax/Dialect/Axis/Ops.td b/src/enzyme_ad/jax/Dialect/Axis/Ops.td index 590cbf7314..a7cd5f4a0d 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Ops.td @@ -9,7 +9,7 @@ include "Dialect.td" include "Interfaces.td" include "Types.td" -def AxisGetAxisOp : AxisOp<"getaxis", [Pure, DeclareOpInterfaceMethods]> { +def AxisGetAxisOp : AxisOp<"getaxis", [MetadataOpTrait, DeclareOpInterfaceMethods]> { let description = [{ Materializes a canonical ShapeAxis value for one ranked, static shape index. }]; @@ -22,7 +22,7 @@ def AxisGetAxisOp : AxisOp<"getaxis", [Pure, DeclareOpInterfaceMethods]> { +def AxisFactorOp : AxisOp<"factor", [MetadataOpTrait, DeclareOpInterfaceMethods]> { let description = [{ Factors one canonical axis into leftmost-major factors with static extents. }]; @@ -35,7 +35,7 @@ def AxisFactorOp : AxisOp<"factor", [Pure, DeclareOpInterfaceMethods]> { +def AxisSegmentOp : AxisOp<"segment", [MetadataOpTrait, DeclareOpInterfaceMethods]> { let description = [{ Splits one canonical axis into left-to-right contiguous linear segments. Segment result index is value-ordered: lower result indices denote lower @@ -51,7 +51,7 @@ def AxisSegmentOp : AxisOp<"segment", [Pure, DeclareOpInterfaceMethods]> { +def AxisProductOp : AxisOp<"product", [MetadataOpTrait, DeclareOpInterfaceMethods]> { let description = [{ Recombines a leftmost-major factor list into one virtual factor product. }]; diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp index 9b1282a2c0..fe17835e4f 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp @@ -189,6 +189,72 @@ bool areFactorsDisjoint(ValueRange factors) { return true; } +// Compares two factor lists as index-space descriptors, ignoring ordering. +// This is multiset equality over (extent, stride, provenance-axis equivalence) +// and is intentionally permutation-invariant. +bool areFactorIndexSpacesEqual(ValueRange lhsFactors, ValueRange rhsFactors) { + if (lhsFactors.size() != rhsFactors.size()) { + return false; + } + + struct FactorInfo { + Value provenance; + unsigned extent; + unsigned stride; + }; + + auto buildFactorInfo = [](ValueRange factors, + SmallVectorImpl &out) -> bool { + out.clear(); + out.reserve(factors.size()); + for (Value factor : factors) { + auto factorType = dyn_cast(factor.getType()); + if (!factorType) { + return false; + } + auto provenance = + getFactorProvenanceAxis(cast>(factor)); + if (failed(provenance)) { + return false; + } + out.push_back( + {*provenance, factorType.getExtent(), factorType.getStride()}); + } + return true; + }; + + SmallVector lhsInfo; + SmallVector rhsInfo; + if (!buildFactorInfo(lhsFactors, lhsInfo) || + !buildFactorInfo(rhsFactors, rhsInfo)) { + return false; + } + + SmallVector rhsMatched(rhsInfo.size(), false); + for (const FactorInfo &lhs : lhsInfo) { + bool foundMatch = false; + for (auto [rhsIndex, rhs] : llvm::enumerate(rhsInfo)) { + if (rhsMatched[rhsIndex]) { + continue; + } + if (lhs.extent != rhs.extent || lhs.stride != rhs.stride) { + continue; + } + if (!areAxesEquivalent(lhs.provenance, rhs.provenance)) { + continue; + } + rhsMatched[rhsIndex] = true; + foundMatch = true; + break; + } + if (!foundMatch) { + return false; + } + } + + return true; +} + // Checks segment pairwise non-overlap metadata. bool arePairwiseSegmentsDisjoint(Value lhsSegment, Value rhsSegment, Value lhsProvenanceAxis, diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h index d031426b05..6e9c8e560f 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h @@ -47,6 +47,12 @@ bool arePairwiseSegmentsDisjoint(::mlir::Value lhsSegment, // Checks that segment intervals are pairwise non-overlapping per source axis. bool areSegmentsDisjoint(::mlir::ValueRange segments); +// Returns true when two factor lists cover the same index space modulo +// permutation order. This checks multiset equality over factor metadata and +// provenance-axis equivalence, but does not require matching list order. +bool areFactorIndexSpacesEqual(::mlir::ValueRange lhsFactors, + ::mlir::ValueRange rhsFactors); + // Checks that factors cover an axis exactly and therefore are disjoint. bool areFactorsComplete(::mlir::Value axis, ::mlir::ValueRange factors); diff --git a/src/enzyme_ad/jax/Dialect/Distributed/FunctionOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/FunctionOps.cpp new file mode 100644 index 0000000000..51f076063c --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Distributed/FunctionOps.cpp @@ -0,0 +1,152 @@ +#include "Dialect.h" +#include "Utilities.h" + +namespace mlir::enzyme::distributed { + +LogicalResult DistributedFunctionOp::verify() { + if (!(*this)->getParentOfType()) { + return emitOpError() << "must be nested in distributed.MeshComputation"; + } + + ArrayAttr argumentTypesAttr = + (*this)->getAttrOfType("argument_types"); + if (!argumentTypesAttr) { + return emitOpError() << "requires argument_types attribute"; + } + for (auto [idx, argumentTypeAttr] : llvm::enumerate(argumentTypesAttr)) { + if (!isa(argumentTypeAttr)) { + return emitOpError() << "requires argument_types[" << idx + << "] to be a TypeAttr"; + } + } + + Block &bodyBlock = getBody().front(); + if (bodyBlock.getNumArguments() != argumentTypesAttr.size()) { + return emitOpError() << "requires body block argument count to equal " + << "argument_types size (" + << bodyBlock.getNumArguments() + << " != " << argumentTypesAttr.size() << ")"; + } + for (auto [idx, argumentTypeAttr] : llvm::enumerate(argumentTypesAttr)) { + Type expectedType = cast(argumentTypeAttr).getValue(); + Type actualType = bodyBlock.getArgument(idx).getType(); + if (expectedType != actualType) { + return emitOpError() << "requires body block argument #" << idx + << " to have type " << expectedType << ", but got " + << actualType; + } + } + + Operation *terminator = bodyBlock.getTerminator(); + auto yieldOp = dyn_cast_or_null(terminator); + if (!yieldOp) { + return emitOpError() + << "requires body to terminate with distributed.DistributedYield"; + } + + ArrayAttr returnTypesAttr = (*this)->getAttrOfType("return_types"); + if (!returnTypesAttr) { + return emitOpError() << "requires return_types attribute"; + } + + ValueRange yieldedValues = yieldOp.getReturns(); + if (yieldedValues.size() != returnTypesAttr.size()) { + return emitOpError() << "requires distributed.DistributedYield to yield " + << returnTypesAttr.size() << " value(s), but got " + << yieldedValues.size(); + } + + for (auto [idx, returnTypeAttr] : llvm::enumerate(returnTypesAttr)) { + auto typeAttr = dyn_cast(returnTypeAttr); + if (!typeAttr) { + return emitOpError() << "requires return_types[" << idx + << "] to be a TypeAttr"; + } + + Type expectedType = typeAttr.getValue(); + Type actualType = yieldedValues[idx].getType(); + if (expectedType != actualType) { + return emitOpError() << "requires distributed.DistributedYield operand #" + << idx << " to have type " << expectedType + << ", but got " << actualType; + } + } + + return success(); +} + +LogicalResult DistributedCallOp::verify() { + FailureOr> callerContext = + getEnclosingExecutionContext(*this); + if (failed(callerContext)) { + return emitOpError() << "must be nested in distributed.function with a " + << "FactorGroupType execution_context"; + } + + FailureOr callee = + resolveSymbolOpFromAttr(*this, getCalleeAttr()); + if (failed(callee)) { + return emitOpError() << "references unknown distributed function symbol " + << getCalleeAttr(); + } + + ArrayAttr calleeReturnTypesAttr = + (*callee)->getAttrOfType("return_types"); + if (!calleeReturnTypesAttr) { + return emitOpError() << "callee " << callee->getSymName() + << " is missing return_types attribute"; + } + if (getReturns().size() != calleeReturnTypesAttr.size()) { + return emitOpError() << "requires call result count to match callee " + << "return_types size (" << getReturns().size() + << " != " << calleeReturnTypesAttr.size() << ")"; + } + for (auto [idx, returnTypeAttr] : llvm::enumerate(calleeReturnTypesAttr)) { + auto typeAttr = dyn_cast(returnTypeAttr); + if (!typeAttr) { + return emitOpError() << "requires callee return_types[" << idx + << "] to be a TypeAttr"; + } + Type expectedType = typeAttr.getValue(); + Type actualType = getReturns()[idx].getType(); + if (expectedType != actualType) { + return emitOpError() << "requires call result #" << idx << " to have " + << "type " << expectedType << ", but got " + << actualType; + } + } + + auto calleeContext = dyn_cast>( + callee->getExecutionContext()); + auto replicateOver = + dyn_cast>(getReplicateOver()); + if (!calleeContext || !replicateOver) { + return emitOpError() << "requires both callee execution_context and " + << "replicate_over to be FactorGroupType"; + } + + FailureOr> callerFactors = + expandExecutionContextFactors(*callerContext); + FailureOr> calleeFactors = + expandExecutionContextFactors(calleeContext); + FailureOr> replicateFactors = + expandExecutionContextFactors(replicateOver); + if (failed(callerFactors) || failed(calleeFactors) || + failed(replicateFactors)) { + return emitOpError() << "requires execution contexts to be produced by " + << "axis.product so factor provenance can be checked"; + } + + llvm::SmallVector expectedCallerFactors = *calleeFactors; + expectedCallerFactors.append(replicateFactors->begin(), + replicateFactors->end()); + if (!axis::areFactorIndexSpacesEqual(*callerFactors, expectedCallerFactors)) { + return emitOpError() + << "requires caller execution context to equal callee " + << "execution_context x replicate_over (permutation-insensitive)"; + } + + return success(); +} + +} // namespace mlir::enzyme::distributed diff --git a/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp index d6600ee47d..6dbb2c201e 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp @@ -5,11 +5,13 @@ namespace mlir::enzyme::distributed { LogicalResult MeshComputationOp::verify() { Block &bodyBlock = getBody().front(); for (Operation &op : bodyBlock.getOperations()) { - if (!op.hasTrait()) { + bool isMetadataOp = op.hasTrait(); + bool isDistributedFunction = isa(op); + if (!isMetadataOp && !isDistributedFunction) { return emitOpError() - << "only static metadata ops are allowed in the mesh body; " - << "operation '" << op.getName() << "' is not marked with " - << "metadata trait"; + << "only distributed.Function and static metadata ops are " + << "allowed in the mesh body; operation '" << op.getName() + << "' is neither"; } } return success(); diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td index 24998c912e..4220b6e408 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -12,6 +12,7 @@ include "Dialect.td" include "Types.td" include "Interfaces.td" include "src/enzyme_ad/jax/Dialect/Axis/Dialect.td" +include "src/enzyme_ad/jax/Dialect/Axis/Types.td" defvar TensorType = HLO_Tensor; defvar ShardedTensorType = HLO_Tensor; // Subject to change, for now shardy just adds attributes on top of a regular tensor. @@ -45,21 +46,58 @@ def MeshComputationOp : DistributedOp<"MeshComputation", [ for optimization decisions. }]; let arguments = (ins - SymbolNameAttr:$computation_name, + SymbolNameAttr:$sym_name, FlatSymbolRefAttr:$physical_mesh ); let regions = (region SizedRegion<1>:$body); let assemblyFormat = [{ - $computation_name `mesh` $physical_mesh $body attr-dict + $sym_name `mesh` $physical_mesh $body attr-dict }]; let hasVerifier = 1; } +def DistributedFunctionOp : DistributedOp<"Function", [Symbol]> { + let description = [{ + Represents a distributed function. Functions have an execution context indicating what subspace of + the logical/physical mesh they are executed on, and can only be called from a matching enclosing + context. All arguments and return values are local values, rather than global tensors with + a sharding annotation. + }]; + let arguments = (ins + SymbolNameAttr:$sym_name, + FactorGroupType:$execution_context, + ArrayAttr:$argument_types, + ArrayAttr:$return_types + ); + let regions = (region SizedRegion<1>:$body); + let assemblyFormat = [{ + $sym_name `context` $execution_context + `:` type($execution_context) + `arg_types` $argument_types `ret_types` $return_types + $body attr-dict + }]; + let hasVerifier = 1; +} + def DistributedYieldOp : DistributedOp<"DistributedYield", [Terminator]> { let description = [{ Terminator for distributed functions. Yields local values back to the caller. }]; - let arguments = (ins Variadic:$returns); + let arguments = (ins Variadic:$returns); let assemblyFormat = "$returns type($returns) attr-dict"; } + +def DistributedCallOp : DistributedOp<"DistributedCall", []> { + let description = [{ + Calls a distributed function. The execution context must match the enclosing context. + }]; + let arguments = (ins + FlatSymbolRefAttr:$callee, + FactorGroupType:$replicate_over, + Variadic:$arguments + ); + let results = (outs Variadic:$returns); + let assemblyFormat = "$callee `replicate_over` $replicate_over $arguments attr-dict `:` type($replicate_over) `,` type($arguments) `->` type($returns)"; + let hasVerifier = 1; +} #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_OPS_TD \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Types.td b/src/enzyme_ad/jax/Dialect/Distributed/Types.td index 0997ab68e9..ea3d64cb11 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Types.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Types.td @@ -5,25 +5,25 @@ include "Dialect.td" include "Interfaces.td" include "src/enzyme_ad/jax/Dialect/Axis/Interfaces.td" -class AxisType traits = []> +class DistributedAxisType traits = []> : DistributedType] # traits>; // A natural axis defined by the hardware topology. -def PhysicalCommAxisType : AxisType<"PhysicalCommAxis", "physical_comm_axis"> { +def PhysicalCommAxisType : DistributedAxisType<"PhysicalCommAxis", "physical_comm_axis"> { let parameters = (ins "unsigned":$extent); let assemblyFormat = "`<` $extent `>`"; } // Axis of a logical mesh where hardware details are abstracted away. -def LogicalMeshAxisType : AxisType<"LogicalMeshAxis", "logical_mesh_axis"> { +def LogicalMeshAxisType : DistributedAxisType<"LogicalMeshAxis", "logical_mesh_axis"> { let parameters = (ins "unsigned":$extent); let assemblyFormat = "`<` $extent `>`"; } // A standalone axis type used to represent replication. -def ReplicationAxisType : AxisType<"ReplicationAxis", "replication_axis"> { +def ReplicationAxisType : DistributedAxisType<"ReplicationAxis", "replication_axis"> { let parameters = (ins "unsigned":$extent); let assemblyFormat = "`<` $extent `>`"; } diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp index 7aa04ef732..bc97e4529f 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp @@ -1 +1,34 @@ -#include "Utilities.h" \ No newline at end of file +#include "Utilities.h" + +namespace mlir::enzyme::distributed { + +using axis::FactorGroupType; +using axis::getProductProvenanceFactors; + +FailureOr> +getEnclosingExecutionContext(Operation *op) { + auto parentFunction = op ? op->getParentOfType() + : DistributedFunctionOp(); + if (!parentFunction) { + return failure(); + } + Value context = parentFunction.getExecutionContext(); + auto typedContext = dyn_cast>(context); + if (!typedContext) { + return failure(); + } + return typedContext; +} + +FailureOr> +expandExecutionContextFactors(TypedValue context) { + FailureOr factors = getProductProvenanceFactors(context); + if (failed(factors)) { + return failure(); + } + llvm::SmallVector copiedFactors; + copiedFactors.append(factors->begin(), factors->end()); + return copiedFactors; +} + +} // namespace mlir::enzyme::distributed \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h index 7abbde75d9..75539ef111 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h @@ -24,15 +24,15 @@ ::mlir::FailureOr resolveSymbolOpFromAttr(::mlir::Operation *from, return typedOp; } -using ::mlir::enzyme::axis::areFactorsComplete; -using ::mlir::enzyme::axis::areFactorsDisjoint; -using ::mlir::enzyme::axis::areSegmentsComplete; -using ::mlir::enzyme::axis::areSegmentsDisjoint; -using ::mlir::enzyme::axis::getAxisExtent; -using ::mlir::enzyme::axis::getFactorExtent; -using ::mlir::enzyme::axis::getFactorProvenanceAxis; -using ::mlir::enzyme::axis::getSegmentExtent; -using ::mlir::enzyme::axis::getSegmentProvenanceAxis; +// Returns the execution-context FactorGroup value from the nearest enclosing +// distributed.function. +::mlir::FailureOr<::mlir::TypedValue<::mlir::enzyme::axis::FactorGroupType>> +getEnclosingExecutionContext(::mlir::Operation *op); + +// Expands an axis.factor_group value into its defining axis.factor list. +::mlir::FailureOr<::llvm::SmallVector<::mlir::Value>> +expandExecutionContextFactors( + ::mlir::TypedValue<::mlir::enzyme::axis::FactorGroupType> context); } // namespace mlir::enzyme::distributed diff --git a/test/lit_tests/distributed/abcda_batch_matmul.mlir b/test/lit_tests/distributed/abcda_batch_matmul.mlir deleted file mode 100644 index 60ed6effae..0000000000 --- a/test/lit_tests/distributed/abcda_batch_matmul.mlir +++ /dev/null @@ -1,43 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline --sdy-insert-explicit-reshards=enable-full-version=true --sdy-reshard-to-collectives --insert-identity-reshard --shardy-to-distributed --localize-distributed-module --cse --distributed-simplify-collectives --distributed-overlap-communication-module %s > /dev/null -// Four batched square matmuls in the chain B(A(C(DA))) where A is reused, -// creating a dependency from the innermost to the third matmul. -// -// Shape: -// t1 = D @ A (innermost) -// t2 = C @ t1 -// t3 = A @ t2 (A reused here) -// t4 = B @ t3 (outermost) -// -// A, B, C, D are all [batch=8, N=512, N=512] square matrices. -// All inputs are batch-sharded on "x"; contracting dims are unsharded so -// no all-reduce arises from a single config. - -sdy.mesh @mesh = <["x"=2]> - -distributed.AxisAllToAll @ax1 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1] - -func.func @bacda_batch_matmul( - %a: tensor<8x512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}, {}]>}, - %b: tensor<8x512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}, {}]>}, - %c: tensor<8x512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}, {}]>}, - %d: tensor<8x512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}, {}]>} -) -> (tensor<8x512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}, {}]>}) { - // t1 = D @ A - %t1 = stablehlo.dot_general %d, %a, - batching_dims = [0] x [0], contracting_dims = [2] x [1] - : (tensor<8x512x512xf32>, tensor<8x512x512xf32>) -> tensor<8x512x512xf32> - // t2 = C @ t1 - %t2 = stablehlo.dot_general %c, %t1, - batching_dims = [0] x [0], contracting_dims = [2] x [1] - : (tensor<8x512x512xf32>, tensor<8x512x512xf32>) -> tensor<8x512x512xf32> - // t3 = A @ t2 (A reused) - %t3 = stablehlo.dot_general %a, %t2, - batching_dims = [0] x [0], contracting_dims = [2] x [1] - : (tensor<8x512x512xf32>, tensor<8x512x512xf32>) -> tensor<8x512x512xf32> - // t4 = B @ t3 - %t4 = stablehlo.dot_general %b, %t3, - batching_dims = [0] x [0], contracting_dims = [2] x [1] - : (tensor<8x512x512xf32>, tensor<8x512x512xf32>) -> tensor<8x512x512xf32> - return %t4 : tensor<8x512x512xf32> -} diff --git a/test/lit_tests/distributed/branched_pipeline_friendly.mlir b/test/lit_tests/distributed/branched_pipeline_friendly.mlir deleted file mode 100644 index 63372abfc8..0000000000 --- a/test/lit_tests/distributed/branched_pipeline_friendly.mlir +++ /dev/null @@ -1,33 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline --sdy-insert-explicit-reshards=enable-full-version=true --sdy-reshard-to-collectives --insert-identity-reshard --shardy-to-distributed --localize-distributed-module --cse --distributed-simplify-collectives --distributed-overlap-communication-module %s > /dev/null -// A branched, dot-only pipeline-friendly test case. -// -// Shape: -// f(g(a, b), h(c, d)) -// -// Design intent: -// 1) Two independent, uneven branches create useful local work before the -// merge. -// 2) Avoid reduction collectives in the computation. -// 3) Force gather-style movement on the right branch at the final merge by -// sharding a contracting dimension of that branch's result. - -sdy.mesh @mesh = <["x"=2]> - -distributed.AxisAllToAll @ax1 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1] - -func.func @branched_pipeline_friendly( - %a: tensor<2048x4096xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, - %b: tensor<4096x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, - %c: tensor<64x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, - %d: tensor<1024x128x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}, {}]>}, - %u: tensor<128xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}]>} -) -> (tensor<2048x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) { - %0 = stablehlo.dot %a, %b : (tensor<2048x4096xf32>, tensor<4096x64xf32>) -> tensor<2048x64xf32> - %1 = stablehlo.dot_general %0, %u, contracting_dims = [] x [] : (tensor<2048x64xf32>, tensor<128xf32>) -> tensor<2048x64x128xf32> - - %2 = stablehlo.dot_general %c, %d, contracting_dims = [1] x [0] : (tensor<64x1024xf32>, tensor<1024x128x256xf32>) -> tensor<64x128x256xf32> - - %3 = stablehlo.dot_general %1, %2, contracting_dims = [1, 2] x [0, 1] : (tensor<2048x64x128xf32>, tensor<64x128x256xf32>) -> tensor<2048x256xf32> - return %3 : tensor<2048x256xf32> -} \ No newline at end of file diff --git a/test/lit_tests/distributed/call_complete_replication.mlir b/test/lit_tests/distributed/call_complete_replication.mlir new file mode 100644 index 0000000000..3bb1ba05e1 --- /dev/null +++ b/test/lit_tests/distributed/call_complete_replication.mlir @@ -0,0 +1,54 @@ +// RUN: enzymexlamlir-opt --split-input-file --verify-diagnostics %s + +// Complete replication: caller context equals callee context x replicate_over. +module { + distributed.MeshComputation @mc0 mesh @mesh0 { + %ax = axis.getaxis tensor<12xf32> 0 + %f0, %f1, %f2 = axis.factor %ax [2, 2, 3] : !axis.shape_axis, 0> + %ctx_callee = axis.product %f0, %f1 : !axis.axis_factor, 0>, 2, 6>, !axis.axis_factor, 0>, 2, 3> + %ctx_rep = axis.product %f2 : !axis.axis_factor, 0>, 3, 1> + %ctx_caller = axis.product %f0, %f1, %f2 : !axis.axis_factor, 0>, 2, 6>, !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> + + distributed.Function @callee context %ctx_callee : !axis.factor_group<4> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx_caller : !axis.factor_group<12> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + %r = distributed.DistributedCall @callee replicate_over %ctx_rep %arg0 : !axis.factor_group<3>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} + +// ----- + +// Elementwise callee context + mismatching replication should fail. +module { + distributed.MeshComputation @mc1 mesh @mesh0 { + %ax4 = axis.getaxis tensor<4xf32> 0 + %f4 = axis.factor %ax4 [4] : !axis.shape_axis, 0> + %ctx4 = axis.product %f4 : !axis.axis_factor, 0>, 4, 1> + + %ax2 = axis.getaxis tensor<2xf32> 0 + %f2 = axis.factor %ax2 [2] : !axis.shape_axis, 0> + %ctx2 = axis.product %f2 : !axis.axis_factor, 0>, 2, 1> + + %ax1 = axis.getaxis tensor<1xf32> 0 + %f1 = axis.factor %ax1 [1] : !axis.shape_axis, 0> + %ctx1 = axis.product %f1 : !axis.axis_factor, 0>, 1, 1> + + distributed.Function @callee context %ctx1 : !axis.factor_group<1> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx4 : !axis.factor_group<4> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + // expected-error @+1 {{requires caller execution context to equal callee execution_context x replicate_over (permutation-insensitive)}} + %r = distributed.DistributedCall @callee replicate_over %ctx2 %arg0 : !axis.factor_group<2>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} diff --git a/test/lit_tests/distributed/call_null_replication.mlir b/test/lit_tests/distributed/call_null_replication.mlir new file mode 100644 index 0000000000..e71a400724 --- /dev/null +++ b/test/lit_tests/distributed/call_null_replication.mlir @@ -0,0 +1,79 @@ +// RUN: enzymexlamlir-opt --split-input-file --verify-diagnostics %s + +// Null replication (modeled as axis.product with zero factors). +module { + distributed.MeshComputation @mc0 mesh @mesh0 { + %ax_main = axis.getaxis tensor<6xf32> 0 + %fmain = axis.factor %ax_main [6] : !axis.shape_axis, 0> + %ctx = axis.product %fmain : !axis.axis_factor, 0>, 6, 1> + + %rep = "axis.product"() : () -> !axis.factor_group<1> + + distributed.Function @callee context %ctx : !axis.factor_group<6> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx : !axis.factor_group<6> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + %r = distributed.DistributedCall @callee replicate_over %rep %arg0 : !axis.factor_group<1>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} + +// ----- + +// Completely mismatching spaces should fail. +module { + distributed.MeshComputation @mc1 mesh @mesh0 { + %ax6 = axis.getaxis tensor<6xf32> 0 + %f6 = axis.factor %ax6 [6] : !axis.shape_axis, 0> + %ctx6 = axis.product %f6 : !axis.axis_factor, 0>, 6, 1> + + %ax3 = axis.getaxis tensor<3xf32> 0 + %f3 = axis.factor %ax3 [3] : !axis.shape_axis, 0> + %ctx3 = axis.product %f3 : !axis.axis_factor, 0>, 3, 1> + + distributed.Function @callee context %ctx3 : !axis.factor_group<3> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx6 : !axis.factor_group<6> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + // expected-error @+1 {{requires caller execution context to equal callee execution_context x replicate_over (permutation-insensitive)}} + %r = distributed.DistributedCall @callee replicate_over %ctx3 %arg0 : !axis.factor_group<3>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} + +// ----- + +// Null replication with mismatching caller/callee context should fail. +module { + distributed.MeshComputation @mc2 mesh @mesh0 { + %ax6 = axis.getaxis tensor<6xf32> 0 + %f6 = axis.factor %ax6 [6] : !axis.shape_axis, 0> + %ctx6 = axis.product %f6 : !axis.axis_factor, 0>, 6, 1> + + %ax3 = axis.getaxis tensor<3xf32> 0 + %f3 = axis.factor %ax3 [3] : !axis.shape_axis, 0> + %ctx3 = axis.product %f3 : !axis.axis_factor, 0>, 3, 1> + + %rep = "axis.product"() : () -> !axis.factor_group<1> + + distributed.Function @callee context %ctx3 : !axis.factor_group<3> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx6 : !axis.factor_group<6> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + // expected-error @+1 {{requires caller execution context to equal callee execution_context x replicate_over (permutation-insensitive)}} + %r = distributed.DistributedCall @callee replicate_over %rep %arg0 : !axis.factor_group<1>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} diff --git a/test/lit_tests/distributed/call_permutation_equiv.mlir b/test/lit_tests/distributed/call_permutation_equiv.mlir new file mode 100644 index 0000000000..e840ad904e --- /dev/null +++ b/test/lit_tests/distributed/call_permutation_equiv.mlir @@ -0,0 +1,52 @@ +// RUN: enzymexlamlir-opt --split-input-file --verify-diagnostics %s + +// Trivial permutation equivalence: same factor multiset in different order. +module { + distributed.MeshComputation @mc0 mesh @mesh0 { + %ax = axis.getaxis tensor<12xf32> 0 + %f0, %f1, %f2 = axis.factor %ax [2, 2, 3] : !axis.shape_axis, 0> + + %ctx_callee = axis.product %f0, %f1 : !axis.axis_factor, 0>, 2, 6>, !axis.axis_factor, 0>, 2, 3> + %ctx_rep = axis.product %f2 : !axis.axis_factor, 0>, 3, 1> + // Permute the caller product order. + %ctx_caller_perm = axis.product %f2, %f1, %f0 : !axis.axis_factor, 0>, 3, 1>, !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 2, 6> + + distributed.Function @callee context %ctx_callee : !axis.factor_group<4> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx_caller_perm : !axis.factor_group<12> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + %r = distributed.DistributedCall @callee replicate_over %ctx_rep %arg0 : !axis.factor_group<3>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} + +// ----- + +// Same extents but different provenance axes should fail equality. +module { + distributed.MeshComputation @mc1 mesh @mesh0 { + %axA = axis.getaxis tensor<6xf32> 0 + %a0, %a1 = axis.factor %axA [2, 3] : !axis.shape_axis, 0> + %ctxA = axis.product %a0, %a1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> + + %axB = axis.getaxis tensor<6xf32> 0 + %b0, %b1 = axis.factor %axB [2, 3] : !axis.shape_axis, 0> + %ctxB = axis.product %b0, %b1 : !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> + + distributed.Function @callee context %ctxA : !axis.factor_group<6> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctxB : !axis.factor_group<6> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + // expected-error @+1 {{requires caller execution context to equal callee execution_context x replicate_over (permutation-insensitive)}} + %r = distributed.DistributedCall @callee replicate_over %ctxB %arg0 : !axis.factor_group<6>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} diff --git a/test/lit_tests/distributed/call_permutation_stress.mlir b/test/lit_tests/distributed/call_permutation_stress.mlir new file mode 100644 index 0000000000..30d2d8b9da --- /dev/null +++ b/test/lit_tests/distributed/call_permutation_stress.mlir @@ -0,0 +1,26 @@ +// RUN: enzymexlamlir-opt --verify-diagnostics %s + +// Stress order-insensitive factor-space equality with many permuted factors. +module { + distributed.MeshComputation @mc mesh @mesh0 { + %ax = axis.getaxis tensor<360xf32> 0 + %f0, %f1, %f2, %f3, %f4, %f5 = axis.factor %ax [2, 2, 2, 3, 3, 5] : !axis.shape_axis, 0> + + %ctx_callee = axis.product %f0, %f1, %f2, %f3 : !axis.axis_factor, 0>, 2, 180>, !axis.axis_factor, 0>, 2, 90>, !axis.axis_factor, 0>, 2, 45>, !axis.axis_factor, 0>, 3, 15> + %ctx_rep = axis.product %f4, %f5 : !axis.axis_factor, 0>, 3, 5>, !axis.axis_factor, 0>, 5, 1> + + // Same total factor multiset as callee x replicate_over, but heavily permuted. + %ctx_caller = axis.product %f5, %f2, %f4, %f1, %f3, %f0 : !axis.axis_factor, 0>, 5, 1>, !axis.axis_factor, 0>, 2, 45>, !axis.axis_factor, 0>, 3, 5>, !axis.axis_factor, 0>, 2, 90>, !axis.axis_factor, 0>, 3, 15>, !axis.axis_factor, 0>, 2, 180> + + distributed.Function @callee context %ctx_callee : !axis.factor_group<24> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx_caller : !axis.factor_group<360> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + %r = distributed.DistributedCall @callee replicate_over %ctx_rep %arg0 : !axis.factor_group<15>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} diff --git a/test/lit_tests/distributed/call_verify_extra.mlir b/test/lit_tests/distributed/call_verify_extra.mlir new file mode 100644 index 0000000000..95bcc4f896 --- /dev/null +++ b/test/lit_tests/distributed/call_verify_extra.mlir @@ -0,0 +1,56 @@ +// RUN: enzymexlamlir-opt --split-input-file --verify-diagnostics %s + +module { + distributed.MeshComputation @mc0 mesh @mesh0 { + %ax = axis.getaxis tensor<4xf32> 0 + %f0 = axis.factor %ax [4] : !axis.shape_axis, 0> + %ctx = axis.product %f0 : !axis.axis_factor, 0>, 4, 1> + + distributed.Function @callee context %ctx : !axis.factor_group<4> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx : !axis.factor_group<4> arg_types [i32] ret_types [f32] { + ^bb0(%arg0: i32): + // expected-error @+1 {{requires call result #0 to have type 'i32', but got 'f32'}} + %r = distributed.DistributedCall @callee replicate_over %ctx %arg0 : !axis.factor_group<4>, i32 -> f32 + distributed.DistributedYield %r f32 + } + } +} + +// ----- + +module { + distributed.MeshComputation @mc1 mesh @mesh0 { + %ax = axis.getaxis tensor<4xf32> 0 + %f0 = axis.factor %ax [4] : !axis.shape_axis, 0> + %ctx = axis.product %f0 : !axis.axis_factor, 0>, 4, 1> + + distributed.Function @caller context %ctx : !axis.factor_group<4> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + // expected-error @+1 {{references unknown distributed function symbol @missing}} + %r = distributed.DistributedCall @missing replicate_over %ctx %arg0 : !axis.factor_group<4>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} + +// ----- + +module { + %ax = axis.getaxis tensor<4xf32> 0 + %f0 = axis.factor %ax [4] : !axis.shape_axis, 0> + %ctx = axis.product %f0 : !axis.axis_factor, 0>, 4, 1> + + distributed.MeshComputation @mc2 mesh @mesh0 { + distributed.Function @callee context %ctx : !axis.factor_group<4> arg_types [!axis.axis_factor, 0>, 4, 1>] ret_types [!axis.axis_factor, 0>, 4, 1>] { + ^bb0(%arg0: !axis.axis_factor, 0>, 4, 1>): + distributed.DistributedYield %arg0 !axis.axis_factor, 0>, 4, 1> + } + } + + // expected-error @+1 {{must be nested in distributed.function with a FactorGroupType execution_context}} + %r = distributed.DistributedCall @callee replicate_over %ctx %f0 : !axis.factor_group<4>, !axis.axis_factor, 0>, 4, 1> -> !axis.axis_factor, 0>, 4, 1> +} diff --git a/test/lit_tests/distributed/concat.mlir b/test/lit_tests/distributed/concat.mlir deleted file mode 100644 index 6d5b93a24a..0000000000 --- a/test/lit_tests/distributed/concat.mlir +++ /dev/null @@ -1,32 +0,0 @@ -// RUN: enzymexlamlir-opt --shardy-to-distributed %s | FileCheck %s - -distributed.AxisAllToAll @ax1 8 3200000000 0 -distributed.AxisAllToAll @ax2 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1, @ax2] - -sdy.mesh @mesh1 = <["z"=1, "x"=4, "y"=4]> -func.func @main1(%arg0: tensor<20x24x80xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}, %arg1: tensor<20x24x80xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}) -> (tensor<20x24x83xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}) { - %0 = stablehlo.slice %arg1 [0:20, 0:24, 0:3] {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x80xf64>) -> tensor<20x24x3xf64> - %1 = stablehlo.concatenate %arg0, %0, dim = 2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x80xf64>, tensor<20x24x3xf64>) -> tensor<20x24x83xf64> - return %1 : tensor<20x24x83xf64> -} - -// CHECK: module { -// CHECK-NEXT: sdy.mesh @mesh1 = <["z"=1, "x"=4, "y"=4]> -// CHECK-NEXT: func.func @main1(%arg0: tensor<20x24x80xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}, %arg1: tensor<20x24x80xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}) -> (tensor<20x24x83xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}) { -// CHECK-NEXT: %cst = stablehlo.constant dense<0.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.slice %arg1 [0:20, 0:24, 0:3] {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x80xf64>) -> tensor<20x24x3xf64> -// CHECK-NEXT: %1 = stablehlo.pad %arg0, %cst, low = [0, 0, 0], high = [0, 0, 3], interior = [0, 0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x80xf64>, tensor) -> tensor<20x24x83xf64> -// CHECK-NEXT: %2 = stablehlo.pad %0, %cst, low = [0, 0, 80], high = [0, 0, 0], interior = [0, 0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x3xf64>, tensor) -> tensor<20x24x83xf64> -// CHECK-NEXT: %3 = stablehlo.add %1, %2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : tensor<20x24x83xf64> -// CHECK-NEXT: return %3 : tensor<20x24x83xf64> -// CHECK-NEXT: } -// CHECK-NEXT: func.func @main2(%arg0: tensor<20x24x80xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}, %arg1: tensor<20x24x80xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}) -> (tensor<20x24x81xf64> {sdy.sharding = #sdy.sharding<@mesh1, [{"z"}, {"y"}, {"x"}]>}) { -// CHECK-NEXT: %cst = stablehlo.constant dense<0.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.slice %arg1 [0:20, 0:24, 0:1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x80xf64>) -> tensor<20x24x1xf64> -// CHECK-NEXT: %1 = stablehlo.pad %0, %cst, low = [0, 0, 0], high = [0, 0, 80], interior = [0, 0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x1xf64>, tensor) -> tensor<20x24x81xf64> -// CHECK-NEXT: %2 = stablehlo.pad %arg0, %cst, low = [0, 0, 1], high = [0, 0, 0], interior = [0, 0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : (tensor<20x24x80xf64>, tensor) -> tensor<20x24x81xf64> -// CHECK-NEXT: %3 = stablehlo.add %1, %2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh1, [{"z"}, {"y"}, {"x"}]>]>} : tensor<20x24x81xf64> -// CHECK-NEXT: return %3 : tensor<20x24x81xf64> -// CHECK-NEXT: } -// CHECK-NEXT: } diff --git a/test/lit_tests/distributed/diamond_pipeline_friendly.mlir b/test/lit_tests/distributed/diamond_pipeline_friendly.mlir deleted file mode 100644 index 9ec8db421b..0000000000 --- a/test/lit_tests/distributed/diamond_pipeline_friendly.mlir +++ /dev/null @@ -1,31 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline --sdy-insert-explicit-reshards=enable-full-version=true --sdy-reshard-to-collectives --insert-identity-reshard --shardy-to-distributed --localize-distributed-module --cse --distributed-simplify-collectives --distributed-overlap-communication-module %s > /dev/null -// A diamond-shaped, dot-only pipeline-friendly test case. -// -// Shape: -// a1 = a0 * x -// b1 = b0 * a1 -// a2 = a0 * y -// b2 = b1 * a2 -// -// Design intent: -// 1) a0 feeds two branches, creating a true dependency diamond. -// 2) The two branches do different work / use different output sizes. -// 3) Avoid all-reduce and prefer gather-style communication on sharded inputs. - -sdy.mesh @mesh = <["x"=2]> - -distributed.AxisAllToAll @ax1 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1] - -func.func @diamond_pipeline_friendly( - %a0: tensor<1024x2048xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, - %b0: tensor<1024x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, - %x: tensor<2048x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, - %y: tensor<2048x96xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>} -) -> (tensor<64x96xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}) { - %0 = stablehlo.dot %a0, %x : (tensor<1024x2048xf32>, tensor<2048x64xf32>) -> tensor<1024x64xf32> - %1 = stablehlo.dot_general %b0, %0, contracting_dims = [1] x [0] : (tensor<1024x1024xf32>, tensor<1024x64xf32>) -> tensor<1024x64xf32> - %2 = stablehlo.dot %a0, %y : (tensor<1024x2048xf32>, tensor<2048x96xf32>) -> tensor<1024x96xf32> - %3 = stablehlo.dot_general %1, %2, contracting_dims = [0] x [0] : (tensor<1024x64xf32>, tensor<1024x96xf32>) -> tensor<64x96xf32> - return %3 : tensor<64x96xf32> -} diff --git a/test/lit_tests/distributed/dotadd.mlir b/test/lit_tests/distributed/dotadd.mlir deleted file mode 100644 index 947c394746..0000000000 --- a/test/lit_tests/distributed/dotadd.mlir +++ /dev/null @@ -1,23 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline '--sdy-insert-explicit-reshards=enable-full-version=true' --sdy-reshard-to-collectives %s | FileCheck %s -sdy.mesh @mesh = <["x"=4, "y"=2]> -func.func @dotadd(%arg0: tensor<32x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {"y"}]>}, - %arg1: tensor<8x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}, {"x"}]>}, - %arg2: tensor<32x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}, {"x"}]>}) -> (tensor<32x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}, {"x"}]>}) { - %0 = stablehlo.dot %arg0, %arg1 : (tensor<32x8xf32>, tensor<8x64xf32>) -> tensor<32x64xf32> - %1 = stablehlo.add %0, %arg2 : (tensor<32x64xf32>, tensor<32x64xf32>) -> tensor<32x64xf32> - return %1 : tensor<32x64xf32> -} - - -func.func @dotadd2(%arg0: tensor<32x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {"y"}]>}, %arg1: tensor<8x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}, {"x"}]>}, %arg2: tensor<32x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}, {"x"}]>}) -> (tensor<32x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {"y"}]>}) { - %0 = stablehlo.dot %arg0, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}, {"x"}]>]>} : (tensor<32x8xf32>, tensor<8x64xf32>) -> tensor<32x64xf32> - %1 = stablehlo.add %0, %arg2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}, {"x"}]>]>} : tensor<32x64xf32> - return %1 : tensor<32x64xf32> -} - -func.func @dottest(%lhs : tensor<8x32xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"x"}, {"y"}]>}, -%rhs : tensor<32x16xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"y"}, {"x"}]>}) -> (tensor<8x16xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"y"}, {"x"}]>}) { - %0 = stablehlo.dot %lhs, %rhs {sdy.sharding=#sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} - : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32> - return %0 : tensor<8x16xf32> -} \ No newline at end of file diff --git a/test/lit_tests/distributed/function_verify.mlir b/test/lit_tests/distributed/function_verify.mlir new file mode 100644 index 0000000000..867a20f3ac --- /dev/null +++ b/test/lit_tests/distributed/function_verify.mlir @@ -0,0 +1,44 @@ +// RUN: enzymexlamlir-opt --split-input-file --verify-diagnostics %s + +module { + distributed.MeshComputation @mc0 mesh @mesh0 { + %axis = axis.getaxis tensor<8xf32> 0 + %f0 = axis.factor %axis [8] : !axis.shape_axis, 0> + %ctx = axis.product %f0 : !axis.axis_factor, 0>, 8, 1> + + // expected-error @+1 {{requires body block argument #0 to have type 'i32', but got 'f32'}} + distributed.Function @arg_mismatch context %ctx : !axis.factor_group<8> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: f32): + distributed.DistributedYield %arg0 f32 + } + } +} + +// ----- + +module { + distributed.MeshComputation @mc1 mesh @mesh0 { + %axis = axis.getaxis tensor<8xf32> 0 + %f0 = axis.factor %axis [8] : !axis.shape_axis, 0> + %ctx = axis.product %f0 : !axis.axis_factor, 0>, 8, 1> + + // expected-error @+1 {{requires distributed.DistributedYield operand #0 to have type 'i32', but got 'f32'}} + distributed.Function @ret_mismatch context %ctx : !axis.factor_group<8> arg_types [f32] ret_types [i32] { + ^bb0(%arg0: f32): + distributed.DistributedYield %arg0 f32 + } + } +} + +// ----- + +module { + %ctx = "axis.product"() : () -> !axis.factor_group<1> + + // expected-error @+1 {{must be nested in distributed.MeshComputation}} + distributed.Function @bad_parent context %ctx : !axis.factor_group<1> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } +} + diff --git a/test/lit_tests/distributed/innerouterinner_dataparallel.mlir b/test/lit_tests/distributed/innerouterinner_dataparallel.mlir deleted file mode 100644 index 0ec0f2208b..0000000000 --- a/test/lit_tests/distributed/innerouterinner_dataparallel.mlir +++ /dev/null @@ -1,24 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline '--sdy-insert-explicit-reshards=enable-full-version=true' --sdy-reshard-to-collectives %s | FileCheck %s -// A somewhat unrealistic example of an inner-product * outer-product * inner-product which should -// be ideal for pipelining, on the basis that only communicating the inner products is much cheaper -// than communicating the outer products. -// (not entirely unrealistic- maybe there is a low-rank approximation being used). - -sdy.mesh @mesh = <["x"=2]> - -distributed.AxisAllToAll @ax1 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1] - -func.func @innerouterinnerdataparallel( - %x: tensor<512x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, - %w1: tensor<1024x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, - %w2: tensor<2xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}]>}, - %w3: tensor<64x2x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}, {}]>}, - %w4: tensor<64x128xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>} -) -> (tensor<512x128xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) { - %0 = stablehlo.dot %x, %w1 : (tensor<512x1024xf32>, tensor<1024x64xf32>) -> tensor<512x64xf32> - %1 = stablehlo.dot_general %0, %w2, contracting_dims = [] x [] : (tensor<512x64xf32>, tensor<2xf32>) -> tensor<512x64x2xf32> - %2 = stablehlo.dot_general %1, %w3, contracting_dims = [1, 2] x [0, 1] : (tensor<512x64x2xf32>, tensor<64x2x64xf32>) -> tensor<512x64xf32> - %3 = stablehlo.dot %2, %w4 : (tensor<512x64xf32>, tensor<64x128xf32>) -> tensor<512x128xf32> - return %3 : tensor<512x128xf32> -} \ No newline at end of file diff --git a/test/lit_tests/distributed/innerouterinner_operatorparallel.mlir b/test/lit_tests/distributed/innerouterinner_operatorparallel.mlir deleted file mode 100644 index bc30dc9d57..0000000000 --- a/test/lit_tests/distributed/innerouterinner_operatorparallel.mlir +++ /dev/null @@ -1,23 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline --sdy-insert-explicit-reshards=enable-full-version=true --sdy-reshard-to-collectives --insert-identity-reshard --shardy-to-distributed --localize-distributed-module --cse --distributed-simplify-collectives %s | FileCheck %s -// A somewhat unrealistic example of an inner-product * outer-product * inner-product which should -// be ideal for pipelining, on the basis that only communicating the inner products is much cheaper -// than communicating the outer products. -// (not entirely unrealistic- maybe there is a low-rank approximation being used). - -sdy.mesh @mesh = <["x"=2]> -distributed.AxisAllToAll @ax1 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1] - -func.func @innerouterinneroperatorparallel( - %x: tensor<512x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {"x"}]>}, - %w1: tensor<1024x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, - %w2: tensor<2xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}]>}, - %w3: tensor<64x2x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}, {}]>} -) -> (tensor<512x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) { - %0 = stablehlo.dot %x, %w1 : (tensor<512x1024xf32>, tensor<1024x64xf32>) -> tensor<512x64xf32> - %1 = stablehlo.dot_general %0, %w2, contracting_dims = [] x [] : (tensor<512x64xf32>, tensor<2xf32>) -> tensor<512x64x2xf32> - %2 = stablehlo.dot_general %1, %w3, contracting_dims = [1, 2] x [0, 1] : (tensor<512x64x2xf32>, tensor<64x2x64xf32>) -> tensor<512x64xf32> - return %2 : tensor<512x64xf32> -} - -// TODO: filecheck expected output \ No newline at end of file diff --git a/test/lit_tests/distributed/innerouterinner_pipeline_friendly.mlir b/test/lit_tests/distributed/innerouterinner_pipeline_friendly.mlir deleted file mode 100644 index 71598b1513..0000000000 --- a/test/lit_tests/distributed/innerouterinner_pipeline_friendly.mlir +++ /dev/null @@ -1,19 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline --sdy-insert-explicit-reshards=enable-full-version=true --sdy-reshard-to-collectives --insert-identity-reshard --shardy-to-distributed --localize-distributed-module --cse --distributed-simplify-collectives --distributed-overlap-communication-module %s > /dev/null -// A dot-only pipeline-friendly test case. - -sdy.mesh @mesh = <["x"=2]> - -distributed.AxisAllToAll @ax1 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1] - -func.func @innerouterinner_pipeline_friendly( - %x: tensor<2048x4096xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, - %w1: tensor<4096x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, - %w2: tensor<512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}]>}, - %w3: tensor<64x512x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}, {}]>} -) -> (tensor<2048x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) { - %0 = stablehlo.dot %x, %w1 : (tensor<2048x4096xf32>, tensor<4096x64xf32>) -> tensor<2048x64xf32> - %1 = stablehlo.dot_general %0, %w2, contracting_dims = [] x [] : (tensor<2048x64xf32>, tensor<512xf32>) -> tensor<2048x64x512xf32> - %2 = stablehlo.dot_general %1, %w3, contracting_dims = [1, 2] x [0, 1] : (tensor<2048x64x512xf32>, tensor<64x512x256xf32>) -> tensor<2048x256xf32> - return %2 : tensor<2048x256xf32> -} diff --git a/test/lit_tests/distributed/mesh_verify.mlir b/test/lit_tests/distributed/mesh_verify.mlir new file mode 100644 index 0000000000..1eaef6cf4d --- /dev/null +++ b/test/lit_tests/distributed/mesh_verify.mlir @@ -0,0 +1,30 @@ +// RUN: enzymexlamlir-opt --split-input-file --verify-diagnostics %s + +// Valid metadata-only mesh body. +module { + distributed.MeshComputation @ok mesh @mesh0 { + ^bb0: + } +} + +// ----- + +// Non-metadata, non-function op in mesh body should fail verification. +module { + // expected-error @+1 {{only distributed.Function and static metadata ops are allowed in the mesh body; operation 'arith.constant' is neither}} + distributed.MeshComputation @bad mesh @mesh0 { + ^bb0: + %axis = axis.getaxis tensor<4xf32> 0 + %f0 = axis.factor %axis [4] : !axis.shape_axis, 0> + %ctx = axis.product %f0 : !axis.axis_factor, 0>, 4, 1> + %c0 = arith.constant 0 : i32 + } +} + +// ----- + +// GetPhysicalAxis verifier should reject unknown symbol references. +module { + // expected-error @+1 {{references unknown physical axis symbol @missing_axis}} + %a = distributed.GetPhysicalAxis @missing_axis : !distributed.physical_comm_axis<4> +} diff --git a/test/lit_tests/distributed/roundtrip.mlir b/test/lit_tests/distributed/roundtrip.mlir new file mode 100644 index 0000000000..554975938b --- /dev/null +++ b/test/lit_tests/distributed/roundtrip.mlir @@ -0,0 +1,40 @@ +// RUN: enzymexlamlir-opt --split-input-file %s | FileCheck %s + +// Roundtrip distributed.Function / distributed.DistributedCall / +// distributed.DistributedYield. +module { + distributed.MeshComputation @mc mesh @mesh0 { + %axis = axis.getaxis tensor<12xf32> 0 + %f0, %f1, %f2 = axis.factor %axis [2, 2, 3] : !axis.shape_axis, 0> + %ctx_callee = axis.product %f0, %f1 : !axis.axis_factor, 0>, 2, 6>, !axis.axis_factor, 0>, 2, 3> + %ctx_rep = axis.product %f2 : !axis.axis_factor, 0>, 3, 1> + %ctx_caller = axis.product %f0, %f1, %f2 : !axis.axis_factor, 0>, 2, 6>, !axis.axis_factor, 0>, 2, 3>, !axis.axis_factor, 0>, 3, 1> + + distributed.Function @callee context %ctx_callee : !axis.factor_group<4> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + distributed.DistributedYield %arg0 i32 + } + + distributed.Function @caller context %ctx_caller : !axis.factor_group<12> arg_types [i32] ret_types [i32] { + ^bb0(%arg0: i32): + %r = distributed.DistributedCall @callee replicate_over %ctx_rep %arg0 : !axis.factor_group<3>, i32 -> i32 + distributed.DistributedYield %r i32 + } + } +} + +// CHECK-LABEL: module { +// CHECK: distributed.Function @callee context +// CHECK: distributed.Function @caller context +// CHECK: distributed.DistributedCall @callee replicate_over + +// ----- + +// Roundtrip distributed.MeshComputation with a minimal metadata-only body. +module { + distributed.MeshComputation @mc mesh @mesh0 { + %a = axis.getaxis tensor<4xf32> 0 + } +} + +// CHECK: distributed.MeshComputation @mc mesh @mesh0 diff --git a/test/lit_tests/distributed/test_timing.mlir b/test/lit_tests/distributed/test_timing.mlir deleted file mode 100644 index 258f8eb528..0000000000 --- a/test/lit_tests/distributed/test_timing.mlir +++ /dev/null @@ -1,47 +0,0 @@ -module { - sdy.mesh @mesh = <["x"=2]> - distributed.AxisAllToAll @ax1 2 1600000000 0 - distributed.PhysicalMesh @phys_mesh [@ax1] - func.func @innerouterinneroperatorparallel(%arg0: tensor<512x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {"x"}]>}, %arg1: tensor<1024x64xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, %arg2: tensor<1x128xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, %arg3: tensor<8192x10xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) -> (tensor<512x10xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) { - %0 = distributed.AxisFactor @ax1 [2 : i32] - %1 = distributed.LogicalMesh %0 {enzyme.shardy_axis_names = ["x"], physical_mesh = @phys_mesh} - %2 = distributed.MeshComputation %1 %arg0, %arg1, %arg2, %arg3 tensor<512x1024xf32>, tensor<1024x64xf32>, tensor<1x128xf32>, tensor<8192x10xf32> -> tensor<512x10xf32> { - ^bb0(%arg4: tensor<512x1024xf32>, %arg5: tensor<1024x64xf32>, %arg6: tensor<1x128xf32>, %arg7: tensor<8192x10xf32>): - %3 = distributed.LogicalMesh %0 {physical_mesh = @phys_mesh} - %4 = "distributed.Sharding"() : () -> !distributed.sharding - %5 = "distributed.Sharding"(%0) : (!distributed.logical_comm_axis) -> !distributed.sharding - %6 = distributed.Collective %1(tensor<512x64xf32>> tensor<512x64xf32>, %4, %4) -> (tensor<512x64xf32>> tensor<512x64xf32>, %5, %4) - %8 = distributed.Collective %1(tensor<8192x10xf32>> tensor<8192x10xf32>, %5, %4) -> (tensor<8192x10xf32>> tensor<8192x10xf32>, %4, %4) - %10 = distributed.RegionComputation %3 %arg4, %arg5, %arg6, %arg7 tensor<512x1024xf32>, tensor<1024x64xf32>, tensor<1x128xf32>, tensor<8192x10xf32> -> tensor<512x10xf32> { - ^bb0(%arg8: tensor<512x1024xf32>, %arg9: tensor<1024x64xf32>, %arg10: tensor<1x128xf32>, %arg11: tensor<8192x10xf32>): - %11 = stablehlo.dot %arg8, %arg9 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>} : (tensor<512x1024xf32>, tensor<1024x64xf32>) -> tensor<512x64xf32> - %12 = sdy.all_reduce {"x"} %11 out_sharding=<@mesh, [{}, {}]> : tensor<512x64xf32> - %13 = distributed.SendRecv %6 %12 tensor<512x64xf32> -> tensor<512x64xf32> {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} - %14 = stablehlo.reshape %13 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<512x64xf32>) -> tensor<32768x1xf32> - %15 = stablehlo.dot %14, %arg10 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<32768x1xf32>, tensor<1x128xf32>) -> tensor<32768x128xf32> - %16 = stablehlo.reshape %15 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<32768x128xf32>) -> tensor<512x8192xf32> - // some dummy ops to make the timing more interesting - %116 = stablehlo.cosine %16 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<512x8192xf32>) -> tensor<512x8192xf32> - %17 = distributed.SendRecv %8 %arg11 tensor<8192x10xf32> -> tensor<8192x10xf32> {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>} - %18 = stablehlo.dot %116, %17 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<512x8192xf32>, tensor<8192x10xf32>) -> tensor<512x10xf32> - distributed.DistributedYield %18 tensor<512x10xf32> - }, { - ^bb0(%arg8: tensor<512x1024xf32>, %arg9: tensor<1024x64xf32>, %arg10: tensor<1x128xf32>, %arg11: tensor<8192x10xf32>): - %11 = stablehlo.dot %arg8, %arg9 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>} : (tensor<512x1024xf32>, tensor<1024x64xf32>) -> tensor<512x64xf32> - // some dummy ops to make the timing more interesting - %111 = stablehlo.add %11, %11 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>} : (tensor<512x64xf32>, tensor<512x64xf32>) -> tensor<512x64xf32> - %211 = stablehlo.abs %111 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>} : (tensor<512x64xf32>) -> tensor<512x64xf32> - %12 = sdy.all_reduce {"x"} %211 out_sharding=<@mesh, [{}, {}]> : tensor<512x64xf32> - %13 = distributed.SendRecv %6 %12 tensor<512x64xf32> -> tensor<512x64xf32> {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} - %14 = stablehlo.reshape %13 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<512x64xf32>) -> tensor<32768x1xf32> - %15 = stablehlo.dot %14, %arg10 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<32768x1xf32>, tensor<1x128xf32>) -> tensor<32768x128xf32> - %16 = stablehlo.reshape %15 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<32768x128xf32>) -> tensor<512x8192xf32> - %17 = distributed.SendRecv %8 %arg11 tensor<8192x10xf32> -> tensor<8192x10xf32> {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>} - %18 = stablehlo.dot %16, %17 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}]>]>} : (tensor<512x8192xf32>, tensor<8192x10xf32>) -> tensor<512x10xf32> - distributed.DistributedYield %18 tensor<512x10xf32> - } - distributed.DistributedYield %10 tensor<512x10xf32> - } - return %2 : tensor<512x10xf32> - } -} \ No newline at end of file diff --git a/test/lit_tests/distributed/two_matmul.mlir b/test/lit_tests/distributed/two_matmul.mlir deleted file mode 100644 index 3d70086990..0000000000 --- a/test/lit_tests/distributed/two_matmul.mlir +++ /dev/null @@ -1,18 +0,0 @@ -// RUN: enzymexlamlir-opt --sdy-propagation-pipeline --sdy-insert-explicit-reshards=enable-full-version=true --sdy-reshard-to-collectives --insert-identity-reshard --shardy-to-distributed --localize-distributed-module --cse --distributed-simplify-collectives --distributed-overlap-communication-module --distributed-sink-recvs-module %s > /dev/null -// Two chained square matmuls: t1 = A @ B, t2 = t1 @ C. -// A is row-sharded; B and C are replicated so no all-reduce arises. - -sdy.mesh @mesh = <["x"=2]> - -distributed.AxisAllToAll @ax1 2 1600000000 0 -distributed.PhysicalMesh @phys_mesh [@ax1] - -func.func @two_matmul( - %a: tensor<1024x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}, - %b: tensor<1024x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, - %c: tensor<1024x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>} -) -> (tensor<1024x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) { - %t1 = stablehlo.dot %a, %b : (tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32> - %t2 = stablehlo.dot %t1, %c : (tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32> - return %t2 : tensor<1024x1024xf32> -} From b5f1ced5fe9f413644fcb34bfb38e75e473aeaae Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 03/11] Fragile shardy discovery --- src/enzyme_ad/jax/BUILD | 2 +- .../FindShardyFunctionsAnalysis.cpp | 255 ++++++++++++++++++ .../Distributed/FindShardyFunctionsAnalysis.h | 42 +++ .../jax/Passes/Distributed/Passes.cpp | 6 - src/enzyme_ad/jax/Passes/Distributed/Passes.h | 6 + .../jax/Passes/Distributed/Passes.td | 12 +- .../print_shardy_function_names.mlir | 67 +++++ 7 files changed, 380 insertions(+), 10 deletions(-) create mode 100644 src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.cpp create mode 100644 src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.h create mode 100644 test/lit_tests/distributed/print_shardy_function_names.mlir diff --git a/src/enzyme_ad/jax/BUILD b/src/enzyme_ad/jax/BUILD index db469a1209..c820a6b179 100644 --- a/src/enzyme_ad/jax/BUILD +++ b/src/enzyme_ad/jax/BUILD @@ -954,7 +954,7 @@ gentbl_cc_library( td_library( name = "DistributedPassesTdFiles", - srcs = [], + srcs = ["Passes/Distributed/Passes.td"], deps = [ "@llvm-project//mlir:PassBaseTdFiles", ], diff --git a/src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.cpp b/src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.cpp new file mode 100644 index 0000000000..b717dd8ffc --- /dev/null +++ b/src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.cpp @@ -0,0 +1,255 @@ +#include "src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.h" +#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h" + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LogicalResult.h" +#include "shardy/dialect/sdy/ir/dialect.h" +#include "shardy/dialect/sdy/ir/utils.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/raw_ostream.h" + +namespace mlir::enzyme::distributed { +#define GEN_PASS_DEF_PRINTSHARDYFUNCTIONNAMESPASS +#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h.inc" + +namespace { + +static bool hasSdyPrefixedAttr(NamedAttribute attr) { + return attr.getName().strref().starts_with("sdy."); +} + +static bool hasSdyAttrs(Operation *op) { + for (NamedAttribute attr : op->getAttrs()) { + if (hasSdyPrefixedAttr(attr)) { + return true; + } + } + return false; +} + +static std::string stringifyAttribute(Attribute attr) { + std::string printed; + llvm::raw_string_ostream os(printed); + os << attr; + return printed; +} + +static void insertUniqueMesh(llvm::SmallVectorImpl &out, + sdy::MeshAttr mesh) { + if (llvm::is_contained(out, mesh)) { + return; + } + out.push_back(std::move(mesh)); +} + +static LogicalResult +collectMeshFromSharding(sdy::TensorShardingAttr sharding, func::FuncOp funcOp, + const SymbolTable &symbolTable, + llvm::SmallVectorImpl &out) { + sdy::MeshAttr mesh = sharding.getMesh(symbolTable); + if (!mesh) { + return funcOp.emitError() + << "failed to resolve Shardy mesh for analysis from " + << sharding.getMeshOrRef(); + } + insertUniqueMesh(out, mesh); + return success(); +} + +static LogicalResult +collectMeshesFromAttr(Attribute attr, func::FuncOp funcOp, + const SymbolTable &symbolTable, + llvm::SmallVectorImpl &out) { + if (auto sharding = dyn_cast_or_null(attr)) { + return collectMeshFromSharding(sharding, funcOp, symbolTable, out); + } + if (auto shardings = + dyn_cast_or_null(attr)) { + for (sdy::TensorShardingAttr sharding : shardings.getShardings()) { + if (failed(collectMeshFromSharding(sharding, funcOp, symbolTable, out))) { + return failure(); + } + } + } + return success(); +} + +static FailureOr> +getShardyFunctionMeshes(func::FuncOp funcOp, const SymbolTable &symbolTable) { + llvm::SmallVector meshes; + + for (NamedAttribute attr : funcOp->getAttrs()) { + if (failed(collectMeshesFromAttr(attr.getValue(), funcOp, symbolTable, + meshes))) { + return failure(); + } + } + + if (ArrayAttr argAttrs = funcOp.getArgAttrsAttr()) { + for (Attribute attrsAttr : argAttrs) { + auto attrs = dyn_cast_or_null(attrsAttr); + if (!attrs) { + continue; + } + for (NamedAttribute attr : attrs) { + if (failed(collectMeshesFromAttr(attr.getValue(), funcOp, symbolTable, + meshes))) { + return failure(); + } + } + } + } + + if (ArrayAttr resultAttrs = funcOp.getResAttrsAttr()) { + for (Attribute attrsAttr : resultAttrs) { + auto attrs = dyn_cast_or_null(attrsAttr); + if (!attrs) { + continue; + } + for (NamedAttribute attr : attrs) { + if (failed(collectMeshesFromAttr(attr.getValue(), funcOp, symbolTable, + meshes))) { + return failure(); + } + } + } + } + + WalkResult walkResult = funcOp.walk([&](Operation *op) { + for (NamedAttribute attr : op->getAttrs()) { + if (failed(collectMeshesFromAttr(attr.getValue(), funcOp, symbolTable, + meshes))) { + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + if (walkResult.wasInterrupted()) { + return failure(); + } + + return meshes; +} + +static bool hasShardyFunctionInfo(func::FuncOp funcOp) { + if (hasSdyAttrs(funcOp)) { + return true; + } + + if (ArrayAttr argAttrs = funcOp.getArgAttrsAttr()) { + for (Attribute attrsAttr : argAttrs) { + auto attrs = dyn_cast_or_null(attrsAttr); + if (!attrs) { + continue; + } + for (NamedAttribute attr : attrs) { + if (hasSdyPrefixedAttr(attr)) { + return true; + } + } + } + } + + if (ArrayAttr resultAttrs = funcOp.getResAttrsAttr()) { + for (Attribute attrsAttr : resultAttrs) { + auto attrs = dyn_cast_or_null(attrsAttr); + if (!attrs) { + continue; + } + for (NamedAttribute attr : attrs) { + if (hasSdyPrefixedAttr(attr)) { + return true; + } + } + } + } + + bool foundShardyBodyIR = false; + funcOp.walk([&](Operation *op) { + if (op == funcOp.getOperation()) { + return WalkResult::advance(); + } + if (op->getName().getStringRef().starts_with("sdy.") || hasSdyAttrs(op)) { + foundShardyBodyIR = true; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + if (foundShardyBodyIR) { + return true; + } + + return false; +} + +} // namespace + +// Fragile, but works for now. +FindShardyFunctionsAnalysis::FindShardyFunctionsAnalysis(ModuleOp module) { + SymbolTable symbolTable(module); + + for (func::FuncOp funcOp : module.getOps()) { + if (!hasShardyFunctionInfo(funcOp)) { + continue; + } + + FailureOr> meshes = + getShardyFunctionMeshes(funcOp, symbolTable); + if (failed(meshes)) { + valid = false; + return; + } + + functionIndices[funcOp.getOperation()] = functions.size(); + functions.push_back(FindShardyFunctionsAnalysis::FunctionInfo{ + funcOp, funcOp.getSymNameAttr(), std::move(*meshes)}); + } +} + +llvm::ArrayRef +FindShardyFunctionsAnalysis::getMeshes(func::FuncOp funcOp) const { + auto it = functionIndices.find(funcOp.getOperation()); + if (it == functionIndices.end()) { + return {}; + } + return functions[it->second].meshes; +} + +namespace { + +// Testing / Debugging purposes. +struct PrintShardyFunctionNamesPass + : public impl::PrintShardyFunctionNamesPassBase< + PrintShardyFunctionNamesPass> { + using PrintShardyFunctionNamesPassBase::PrintShardyFunctionNamesPassBase; + + void runOnOperation() override { + const FindShardyFunctionsAnalysis &analysis = + getAnalysis(); + if (!analysis.isValid()) { + signalPassFailure(); + return; + } + + for (const FindShardyFunctionsAnalysis::FunctionInfo &info : + analysis.getShardyFunctions()) { + llvm::outs() << info.symName.getValue(); + if (!info.meshes.empty()) { + llvm::outs() << ": "; + llvm::interleaveComma(info.meshes, llvm::outs(), + [&](sdy::MeshAttr mesh) { + llvm::outs() << stringifyAttribute(mesh); + }); + } + llvm::outs() << "\n"; + } + } +}; + +} // namespace +} // namespace mlir::enzyme::distributed diff --git a/src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.h b/src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.h new file mode 100644 index 0000000000..8f13df2350 --- /dev/null +++ b/src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.h @@ -0,0 +1,42 @@ +#ifndef ENZYME_AD_JAX_PASSES_DISTRIBUTED_FINDSHARDYFUNCTIONSANALYSIS_H +#define ENZYME_AD_JAX_PASSES_DISTRIBUTED_FINDSHARDYFUNCTIONSANALYSIS_H + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "shardy/dialect/sdy/ir/dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir::enzyme::distributed { + +class FindShardyFunctionsAnalysis { +public: + struct FunctionInfo { + func::FuncOp funcOp; + StringAttr symName; + llvm::SmallVector meshes; + }; + + explicit FindShardyFunctionsAnalysis(ModuleOp module); + + llvm::ArrayRef getShardyFunctions() const { return functions; } + + bool isShardyFunction(func::FuncOp funcOp) const { + return functionIndices.contains(funcOp.getOperation()); + } + + llvm::ArrayRef getMeshes(func::FuncOp funcOp) const; + + bool isValid() const { return valid; } + +private: + bool valid = true; + llvm::SmallVector functions; + llvm::DenseMap functionIndices; +}; + +} // namespace mlir::enzyme::distributed + +#endif diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp b/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp index 05c82777d3..1fcfa92552 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp @@ -1,7 +1 @@ #include "src/enzyme_ad/jax/Passes/Distributed/Passes.h" - -namespace mlir::enzyme::distributed { - -void registerdistributedPasses() {} - -} // namespace mlir::enzyme::distributed diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.h b/src/enzyme_ad/jax/Passes/Distributed/Passes.h index aad045415c..ae8661f140 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.h +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.h @@ -7,6 +7,12 @@ namespace mlir { namespace enzyme { namespace distributed { +#define GEN_PASS_DECL +#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h.inc" + +#define GEN_PASS_REGISTRATION +#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h.inc" + void registerdistributedPasses(); } // namespace distributed diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.td b/src/enzyme_ad/jax/Passes/Distributed/Passes.td index 96c8254b61..6327ff9fb2 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.td +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.td @@ -3,8 +3,14 @@ include "mlir/Pass/PassBase.td" -// Distributed pass definitions intentionally removed during the major -// refactor. Keep this file and the surrounding build targets in place so new -// pass defs can be added back without recreating the infrastructure. +def PrintShardyFunctionNamesPass + : Pass<"print-shardy-function-names", "::mlir::ModuleOp"> { + let summary = "Print names of functions that carry Shardy IR"; + let description = [{ + Scans func.func operations and prints the symbol names of functions that + either carry sdy-prefixed attributes on the function/signature or contain + sdy dialect operations/attributes in their bodies. + }]; +} #endif // ENZYME_AD_JAX_DISTRIBUTED_PASSES diff --git a/test/lit_tests/distributed/print_shardy_function_names.mlir b/test/lit_tests/distributed/print_shardy_function_names.mlir new file mode 100644 index 0000000000..0e4e973238 --- /dev/null +++ b/test/lit_tests/distributed/print_shardy_function_names.mlir @@ -0,0 +1,67 @@ +// RUN: enzymexlamlir-opt --print-shardy-function-names -o /dev/null %s | FileCheck %s + +module { + sdy.mesh @mesh = <["x"=2]> + sdy.mesh @mesh_ab = <["a"=4, "b"=2]> + sdy.mesh @mesh_xyz = <["x"=2, "y"=2, "z"=2]> + + func.func @plain(%arg0: tensor<4xf32>) -> tensor<4xf32> { + return %arg0 : tensor<4xf32> + } + + func.func @sharded_arg(%arg0: tensor<4xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> tensor<4xf32> { + return %arg0 : tensor<4xf32> + } + + func.func @sharded_arg_2(%arg0: tensor<4xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> tensor<4xf32> { + return %arg0 : tensor<4xf32> + } + + func.func @body_only(%arg0: tensor<4xf32>) -> tensor<4xf32> { + %0 = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}]>] >} : tensor<4xf32> + return %0 : tensor<4xf32> + } + + func.func @multi_mesh_signature( + %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_ab, [{"a"}, {"b"}]>}) + -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_xyz, [{"x"}, {"z", "y"}]>}) { + return %arg0 : tensor<8x8xf32> + } + + func.func @multi_mesh_body(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { + %0 = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_ab, [{"a"}, {"b"}]>]>} : tensor<8x8xf32> + %1 = stablehlo.negate %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_xyz, [{"x"}, {"z", "y"}]>]>} : tensor<8x8xf32> + return %1 : tensor<8x8xf32> + } + + func.func @multi_mesh_mixed( + %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) + -> tensor<8x8xf32> { + %0 = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_ab, [{"a"}, {"b"}]>]>} : tensor<8x8xf32> + %1 = stablehlo.negate %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_xyz, [{"x"}, {"z", "y"}]>]>} : tensor<8x8xf32> + return %1 : tensor<8x8xf32> + } + + func.func @anonymous_mesh_signature( + %arg0: tensor<4xf32> {sdy.sharding = #sdy.sharding, [{"u"}]>}) + -> (tensor<4xf32> {sdy.sharding = #sdy.sharding, [{"u"}]>}) { + return %arg0 : tensor<4xf32> + } + + func.func @named_and_anonymous_same_mesh( + %arg0: tensor<4xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) + -> tensor<4xf32> { + %0 = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[, [{"x"}]>]>} : tensor<4xf32> + return %0 : tensor<4xf32> + } +} + +// CHECK: sharded_arg: #sdy.mesh<["x"=2]> +// CHECK: sharded_arg_2: #sdy.mesh<["x"=2]> +// CHECK: body_only: #sdy.mesh<["x"=2]> +// CHECK: multi_mesh_signature: #sdy.mesh<["a"=4, "b"=2]>, #sdy.mesh<["x"=2, "y"=2, "z"=2]> +// CHECK: multi_mesh_body: #sdy.mesh<["a"=4, "b"=2]>, #sdy.mesh<["x"=2, "y"=2, "z"=2]> +// CHECK: multi_mesh_mixed: #sdy.mesh<["x"=2]>, #sdy.mesh<["a"=4, "b"=2]>, #sdy.mesh<["x"=2, "y"=2, "z"=2]> +// CHECK: anonymous_mesh_signature: #sdy.mesh<["u"=2]> +// CHECK: named_and_anonymous_same_mesh: #sdy.mesh<["x"=2]> +// CHECK-NOT: plain From bf458b9afe2bcf1062c7f612aa7dfa4034f1edc7 Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 04/11] Physical mesh op --- .../jax/Dialect/Distributed/AxisOps.cpp | 69 +++++++++++++++---- src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 23 +++++-- .../jax/Dialect/Distributed/Types.td | 10 --- .../jax/Dialect/Distributed/Utilities.h | 21 +++--- test/lit_tests/distributed/roundtrip.mlir | 11 +++ 5 files changed, 97 insertions(+), 37 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp index f497e333f3..3fda779626 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp @@ -3,25 +3,66 @@ namespace mlir::enzyme::distributed { -LogicalResult GetPhysicalAxisOp::verify() { - FailureOr physicalAxis = - resolveSymbolOpFromAttr( - *this, getPhysicalAxisAttr()); - if (failed(physicalAxis)) { - return emitOpError() << "references unknown physical axis symbol " - << getPhysicalAxisAttr(); +LogicalResult PhysicalMeshOp::verify() { + for (auto [idx, axisAttr] : llvm::enumerate(getAxesAttr())) { + auto typeAttr = dyn_cast(axisAttr); + if (!typeAttr) { + return emitOpError() << "requires axes[" << idx << "] to be a TypeAttr"; + } + if (!isa(typeAttr.getValue())) { + return emitOpError() << "requires axes[" << idx + << "] to be a PhysicalCommAxisType attribute"; + } } - unsigned expectedExtent = - static_cast(physicalAxis->getPhysicalAxisExtent()); - auto axisType = getAxis().getType(); - if (axisType.getExtent() != expectedExtent) { - return emitOpError() << "requires result type extent to match referenced " - "physical axis extent (" - << axisType.getExtent() << " != " << expectedExtent + return success(); +} + +LogicalResult GetPhysicalMeshAxesOp::verify() { + FailureOr physicalMesh = + resolveSymbolOpFromAttr(*this, getPhysicalMeshAttr()); + if (failed(physicalMesh)) { + return emitOpError() << "references unknown physical mesh symbol " + << getPhysicalMeshAttr(); + } + + ArrayAttr axisAttrs = physicalMesh->getAxesAttr(); + if (getAxes().size() != axisAttrs.size()) { + return emitOpError() << "requires result count to match referenced " + "physical mesh axes size (" + << getAxes().size() << " != " << axisAttrs.size() << ")"; } + for (auto [idx, axisAttr] : llvm::enumerate(axisAttrs)) { + auto typeAttr = dyn_cast(axisAttr); + if (!typeAttr) { + return emitOpError() + << "requires referenced physical mesh axes[" << idx + << "] to be a TypeAttr"; + } + + auto expectedAxisType = dyn_cast(typeAttr.getValue()); + if (!expectedAxisType) { + return emitOpError() + << "requires referenced physical mesh axes[" << idx + << "] to be a PhysicalCommAxisType attribute"; + } + + auto actualAxisType = + dyn_cast(getAxes()[idx].getType()); + if (!actualAxisType) { + return emitOpError() << "requires result #" << idx + << " to have PhysicalCommAxisType"; + } + + if (actualAxisType != expectedAxisType) { + return emitOpError() << "requires result #" << idx << " to have type " + << expectedAxisType << ", but got " + << actualAxisType; + } + } + return success(); } diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td index 4220b6e408..8a338dc961 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -21,13 +21,26 @@ defvar ShardingAttribute = Sdy_TensorSharding; // Marker for distributed metadata ops. Includes builtin Pure. def PureDistributedMetadata : TraitList<[MetadataOpTrait]>; -def GetPhysicalAxisOp : DistributedOp<"GetPhysicalAxis", [PureDistributedMetadata]> { +def PhysicalMeshOp : DistributedOp<"PhysicalMesh", [PureDistributedMetadata, Symbol]> { let description = [{ - Resolves a module-level physical axis symbol into an SSA axis value. + Declares a mesh of device targets along with their local communication hierarchy. + }]; + let arguments = (ins + SymbolNameAttr:$sym_name, + StrAttr:$device_target, + ArrayAttr:$axes // array of physical axis types. + ); + let assemblyFormat = "$sym_name `device_target` $device_target `axes` $axes attr-dict"; + let hasVerifier = 1; +} + +def GetPhysicalMeshAxesOp : DistributedOp<"GetPhysicalMeshAxes", [PureDistributedMetadata]> { + let description = [{ + Resolves a module-level physical mesh symbol into one SSA axis value per mesh axis. }]; - let arguments = (ins FlatSymbolRefAttr:$physical_axis); - let results = (outs PhysicalCommAxisType:$axis); - let assemblyFormat = "$physical_axis attr-dict `:` type($axis)"; + let arguments = (ins FlatSymbolRefAttr:$physical_mesh); + let results = (outs Variadic:$axes); + let assemblyFormat = "$physical_mesh attr-dict `:` type($axes)"; let hasVerifier = 1; } diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Types.td b/src/enzyme_ad/jax/Dialect/Distributed/Types.td index ea3d64cb11..1b7ab395cb 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Types.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Types.td @@ -28,13 +28,6 @@ def ReplicationAxisType : DistributedAxisType<"ReplicationAxis", "replication_ax let assemblyFormat = "`<` $extent `>`"; } - -// Small (unecessary?) type wrapping products of logical axis and "reduce" for sharding annotations. -def ShardingAxisType : DistributedType<"Sharding", "sharding">; - -def PhysicalMeshType : DistributedType<"PhysicalMesh", "physical_mesh">; -def LogicalMeshType : DistributedType<"LogicalMesh", "logical_mesh">; - // Wraps a value type that may be synchronized at a later point. def AsynchHandleType : DistributedType<"AsynchHandle", "asynch_handle"> { let summary = "Handle to a value that can be synchronized later"; @@ -42,7 +35,4 @@ def AsynchHandleType : DistributedType<"AsynchHandle", "asynch_handle"> { let assemblyFormat = "`<` $valueType `>`"; } -// Single token type used for collective declaration and message events. -def MessageTokenType : DistributedType<"MessageToken", "message_token">; - #endif // ENZYME_DISTRIBUTED_DIALECT_TYPES_H \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h index 75539ef111..4a926dbdbe 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h @@ -13,15 +13,20 @@ ::mlir::FailureOr resolveSymbolOpFromAttr(::mlir::Operation *from, if (!symRef) { return ::mlir::failure(); } - auto *op = ::mlir::SymbolTable::lookupNearestSymbolFrom(from, symRef); - if (!op) { - return ::mlir::failure(); - } - auto typedOp = llvm::dyn_cast(op); - if (!typedOp) { - return ::mlir::failure(); + + for (auto *scope = from; scope; scope = scope->getParentOp()) { + if (!scope->hasTrait<::mlir::OpTrait::SymbolTable>()) { + continue; + } + if (auto *op = ::mlir::SymbolTable::lookupSymbolIn(scope, symRef)) { + if (auto typedOp = llvm::dyn_cast(op)) { + return typedOp; + } + return ::mlir::failure(); + } } - return typedOp; + + return ::mlir::failure(); } // Returns the execution-context FactorGroup value from the nearest enclosing diff --git a/test/lit_tests/distributed/roundtrip.mlir b/test/lit_tests/distributed/roundtrip.mlir index 554975938b..d538782e23 100644 --- a/test/lit_tests/distributed/roundtrip.mlir +++ b/test/lit_tests/distributed/roundtrip.mlir @@ -3,7 +3,11 @@ // Roundtrip distributed.Function / distributed.DistributedCall / // distributed.DistributedYield. module { + distributed.PhysicalMesh @mesh0 device_target "cpu" axes [!distributed.physical_comm_axis<2>, !distributed.physical_comm_axis<3>] + distributed.MeshComputation @mc mesh @mesh0 { + %p0, %p1 = distributed.GetPhysicalMeshAxes @mesh0 : !distributed.physical_comm_axis<2>, !distributed.physical_comm_axis<3> + %axis = axis.getaxis tensor<12xf32> 0 %f0, %f1, %f2 = axis.factor %axis [2, 2, 3] : !axis.shape_axis, 0> %ctx_callee = axis.product %f0, %f1 : !axis.axis_factor, 0>, 2, 6>, !axis.axis_factor, 0>, 2, 3> @@ -24,6 +28,8 @@ module { } // CHECK-LABEL: module { +// CHECK: distributed.PhysicalMesh @mesh0 device_target "cpu" axes [!distributed.physical_comm_axis<2>, !distributed.physical_comm_axis<3>] +// CHECK: %{{.*}}:2 = distributed.GetPhysicalMeshAxes @mesh0 : !distributed.physical_comm_axis<2>, !distributed.physical_comm_axis<3> // CHECK: distributed.Function @callee context // CHECK: distributed.Function @caller context // CHECK: distributed.DistributedCall @callee replicate_over @@ -32,9 +38,14 @@ module { // Roundtrip distributed.MeshComputation with a minimal metadata-only body. module { + distributed.PhysicalMesh @mesh0 device_target "cpu" axes [!distributed.physical_comm_axis<4>] + distributed.MeshComputation @mc mesh @mesh0 { + %p0 = distributed.GetPhysicalMeshAxes @mesh0 : !distributed.physical_comm_axis<4> %a = axis.getaxis tensor<4xf32> 0 } } +// CHECK: distributed.PhysicalMesh @mesh0 device_target "cpu" axes [!distributed.physical_comm_axis<4>] // CHECK: distributed.MeshComputation @mc mesh @mesh0 +// CHECK: %{{.*}} = distributed.GetPhysicalMeshAxes @mesh0 : !distributed.physical_comm_axis<4> From b84c49e7b31e2b723005f80965e4fe3ff726dc69 Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 05/11] WIP: lowering tests --- .../jax/Dialect/Distributed/AxisOps.cpp | 66 ++++++++- src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 13 ++ .../jax/Passes/Distributed/Passes.cpp | 1 - src/enzyme_ad/jax/Passes/Distributed/Passes.h | 2 - .../jax/Passes/Distributed/Passes.td | 19 +++ .../Distributed/ShardyToDistributed.cpp | 140 ++++++++++++++++++ ...ardy_transformer_block_partial_export.mlir | 48 ++++++ .../shardy_transformer_block_pre_export.mlir | 117 +++++++++++++++ .../shardy_transformer_exported.mlir | 121 +++++++++++++++ 9 files changed, 518 insertions(+), 9 deletions(-) delete mode 100644 src/enzyme_ad/jax/Passes/Distributed/Passes.cpp create mode 100644 src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp create mode 100644 test/lit_tests/distributed/shardy_transformer_block_partial_export.mlir create mode 100644 test/lit_tests/distributed/shardy_transformer_block_pre_export.mlir create mode 100644 test/lit_tests/distributed/shardy_transformer_exported.mlir diff --git a/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp index 3fda779626..cc0e82ea5e 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp @@ -37,16 +37,14 @@ LogicalResult GetPhysicalMeshAxesOp::verify() { for (auto [idx, axisAttr] : llvm::enumerate(axisAttrs)) { auto typeAttr = dyn_cast(axisAttr); if (!typeAttr) { - return emitOpError() - << "requires referenced physical mesh axes[" << idx - << "] to be a TypeAttr"; + return emitOpError() << "requires referenced physical mesh axes[" << idx + << "] to be a TypeAttr"; } auto expectedAxisType = dyn_cast(typeAttr.getValue()); if (!expectedAxisType) { - return emitOpError() - << "requires referenced physical mesh axes[" << idx - << "] to be a PhysicalCommAxisType attribute"; + return emitOpError() << "requires referenced physical mesh axes[" << idx + << "] to be a PhysicalCommAxisType attribute"; } auto actualAxisType = @@ -66,4 +64,60 @@ LogicalResult GetPhysicalMeshAxesOp::verify() { return success(); } +LogicalResult LogicalMeshAxesOp::verify() { + ArrayRef axisExtents = getAxisExtents(); + if (getAxes().size() != axisExtents.size()) { + return emitOpError() << "requires number of results to match number of " + "axis_extents (" + << getAxes().size() << " != " << axisExtents.size() + << ")"; + } + + for (auto [idx, axisExtent] : llvm::enumerate(axisExtents)) { + if (axisExtent <= 0) { + return emitOpError() << "requires axis_extents[" << idx + << "] to be positive, got " << axisExtent; + } + + auto actualAxisType = + dyn_cast(getAxes()[idx].getType()); + if (!actualAxisType) { + return emitOpError() << "requires result #" << idx + << " to have LogicalMeshAxisType"; + } + + if (actualAxisType.getExtent() != static_cast(axisExtent)) { + return emitOpError() << "requires result #" << idx << " to have extent " + << axisExtent << ", but got " + << actualAxisType.getExtent(); + } + } + + return success(); +} + +LogicalResult LogicalMeshAxesOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + LogicalMeshAxesOpAdaptor adaptor(operands, attributes, properties, regions); + + ArrayRef axisExtents = adaptor.getAxisExtents(); + inferredReturnTypes.reserve(axisExtents.size()); + for (auto [idx, axisExtent] : llvm::enumerate(axisExtents)) { + if (axisExtent <= 0) { + if (location) { + mlir::emitError(*location) << "requires axis_extents[" << idx + << "] to be positive, got " << axisExtent; + } + return failure(); + } + + inferredReturnTypes.push_back( + LogicalMeshAxisType::get(context, static_cast(axisExtent))); + } + + return success(); +} + } // namespace mlir::enzyme::distributed diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td index 8a338dc961..6a8fea98c7 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -44,6 +44,19 @@ def GetPhysicalMeshAxesOp : DistributedOp<"GetPhysicalMeshAxes", [PureDistribute let hasVerifier = 1; } +def LogicalMeshAxesOp : DistributedOp<"LogicalMeshAxes", [PureDistributedMetadata, DeclareOpInterfaceMethods]> { + let description = [{ + Declares a new logical mesh of axes, which may or may not correspond to a + physical mesh. + }]; + let arguments = (ins + DenseI32ArrayAttr:$axis_extents + ); + let results = (outs Variadic:$axes); + let assemblyFormat = "$axis_extents attr-dict `:` type($axes)"; + let hasVerifier = 1; +} + def MeshComputationOp : DistributedOp<"MeshComputation", [ HasOnlyGraphRegion, NoTerminator, diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp b/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp deleted file mode 100644 index 1fcfa92552..0000000000 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp +++ /dev/null @@ -1 +0,0 @@ -#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h" diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.h b/src/enzyme_ad/jax/Passes/Distributed/Passes.h index ae8661f140..5e6b13040e 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.h +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.h @@ -13,8 +13,6 @@ namespace distributed { #define GEN_PASS_REGISTRATION #include "src/enzyme_ad/jax/Passes/Distributed/Passes.h.inc" -void registerdistributedPasses(); - } // namespace distributed } // namespace enzyme } // namespace mlir diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.td b/src/enzyme_ad/jax/Passes/Distributed/Passes.td index 6327ff9fb2..40854e15e3 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.td +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.td @@ -13,4 +13,23 @@ def PrintShardyFunctionNamesPass }]; } +def ShardyToDistributedPass : Pass<"shardy-to-distributed", "::mlir::ModuleOp"> { + let summary = "Convert Shardy functions to Distributed mesh computations."; + let description = [{ + Converts from Shardy to Distributed IR. This pass requires all Shardy functions + to be available in the module: calling external functions is not supported. + Also requires exactly one physical mesh to be declared in the module. + Shardy functions on multiple meshes will be supported in the future, but + for now this pass will fail if multiple meshes are declared. + + Note: this pass expects local tensor sizes as input and all + layout mismatches to have explicit layout conversion ops, either + reshard or reduce. This can be achieved with some collection of existing + shardy passes: + --sdy-propagation-pipeline + --sdy-export-pipeline='enable-insert-explicit-collectives=true' + --sdy-convert-global-to-local + }]; +} + #endif // ENZYME_AD_JAX_DISTRIBUTED_PASSES diff --git a/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp b/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp new file mode 100644 index 0000000000..cd17acf962 --- /dev/null +++ b/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp @@ -0,0 +1,140 @@ +#include "src/enzyme_ad/jax/Dialect/Distributed/Dialect.h" +#include "src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.h" +#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir::enzyme::distributed { + +#define GEN_PASS_DEF_SHARDYTODISTRIBUTEDPASS +#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h.inc" + +namespace { + +struct ShardyToDistributedPass + : public impl::ShardyToDistributedPassBase { + using ShardyToDistributedPassBase::ShardyToDistributedPassBase; + + // For now we simplify our life by requiring all shardy functions to have the + // same mesh. This will be relaxed at need as implementation progresses. + sdy::MeshAttr getCommonMesh(const FindShardyFunctionsAnalysis &analysis) { + auto shardyFunctions = analysis.getShardyFunctions(); + sdy::MeshAttr commonMesh = nullptr; + for (const FindShardyFunctionsAnalysis::FunctionInfo &info : + shardyFunctions) { + if (info.meshes.size() != 1) { + getOperation()->emitError() + << "expected shardy function to have exactly one mesh, found " + << info.meshes.size(); + signalPassFailure(); + return nullptr; + } + if (!commonMesh) { + commonMesh = info.meshes[0]; + } else if (commonMesh != info.meshes[0]) { + getOperation()->emitError() + << "expected all shardy functions to have the same mesh, found " + << commonMesh << " and " << info.meshes[0]; + signalPassFailure(); + return nullptr; + } + } + return commonMesh; + } + + void runOnOperation() override { + ModuleOp module = getOperation(); + + unsigned physicalMeshCount = 0; + distributed::PhysicalMeshOp physicalMesh; + for (distributed::PhysicalMeshOp meshOp : + module.getOps()) { + ++physicalMeshCount; + if (physicalMeshCount == 1) { + physicalMesh = meshOp; + } + if (physicalMeshCount > 1) { + module.emitError() + << "expected exactly one distributed physical mesh in module, " + "found " + << physicalMeshCount; + signalPassFailure(); + return; + } + } + + if (physicalMeshCount == 0) { + module.emitError() + << "expected exactly one distributed physical mesh in module, " + "found 0"; + signalPassFailure(); + return; + } + + const FindShardyFunctionsAnalysis &analysis = + getAnalysis(); + if (!analysis.isValid()) { + signalPassFailure(); + return; + } + if (analysis.getShardyFunctions().empty()) { + module.emitRemark() << "no shardy functions found, skipping pass"; + return; + } + + sdy::MeshAttr commonMesh = getCommonMesh(analysis); + if (!commonMesh) { + return; + } + + // Create a new mesh computation using the modules pysical mesh + OpBuilder builder(module.getContext()); + builder.setInsertionPointAfter( + physicalMesh); // graph region, doesn't matter + distributed::MeshComputationOp meshComputation = + builder.create( + module.getLoc(), builder.getStringAttr("mesh_computation"), + FlatSymbolRefAttr::get(physicalMesh.getSymNameAttr())); + Region &meshComputationBody = meshComputation.getBody(); + if (meshComputationBody.empty()) { + meshComputationBody.emplaceBlock(); + } + builder.setInsertionPointToStart(&meshComputationBody.front()); + + // Create a new logical mesh inside of the mesh computation using the common + // shardy mesh for axis sizes. Internally record a mapping from shardy axis + // names to distributed axis values. + llvm::SmallVector logicalAxisExtents; + llvm::SmallVector shardyAxisNames; + for (sdy::MeshAxisAttr axis : commonMesh.getAxes()) { + int64_t axisSize = axis.getSize(); + if (axisSize <= 0 || axisSize > std::numeric_limits::max()) { + module.emitError() << "unsupported shardy mesh axis size " << axisSize + << " for axis " << axis.getName(); + signalPassFailure(); + return; + } + logicalAxisExtents.push_back(static_cast(axisSize)); + shardyAxisNames.push_back(builder.getStringAttr(axis.getName())); + } + + distributed::LogicalMeshAxesOp logicalMeshAxes = + builder.create( + module.getLoc(), builder.getDenseI32ArrayAttr(logicalAxisExtents)); + + llvm::DenseMap shardyToDistributedAxis; + for (auto [idx, distributedAxis] : + llvm::enumerate(logicalMeshAxes.getAxes())) { + shardyToDistributedAxis[shardyAxisNames[idx]] = distributedAxis; + } + + (void)shardyToDistributedAxis; + (void)logicalMeshAxes; + (void)meshComputation; + } +}; + +} // namespace + +} // namespace mlir::enzyme::distributed \ No newline at end of file diff --git a/test/lit_tests/distributed/shardy_transformer_block_partial_export.mlir b/test/lit_tests/distributed/shardy_transformer_block_partial_export.mlir new file mode 100644 index 0000000000..b21a6f8494 --- /dev/null +++ b/test/lit_tests/distributed/shardy_transformer_block_partial_export.mlir @@ -0,0 +1,48 @@ +module @shardy_transformer_block_pre_export { + sdy.mesh @mesh = <["data"=4, "tile"=4, "model"=2]> {stablehlo.mesh = {axes = [{name = "data", size = 4 : i64}, {name = "tile", size = 4 : i64}, {name = "model", size = 2 : i64}]}} + func.func @transformer_block(%arg0: tensor<128x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {"tile"}]>}, %arg1: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg2: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg3: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg4: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg5: tensor<512x2048xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg6: tensor<2048x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}, {"tile"}]>}, %arg7: tensor<512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}]>}, %arg8: tensor<512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}]>}) -> (tensor<128x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {"model"}]>}) { + %0 = sdy.constant dense<0.000000e+00> : tensor + %1 = sdy.constant dense<9.99999974E-6> : tensor + %2 = sdy.constant dense<1.250000e-01> : tensor + %3 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + %4 = stablehlo.dot_general %arg0, %arg2, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + %5 = stablehlo.dot_general %arg0, %arg3, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + %6 = stablehlo.reshape %3 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<128x512xf32>) -> tensor<8x8x16x64xf32> + %7 = stablehlo.reshape %4 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<128x512xf32>) -> tensor<8x8x16x64xf32> + %8 = stablehlo.reshape %5 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<128x512xf32>) -> tensor<8x8x16x64xf32> + %9 = stablehlo.transpose %7, dims = [0, 1, 3, 2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<8x8x16x64xf32>) -> tensor<8x8x64x16xf32> + %10 = stablehlo.dot_general %6, %9, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<8x8x16x64xf32>, tensor<8x8x64x16xf32>) -> tensor<8x8x16x16xf32> + %11 = stablehlo.broadcast_in_dim %2, dims = [] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor) -> tensor<8x8x16x16xf32> + %12 = stablehlo.multiply %10, %11 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : tensor<8x8x16x16xf32> + %13 = stablehlo.exponential %12 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : tensor<8x8x16x16xf32> + %14 = stablehlo.reduce(%13 init: %0) applies stablehlo.add across dimensions = [3] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}]>]>} : (tensor<8x8x16x16xf32>, tensor) -> tensor<8x8x16xf32> + %15 = stablehlo.broadcast_in_dim %14, dims = [0, 1, 2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<8x8x16xf32>) -> tensor<8x8x16x16xf32> + %16 = stablehlo.divide %13, %15 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : tensor<8x8x16x16xf32> + %17 = stablehlo.dot_general %16, %8, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<8x8x16x16xf32>, tensor<8x8x16x64xf32>) -> tensor<8x8x16x64xf32> + %18 = stablehlo.reshape %17 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<8x8x16x64xf32>) -> tensor<128x512xf32> + %19 = stablehlo.dot_general %18, %arg4, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + %20 = sdy.reshard %19 <@mesh, [{"data"}, {"tile"}]> : tensor<128x512xf32> + %21 = stablehlo.add %arg0, %20 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<128x512xf32> + %22 = stablehlo.reduce(%21 init: %0) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor<128x512xf32>, tensor) -> tensor<128xf32> + %23 = stablehlo.broadcast_in_dim %22, dims = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<128xf32>) -> tensor<128x512xf32> + %24 = stablehlo.subtract %21, %23 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<128x512xf32> + %25 = stablehlo.multiply %24, %24 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<128x512xf32> + %26 = stablehlo.reduce(%25 init: %0) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor<128x512xf32>, tensor) -> tensor<128xf32> + %27 = stablehlo.broadcast_in_dim %1, dims = [] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor) -> tensor<128xf32> + %28 = stablehlo.add %26, %27 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : tensor<128xf32> + %29 = stablehlo.rsqrt %28 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : tensor<128xf32> + %30 = stablehlo.broadcast_in_dim %29, dims = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<128xf32>) -> tensor<128x512xf32> + %31 = stablehlo.multiply %24, %30 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<128x512xf32> + %32 = stablehlo.broadcast_in_dim %arg7, dims = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<512xf32>) -> tensor<128x512xf32> + %33 = stablehlo.broadcast_in_dim %arg8, dims = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<512xf32>) -> tensor<128x512xf32> + %34 = stablehlo.multiply %31, %32 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<128x512xf32> + %35 = stablehlo.add %34, %33 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<128x512xf32> + %36 = stablehlo.dot_general %35, %arg5, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<128x512xf32>, tensor<512x2048xf32>) -> tensor<128x2048xf32> + %37 = stablehlo.broadcast_in_dim %0, dims = [] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor) -> tensor<128x2048xf32> + %38 = stablehlo.maximum %36, %37 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : tensor<128x2048xf32> + %39 = stablehlo.dot_general %38, %arg6, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<128x2048xf32>, tensor<2048x512xf32>) -> tensor<128x512xf32> + %40 = stablehlo.add %35, %39 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : tensor<128x512xf32> + return %40 : tensor<128x512xf32> + } +} + diff --git a/test/lit_tests/distributed/shardy_transformer_block_pre_export.mlir b/test/lit_tests/distributed/shardy_transformer_block_pre_export.mlir new file mode 100644 index 0000000000..b777579802 --- /dev/null +++ b/test/lit_tests/distributed/shardy_transformer_block_pre_export.mlir @@ -0,0 +1,117 @@ +// RUN: enzymexlamlir-opt %s | FileCheck %s +// RUN: enzymexlamlir-opt --pass-pipeline='builtin.module(sdy-propagation-pipeline,sdy-close-shardings,func.func(sdy-insert-explicit-reshards))' %s | FileCheck %s --check-prefix=PIPE + +module @shardy_transformer_block_pre_export { + sdy.mesh @mesh = <["data"=4, "tile"=4, "model"=2]> + + func.func @transformer_block( + %arg0: tensor<128x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {"tile"}]>}, + %arg1: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, + %arg2: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, + %arg3: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, + %arg4: tensor<512x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, + %arg5: tensor<512x2048xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, + %arg6: tensor<2048x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}, {"tile"}]>}, + %arg7: tensor<512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}]>}, + %arg8: tensor<512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}]>}) + -> (tensor<128x512xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {"model"}]>}) { + %c0 = stablehlo.constant dense<0.000000e+00> : tensor + %ceps = stablehlo.constant dense<1.000000e-05> : tensor + %cscale = stablehlo.constant dense<1.250000e-01> : tensor + + // QKV projections. + %q = stablehlo.dot_general %arg0, %arg1, batching_dims = [] x [], contracting_dims = [1] x [0] : + (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + %k = stablehlo.dot_general %arg0, %arg2, batching_dims = [] x [], contracting_dims = [1] x [0] : + (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + %v = stablehlo.dot_general %arg0, %arg3, batching_dims = [] x [], contracting_dims = [1] x [0] : + (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + + // Multi-head attention as [batch=8, heads=8, seq=16, head_dim=64]. + %q_heads = stablehlo.reshape %q : (tensor<128x512xf32>) -> tensor<8x8x16x64xf32> + %k_heads = stablehlo.reshape %k : (tensor<128x512xf32>) -> tensor<8x8x16x64xf32> + %v_heads = stablehlo.reshape %v : (tensor<128x512xf32>) -> tensor<8x8x16x64xf32> + + %k_t = stablehlo.transpose %k_heads, dims = [0, 1, 3, 2] : + (tensor<8x8x16x64xf32>) -> tensor<8x8x64x16xf32> + %scores = stablehlo.dot_general %q_heads, %k_t, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] : + (tensor<8x8x16x64xf32>, tensor<8x8x64x16xf32>) -> tensor<8x8x16x16xf32> + + %scale_bcast = stablehlo.broadcast_in_dim %cscale, dims = [] : + (tensor) -> tensor<8x8x16x16xf32> + %scores_scaled = stablehlo.multiply %scores, %scale_bcast : tensor<8x8x16x16xf32> + + // Softmax-like normalization using exp/reduce/divide. + %weights = stablehlo.exponential %scores_scaled : tensor<8x8x16x16xf32> + %weights_sum = stablehlo.reduce(%weights init: %c0) applies stablehlo.add across dimensions = [3] : + (tensor<8x8x16x16xf32>, tensor) -> tensor<8x8x16xf32> + %weights_sum_bcast = stablehlo.broadcast_in_dim %weights_sum, dims = [0, 1, 2] : + (tensor<8x8x16xf32>) -> tensor<8x8x16x16xf32> + %weights_norm = stablehlo.divide %weights, %weights_sum_bcast : tensor<8x8x16x16xf32> + + %ctx_heads = stablehlo.dot_general %weights_norm, %v_heads, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] : + (tensor<8x8x16x16xf32>, tensor<8x8x16x64xf32>) -> tensor<8x8x16x64xf32> + %ctx = stablehlo.reshape %ctx_heads : (tensor<8x8x16x64xf32>) -> tensor<128x512xf32> + + // Output projection + residual. + %attn_out = stablehlo.dot_general %ctx, %arg4, batching_dims = [] x [], contracting_dims = [1] x [0] : + (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32> + %resid0 = stablehlo.add %arg0, %attn_out : tensor<128x512xf32> + + // LayerNorm-like normalization. + %mean = stablehlo.reduce(%resid0 init: %c0) applies stablehlo.add across dimensions = [1] : + (tensor<128x512xf32>, tensor) -> tensor<128xf32> + %mean_bcast = stablehlo.broadcast_in_dim %mean, dims = [0] : + (tensor<128xf32>) -> tensor<128x512xf32> + %centered = stablehlo.subtract %resid0, %mean_bcast : tensor<128x512xf32> + %sq = stablehlo.multiply %centered, %centered : tensor<128x512xf32> + %var = stablehlo.reduce(%sq init: %c0) applies stablehlo.add across dimensions = [1] : + (tensor<128x512xf32>, tensor) -> tensor<128xf32> + %eps_bcast = stablehlo.broadcast_in_dim %ceps, dims = [] : + (tensor) -> tensor<128xf32> + %var_eps = stablehlo.add %var, %eps_bcast : tensor<128xf32> + %inv_std = stablehlo.rsqrt %var_eps : tensor<128xf32> + %inv_std_bcast = stablehlo.broadcast_in_dim %inv_std, dims = [0] : + (tensor<128xf32>) -> tensor<128x512xf32> + %norm = stablehlo.multiply %centered, %inv_std_bcast : tensor<128x512xf32> + + %gamma = stablehlo.broadcast_in_dim %arg7, dims = [1] : + (tensor<512xf32>) -> tensor<128x512xf32> + %beta = stablehlo.broadcast_in_dim %arg8, dims = [1] : + (tensor<512xf32>) -> tensor<128x512xf32> + %norm_scaled = stablehlo.multiply %norm, %gamma : tensor<128x512xf32> + %norm_out = stablehlo.add %norm_scaled, %beta : tensor<128x512xf32> + + // MLP block + residual. + %ff1 = stablehlo.dot_general %norm_out, %arg5, batching_dims = [] x [], contracting_dims = [1] x [0] : + (tensor<128x512xf32>, tensor<512x2048xf32>) -> tensor<128x2048xf32> + %zero_ff = stablehlo.broadcast_in_dim %c0, dims = [] : + (tensor) -> tensor<128x2048xf32> + %ff1_relu = stablehlo.maximum %ff1, %zero_ff : tensor<128x2048xf32> + %ff2 = stablehlo.dot_general %ff1_relu, %arg6, batching_dims = [] x [], contracting_dims = [1] x [0] : + (tensor<128x2048xf32>, tensor<2048x512xf32>) -> tensor<128x512xf32> + %out = stablehlo.add %norm_out, %ff2 : tensor<128x512xf32> + + return %out : tensor<128x512xf32> + } +} + +// CHECK-LABEL: module @shardy_transformer_block_pre_export +// CHECK: sdy.mesh @mesh = <["data"=4, "tile"=4, "model"=2]> +// CHECK-LABEL: func.func @transformer_block( +// CHECK: stablehlo.dot_general %arg0, %arg1 +// CHECK: stablehlo.dot_general %arg0, %arg2 +// CHECK: stablehlo.dot_general %arg0, %arg3 +// CHECK: stablehlo.transpose %{{.*}}, dims = [0, 1, 3, 2] +// CHECK: stablehlo.reduce(%{{.*}} init: %{{.*}}) applies stablehlo.add across dimensions = [3] +// CHECK: stablehlo.rsqrt %{{.*}} : tensor<128xf32> +// CHECK: stablehlo.maximum %{{.*}}, %{{.*}} : tensor<128x2048xf32> +// CHECK: return %{{.*}} : tensor<128x512xf32> + +// PIPE-LABEL: module @shardy_transformer_block_pre_export +// PIPE: sdy.mesh @mesh = <["data"=4, "tile"=4, "model"=2]> {stablehlo.mesh = +// PIPE-LABEL: func.func @transformer_block( +// PIPE: stablehlo.dot_general %arg0, %arg1{{.*}}{sdy.sharding = +// PIPE: stablehlo.dot_general %{{.*}}, %{{.*}} batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2]{{.*}}{sdy.sharding = +// PIPE: stablehlo.rsqrt %{{.*}} {{.*}}{sdy.sharding = +// PIPE: return %{{.*}} : tensor<128x512xf32> diff --git a/test/lit_tests/distributed/shardy_transformer_exported.mlir b/test/lit_tests/distributed/shardy_transformer_exported.mlir new file mode 100644 index 0000000000..0dea1b7b49 --- /dev/null +++ b/test/lit_tests/distributed/shardy_transformer_exported.mlir @@ -0,0 +1,121 @@ +module @shardy_transformer_block_pre_export { + sdy.mesh @mesh = <["data"=4, "tile"=4, "model"=2]> {stablehlo.mesh = {axes = [{name = "data", size = 4 : i64}, {name = "tile", size = 4 : i64}, {name = "model", size = 2 : i64}]}} + func.func @transformer_block(%arg0: tensor<32x128xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {"tile"}]>}, %arg1: tensor<128x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg2: tensor<128x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg3: tensor<128x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg4: tensor<128x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg5: tensor<128x1024xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"tile"}, {"model"}]>}, %arg6: tensor<1024x128xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}, {"tile"}]>}, %arg7: tensor<256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}]>}, %arg8: tensor<256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}]>}) -> (tensor<32x256xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {"model"}]>}) { + %cst = stablehlo.constant dense<0.000000e+00> : tensor + %cst_0 = stablehlo.constant dense<9.99999974E-6> : tensor + %cst_1 = stablehlo.constant dense<1.250000e-01> : tensor + %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> + %1 = "stablehlo.all_reduce"(%0) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32x256xf32>) -> tensor<32x256xf32> + %2 = stablehlo.dot_general %arg0, %arg2, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> + %3 = "stablehlo.all_reduce"(%2) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32x256xf32>) -> tensor<32x256xf32> + %4 = stablehlo.dot_general %arg0, %arg3, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> + %5 = "stablehlo.all_reduce"(%4) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32x256xf32>) -> tensor<32x256xf32> + %6 = "stablehlo.all_gather"(%1) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> + %7 = stablehlo.reshape %6 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<32x512xf32>) -> tensor<2x8x16x64xf32> + %8 = "stablehlo.all_gather"(%3) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> + %9 = stablehlo.reshape %8 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<32x512xf32>) -> tensor<2x8x16x64xf32> + %10 = "stablehlo.all_gather"(%5) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> + %11 = stablehlo.reshape %10 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<32x512xf32>) -> tensor<2x8x16x64xf32> + %12 = stablehlo.transpose %9, dims = [0, 1, 3, 2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<2x8x16x64xf32>) -> tensor<2x8x64x16xf32> + %13 = stablehlo.dot_general %7, %12, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<2x8x16x64xf32>, tensor<2x8x64x16xf32>) -> tensor<2x8x16x16xf32> + %14 = stablehlo.broadcast_in_dim %cst_1, dims = [] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor) -> tensor<2x8x16x16xf32> + %15 = stablehlo.multiply %13, %14 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : tensor<2x8x16x16xf32> + %16 = stablehlo.exponential %15 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : tensor<2x8x16x16xf32> + %17 = stablehlo.reduce(%16 init: %cst) applies stablehlo.add across dimensions = [3] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}]>]>} : (tensor<2x8x16x16xf32>, tensor) -> tensor<2x8x16xf32> + %18 = stablehlo.broadcast_in_dim %17, dims = [0, 1, 2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<2x8x16xf32>) -> tensor<2x8x16x16xf32> + %19 = stablehlo.divide %16, %18 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : tensor<2x8x16x16xf32> + %20 = stablehlo.dot_general %19, %11, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<2x8x16x16xf32>, tensor<2x8x16x64xf32>) -> tensor<2x8x16x64xf32> + %21 = stablehlo.reshape %20 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}]>]>} : (tensor<2x8x16x64xf32>) -> tensor<32x512xf32> + %c = stablehlo.constant dense<0> : tensor + %22 = stablehlo.partition_id : tensor + %23 = stablehlo.convert %22 : (tensor) -> tensor + %c_2 = stablehlo.constant dense<[0, 0, 128, 128, 256, 256, 384, 384, 0, 0, 128, 128, 256, 256, 384, 384, 0, 0, 128, 128, 256, 256, 384, 384, 0, 0, 128, 128, 256, 256, 384, 384]> : tensor<32xi64> + %24 = stablehlo.dynamic_slice %c_2, %23, sizes = [1] : (tensor<32xi64>, tensor) -> tensor<1xi64> + %25 = stablehlo.reshape %24 : (tensor<1xi64>) -> tensor + %26 = stablehlo.dynamic_slice %21, %c, %25, sizes = [32, 128] : (tensor<32x512xf32>, tensor, tensor) -> tensor<32x128xf32> + %27 = stablehlo.dot_general %26, %arg4, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> + %28 = "stablehlo.all_reduce"(%27) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32x256xf32>) -> tensor<32x256xf32> + %c_3 = stablehlo.constant dense<0> : tensor + %29 = stablehlo.partition_id : tensor + %30 = stablehlo.convert %29 : (tensor) -> tensor + %c_4 = stablehlo.constant dense<[0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128]> : tensor<32xi64> + %31 = stablehlo.dynamic_slice %c_4, %30, sizes = [1] : (tensor<32xi64>, tensor) -> tensor<1xi64> + %32 = stablehlo.reshape %31 : (tensor<1xi64>) -> tensor + %33 = stablehlo.dynamic_slice %28, %c_3, %32, sizes = [32, 128] : (tensor<32x256xf32>, tensor, tensor) -> tensor<32x128xf32> + %34 = "stablehlo.collective_permute"(%33) <{channel_handle = #stablehlo.channel_handle, source_target_pairs = dense<[[0, 0], [1, 2], [2, 4], [3, 6], [4, 1], [5, 3], [6, 5], [7, 7], [8, 8], [9, 10], [10, 12], [11, 14], [12, 9], [13, 11], [14, 13], [15, 15], [16, 16], [17, 18], [18, 20], [19, 22], [20, 17], [21, 19], [22, 21], [23, 23], [24, 24], [25, 26], [26, 28], [27, 30], [28, 25], [29, 27], [30, 29], [31, 31]]> : tensor<32x2xi64>}> : (tensor<32x128xf32>) -> tensor<32x128xf32> + %35 = stablehlo.add %arg0, %34 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> + %36 = stablehlo.reduce(%35 init: %cst) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor<32x128xf32>, tensor) -> tensor<32xf32> + %37 = "stablehlo.all_reduce"(%36) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32xf32>) -> tensor<32xf32> + %38 = stablehlo.broadcast_in_dim %37, dims = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<32xf32>) -> tensor<32x128xf32> + %39 = stablehlo.subtract %35, %38 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> + %40 = stablehlo.multiply %39, %39 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> + %41 = stablehlo.reduce(%40 init: %cst) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor<32x128xf32>, tensor) -> tensor<32xf32> + %42 = "stablehlo.all_reduce"(%41) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32xf32>) -> tensor<32xf32> + %43 = stablehlo.broadcast_in_dim %cst_0, dims = [] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor) -> tensor<32xf32> + %44 = stablehlo.add %42, %43 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : tensor<32xf32> + %45 = stablehlo.rsqrt %44 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : tensor<32xf32> + %46 = stablehlo.broadcast_in_dim %45, dims = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<32xf32>) -> tensor<32x128xf32> + %47 = stablehlo.multiply %39, %46 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> + %48 = stablehlo.partition_id : tensor + %49 = stablehlo.convert %48 : (tensor) -> tensor + %c_5 = stablehlo.constant dense<[0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128]> : tensor<32xi64> + %50 = stablehlo.dynamic_slice %c_5, %49, sizes = [1] : (tensor<32xi64>, tensor) -> tensor<1xi64> + %51 = stablehlo.reshape %50 : (tensor<1xi64>) -> tensor + %52 = stablehlo.dynamic_slice %arg7, %51, sizes = [128] : (tensor<256xf32>, tensor) -> tensor<128xf32> + %53 = "stablehlo.collective_permute"(%52) <{channel_handle = #stablehlo.channel_handle, source_target_pairs = dense<[[0, 0], [1, 2], [2, 4], [3, 6], [4, 1], [5, 3], [6, 5], [7, 7], [8, 8], [9, 10], [10, 12], [11, 14], [12, 9], [13, 11], [14, 13], [15, 15], [16, 16], [17, 18], [18, 20], [19, 22], [20, 17], [21, 19], [22, 21], [23, 23], [24, 24], [25, 26], [26, 28], [27, 30], [28, 25], [29, 27], [30, 29], [31, 31]]> : tensor<32x2xi64>}> : (tensor<128xf32>) -> tensor<128xf32> + %54 = stablehlo.broadcast_in_dim %53, dims = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<128xf32>) -> tensor<32x128xf32> + %55 = stablehlo.partition_id : tensor + %56 = stablehlo.convert %55 : (tensor) -> tensor + %c_6 = stablehlo.constant dense<[0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128, 0, 0, 128, 128]> : tensor<32xi64> + %57 = stablehlo.dynamic_slice %c_6, %56, sizes = [1] : (tensor<32xi64>, tensor) -> tensor<1xi64> + %58 = stablehlo.reshape %57 : (tensor<1xi64>) -> tensor + %59 = stablehlo.dynamic_slice %arg8, %58, sizes = [128] : (tensor<256xf32>, tensor) -> tensor<128xf32> + %60 = "stablehlo.collective_permute"(%59) <{channel_handle = #stablehlo.channel_handle, source_target_pairs = dense<[[0, 0], [1, 2], [2, 4], [3, 6], [4, 1], [5, 3], [6, 5], [7, 7], [8, 8], [9, 10], [10, 12], [11, 14], [12, 9], [13, 11], [14, 13], [15, 15], [16, 16], [17, 18], [18, 20], [19, 22], [20, 17], [21, 19], [22, 21], [23, 23], [24, 24], [25, 26], [26, 28], [27, 30], [28, 25], [29, 27], [30, 29], [31, 31]]> : tensor<32x2xi64>}> : (tensor<128xf32>) -> tensor<128xf32> + %61 = stablehlo.broadcast_in_dim %60, dims = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<128xf32>) -> tensor<32x128xf32> + %62 = stablehlo.multiply %47, %54 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> + %63 = stablehlo.add %62, %61 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> + %64 = stablehlo.dot_general %63, %arg5, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x1024xf32>) -> tensor<32x1024xf32> + %65 = "stablehlo.all_reduce"(%64) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32x1024xf32>) -> tensor<32x1024xf32> + %66 = stablehlo.broadcast_in_dim %cst, dims = [] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor) -> tensor<32x1024xf32> + %67 = stablehlo.maximum %65, %66 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : tensor<32x1024xf32> + %68 = stablehlo.dot_general %67, %arg6, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<32x1024xf32>, tensor<1024x128xf32>) -> tensor<32x128xf32> + %69 = "stablehlo.all_reduce"(%68) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> ({ + ^bb0(%arg9: tensor, %arg10: tensor): + %73 = stablehlo.add %arg9, %arg10 : tensor + stablehlo.return %73 : tensor + }) : (tensor<32x128xf32>) -> tensor<32x128xf32> + %70 = stablehlo.add %63, %69 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> + %71 = "stablehlo.collective_permute"(%70) <{channel_handle = #stablehlo.channel_handle, source_target_pairs = dense<[[0, 0], [1, 2], [2, 1], [3, 3], [4, 4], [5, 6], [6, 5], [7, 7], [8, 8], [9, 10], [10, 9], [11, 11], [12, 12], [13, 14], [14, 13], [15, 15], [16, 16], [17, 18], [18, 17], [19, 19], [20, 20], [21, 22], [22, 21], [23, 23], [24, 24], [25, 26], [26, 25], [27, 27], [28, 28], [29, 30], [30, 29], [31, 31]]> : tensor<32x2xi64>}> : (tensor<32x128xf32>) -> tensor<32x128xf32> + %72 = "stablehlo.all_gather"(%71) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7], [8, 10], [9, 11], [12, 14], [13, 15], [16, 18], [17, 19], [20, 22], [21, 23], [24, 26], [25, 27], [28, 30], [29, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x128xf32>) -> tensor<32x256xf32> + return %72 : tensor<32x256xf32> + } +} + From b2b2d75ca2b3e732bd323e49a176c398a7b9883d Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 06/11] Collective parsing (verifier failure) --- src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp | 205 ++++++++++++------ src/enzyme_ad/jax/Dialect/Axis/Utilities.h | 27 ++- .../jax/Dialect/Distributed/CollectiveOps.cpp | 187 ++++++++++++++++ .../jax/Dialect/Distributed/CollectiveOps.h | 119 ++++++++++ .../Distributed/MeshComputationOps.cpp | 1 + src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 83 ++++++- .../jax/Dialect/Distributed/Types.cpp | 13 +- .../jax/Dialect/Distributed/Utilities.cpp | 35 ++- .../jax/Dialect/Distributed/Utilities.h | 22 +- test/lit_tests/distributed/roundtrip.mlir | 43 ++++ 10 files changed, 642 insertions(+), 93 deletions(-) create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.cpp create mode 100644 src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.h diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp index fe17835e4f..868696115a 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp @@ -2,8 +2,26 @@ #include +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/raw_ostream.h" + namespace mlir::enzyme::axis { +template +static TypedValue castTypedValue(Value value, llvm::StringRef expectedType) { + if (auto typed = dyn_cast>(value)) { + return typed; + } + + std::string typeString; + llvm::raw_string_ostream os(typeString); + value.getType().print(os); + os.flush(); + llvm::errs() << "castTypedValue failed: expected " << expectedType + << ", got value type " << typeString << "\n"; + llvm::report_fatal_error("invalid typed value cast"); +} + // Dispatches alias checks for canonical axes. Canonical axes are // either equivalent or wholly disjoint. static bool areAxesEquivalent(Value lhs, Value rhs) { @@ -27,18 +45,14 @@ static bool areAxesEquivalent(Value lhs, Value rhs) { bool arePairwiseFactorsDisjoint(Value lhsFactor, Value rhsFactor, Value lhsProvenanceAxis, Value rhsProvenanceAxis) { - auto lhsType = dyn_cast(lhsFactor.getType()); - auto rhsType = dyn_cast(rhsFactor.getType()); - assert(lhsType && "factor value must have AxisFactorType"); - assert(rhsType && "factor value must have AxisFactorType"); - if (!lhsType || !rhsType) { - return false; - } + auto lhsTyped = castTypedValue(lhsFactor, "AxisFactorType"); + auto rhsTyped = castTypedValue(rhsFactor, "AxisFactorType"); + auto lhsType = lhsTyped.getType(); + auto rhsType = rhsTyped.getType(); Value lhsAxis = lhsProvenanceAxis; if (!lhsAxis) { - auto lhsProvenance = - getFactorProvenanceAxis(cast>(lhsFactor)); + auto lhsProvenance = getFactorProvenanceAxis(lhsTyped); assert(succeeded(lhsProvenance) && "factor must have a provenance axis"); if (failed(lhsProvenance)) { return false; @@ -48,8 +62,7 @@ bool arePairwiseFactorsDisjoint(Value lhsFactor, Value rhsFactor, Value rhsAxis = rhsProvenanceAxis; if (!rhsAxis) { - auto rhsProvenance = - getFactorProvenanceAxis(cast>(rhsFactor)); + auto rhsProvenance = getFactorProvenanceAxis(rhsTyped); assert(succeeded(rhsProvenance) && "factor must have a provenance axis"); if (failed(rhsProvenance)) { return false; @@ -83,24 +96,18 @@ bool arePairwiseFactorsDisjoint(Value lhsFactor, Value rhsFactor, } // Asserts an axis (not factor) type and gets the extent. -int getAxisExtent(Value axis) { - auto axisInterface = dyn_cast(axis.getType()); - assert(axisInterface && "axis type must implement AxisTypeInterface"); - return static_cast(axisInterface.extent()); +int getAxisExtent(TypedValue axis) { + return static_cast(axis.getType().extent()); } // Asserts a factor type and gets the extent. -int getFactorExtent(Value factor) { - auto factorType = dyn_cast(factor.getType()); - assert(factorType && "factor type must be AxisFactorType"); - return static_cast(factorType.getExtent()); +int getFactorExtent(TypedValue factor) { + return static_cast(factor.getType().getExtent()); } // Asserts a segment type and gets the extent. -int getSegmentExtent(Value segment) { - auto segmentType = dyn_cast(segment.getType()); - assert(segmentType && "segment type must be AxisSegmentType"); - return static_cast(segmentType.getExtent()); +int getSegmentExtent(TypedValue segment) { + return static_cast(segment.getType().getExtent()); } // Returns the defining op for a canonical axis SSA value. @@ -140,6 +147,26 @@ getProductProvenanceFactors(TypedValue factorProduct) { return productOp.getFactors(); } +// Returns the product of extents for a factor-product SSA value. +FailureOr +getFactorGroupExtent(TypedValue factorProduct) { + auto factors = getProductProvenanceFactors(factorProduct); + if (failed(factors)) { + return failure(); + } + + uint64_t extent = 1; + for (Value factor : *factors) { + auto factorType = dyn_cast(factor.getType()); + if (!factorType) { + return failure(); + } + extent *= static_cast(factorType.getExtent()); + } + + return extent; +} + // Checks factor compatibility and pairwise non-overlap metadata. bool areFactorsDisjoint(ValueRange factors) { if (factors.empty()) { @@ -157,16 +184,12 @@ bool areFactorsDisjoint(ValueRange factors) { SmallVector cachedFactors; cachedFactors.reserve(factors.size()); for (Value factor : factors) { - auto factorType = dyn_cast(factor.getType()); - assert(factorType && "factor value must have AxisFactorType"); - if (!factorType) { - return false; - } + auto factorTyped = castTypedValue(factor, "AxisFactorType"); + auto factorType = factorTyped.getType(); assert(factorType.getExtent() > 0 && "factor extent must be positive"); assert(factorType.getStride() > 0 && "factor stride must be positive"); - auto provenance = - getFactorProvenanceAxis(cast>(factor)); + auto provenance = getFactorProvenanceAxis(factorTyped); assert(succeeded(provenance) && "factor must have a provenance axis"); cachedFactors.push_back({*provenance}); } @@ -208,12 +231,10 @@ bool areFactorIndexSpacesEqual(ValueRange lhsFactors, ValueRange rhsFactors) { out.clear(); out.reserve(factors.size()); for (Value factor : factors) { - auto factorType = dyn_cast(factor.getType()); - if (!factorType) { - return false; - } - auto provenance = - getFactorProvenanceAxis(cast>(factor)); + auto factorTyped = + castTypedValue(factor, "AxisFactorType"); + auto factorType = factorTyped.getType(); + auto provenance = getFactorProvenanceAxis(factorTyped); if (failed(provenance)) { return false; } @@ -259,18 +280,16 @@ bool areFactorIndexSpacesEqual(ValueRange lhsFactors, ValueRange rhsFactors) { bool arePairwiseSegmentsDisjoint(Value lhsSegment, Value rhsSegment, Value lhsProvenanceAxis, Value rhsProvenanceAxis) { - auto lhsType = dyn_cast(lhsSegment.getType()); - auto rhsType = dyn_cast(rhsSegment.getType()); - assert(lhsType && "segment value must have AxisSegmentType"); - assert(rhsType && "segment value must have AxisSegmentType"); - if (!lhsType || !rhsType) { - return false; - } + auto lhsTyped = + castTypedValue(lhsSegment, "AxisSegmentType"); + auto rhsTyped = + castTypedValue(rhsSegment, "AxisSegmentType"); + auto lhsType = lhsTyped.getType(); + auto rhsType = rhsTyped.getType(); Value lhsAxis = lhsProvenanceAxis; if (!lhsAxis) { - auto lhsProvenance = - getSegmentProvenanceAxis(cast>(lhsSegment)); + auto lhsProvenance = getSegmentProvenanceAxis(lhsTyped); assert(succeeded(lhsProvenance) && "segment must have a provenance axis"); if (failed(lhsProvenance)) { return false; @@ -280,8 +299,7 @@ bool arePairwiseSegmentsDisjoint(Value lhsSegment, Value rhsSegment, Value rhsAxis = rhsProvenanceAxis; if (!rhsAxis) { - auto rhsProvenance = - getSegmentProvenanceAxis(cast>(rhsSegment)); + auto rhsProvenance = getSegmentProvenanceAxis(rhsTyped); assert(succeeded(rhsProvenance) && "segment must have a provenance axis"); if (failed(rhsProvenance)) { return false; @@ -313,15 +331,12 @@ bool areSegmentsDisjoint(ValueRange segments) { SmallVector provenanceAxes; provenanceAxes.reserve(segments.size()); for (Value segment : segments) { - auto segmentType = dyn_cast(segment.getType()); - assert(segmentType && "segment value must have AxisSegmentType"); - if (!segmentType) { - return false; - } + auto segmentTyped = + castTypedValue(segment, "AxisSegmentType"); + auto segmentType = segmentTyped.getType(); assert(segmentType.getExtent() > 0 && "segment extent must be positive"); - auto provenance = - getSegmentProvenanceAxis(cast>(segment)); + auto provenance = getSegmentProvenanceAxis(segmentTyped); assert(succeeded(provenance) && "segment must have a provenance axis"); if (failed(provenance)) { return false; @@ -354,23 +369,20 @@ bool areFactorsComplete(Value axis, ValueRange factors) { // axis and their extents cover the whole source-axis extent. uint64_t product = 1; for (Value factor : factors) { - auto factorType = dyn_cast(factor.getType()); - assert(factorType && "factor value must have AxisFactorType"); - if (!factorType) { - return false; - } + auto factorTyped = castTypedValue(factor, "AxisFactorType"); - auto provenance = - getFactorProvenanceAxis(cast>(factor)); + auto provenance = getFactorProvenanceAxis(factorTyped); assert(succeeded(provenance) && "factor must have a provenance axis"); if (failed(provenance) || *provenance != axis) { return false; } - product *= static_cast(getFactorExtent(factor)); + product *= static_cast(getFactorExtent(factorTyped)); } - return product == static_cast(getAxisExtent(axis)); + return product == + static_cast(getAxisExtent( + castTypedValue(axis, "AxisTypeInterface"))); } // Checks that segments reconstruct the full source axis interval [0, extent). @@ -383,14 +395,11 @@ bool areSegmentsComplete(Value axis, ValueRange segments) { intervals.reserve(segments.size()); for (Value segment : segments) { - auto segmentType = dyn_cast(segment.getType()); - assert(segmentType && "segment value must have AxisSegmentType"); - if (!segmentType) { - return false; - } + auto segmentTyped = + castTypedValue(segment, "AxisSegmentType"); + auto segmentType = segmentTyped.getType(); - auto provenance = - getSegmentProvenanceAxis(cast>(segment)); + auto provenance = getSegmentProvenanceAxis(segmentTyped); assert(succeeded(provenance) && "segment must have a provenance axis"); if (failed(provenance) || *provenance != axis) { return false; @@ -414,7 +423,61 @@ bool areSegmentsComplete(Value axis, ValueRange segments) { cursor = end; } - return cursor == static_cast(getAxisExtent(axis)); + return cursor == + static_cast(getAxisExtent( + castTypedValue(axis, "AxisTypeInterface"))); +} + +llvm::SmallVector<::mlir::Value> +flattenGroupsToFactors(::mlir::ValueRange factorGroups) { + llvm::SmallVector<::mlir::Value> flattenedFactors; + for (auto group : factorGroups) { + auto typedGroup = + cast<::mlir::TypedValue<::mlir::enzyme::axis::FactorGroupType>>(group); + auto factors = getProductProvenanceFactors(typedGroup); + if (failed(factors)) { + llvm::report_fatal_error( + "flattenGroupsToFactors failed to get factors from FactorGroupType"); + } + flattenedFactors.append(factors->begin(), factors->end()); + } + return flattenedFactors; +} + +bool areFactorGroupsDisjoint(::mlir::ValueRange factorGroups) { + auto flattenedFactors = flattenGroupsToFactors(factorGroups); + return areFactorsDisjoint(ValueRange(flattenedFactors)); } +llvm::SmallVector<::mlir::Value> +createAxesForRankedShape(::mlir::Type shapeType, ::mlir::OpBuilder &builder, + ::mlir::Location loc) { + auto rankedShapeType = cast(shapeType); + auto type_attr = TypeAttr::get(rankedShapeType); + int rank = rankedShapeType.getRank(); + llvm::SmallVector<::mlir::Value> axes; + axes.reserve(rank); + for (int i = 0; i < rank; ++i) { + auto rank_attr = builder.getI32IntegerAttr(i); + auto axis = builder.create(loc, type_attr, rank_attr); + axes.push_back(axis); + } + return axes; +} + +llvm::SmallVector<::mlir::Value> viewAxesAsFactors(::mlir::ValueRange axes, + ::mlir::OpBuilder &builder, + ::mlir::Location loc) { + llvm::SmallVector<::mlir::Value> factors; + factors.reserve(axes.size()); + for (auto axis : axes) { + auto axis_typed = + castTypedValue(axis, "AxisTypeInterface"); + int extent = getAxisExtent(axis_typed); + auto factor = + builder.create(loc, axis, ArrayRef{extent}); + factors.push_back(factor.getResult(0)); + } + return factors; +} } // namespace mlir::enzyme::axis diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h index 6e9c8e560f..09993d7aec 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h @@ -6,13 +6,13 @@ namespace mlir::enzyme::axis { // Returns the static extent for any canonical axis SSA value. -int getAxisExtent(::mlir::Value axis); +int getAxisExtent(::mlir::TypedValue axis); // Returns the static extent for any factor SSA value. -int getFactorExtent(::mlir::Value factor); +int getFactorExtent(::mlir::TypedValue factor); // Returns the static extent for any segment SSA value. -int getSegmentExtent(::mlir::Value segment); +int getSegmentExtent(::mlir::TypedValue segment); // Returns the defining op for a canonical axis SSA value. ::mlir::FailureOr<::mlir::Operation *> getAxisProvenanceOp(::mlir::Value axis); @@ -29,6 +29,10 @@ getSegmentProvenanceAxis(::mlir::TypedValue segment); ::mlir::FailureOr<::mlir::ValueRange> getProductProvenanceFactors(::mlir::TypedValue factorProduct); +// Returns the product of extents for a factor-product SSA value. +::mlir::FailureOr +getFactorGroupExtent(::mlir::TypedValue factorProduct); + // Checks that factors are pairwise non-overlapping for one source axis. bool arePairwiseFactorsDisjoint(::mlir::Value lhsFactor, ::mlir::Value rhsFactor, @@ -59,6 +63,23 @@ bool areFactorsComplete(::mlir::Value axis, ::mlir::ValueRange factors); // Checks that segments exactly cover [0, axis extent) with no overlap or gaps. bool areSegmentsComplete(::mlir::Value axis, ::mlir::ValueRange segments); +// Small utlity for extracting all factors from one or more factor groups. +llvm::SmallVector<::mlir::Value> +flattenGroupsToFactors(::mlir::ValueRange factorGroups); + +// Shortcut for calling distjoint on a flattened factor list +bool areFactorGroupsDisjoint(::mlir::ValueRange factorGroups); + +// Given a shapeType with a known rank, returns a list of canonical axes for +// each dimension of that shape. Use a builder without an insertion point and +// an unknown location for ephemeral axes. +llvm::SmallVector<::mlir::Value> +createAxesForRankedShape(::mlir::Type shapeType, ::mlir::OpBuilder &builder, + ::mlir::Location loc); +// Creates a single factor for each axis, with full extent and stride 1. +llvm::SmallVector<::mlir::Value> viewAxesAsFactors(::mlir::ValueRange axes, + ::mlir::OpBuilder &builder, + ::mlir::Location loc); } // namespace mlir::enzyme::axis #endif // ENZYME_AD_JAX_DIALECT_AXIS_UTILITIES_H diff --git a/src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.cpp new file mode 100644 index 0000000000..35b6a92bcd --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.cpp @@ -0,0 +1,187 @@ +#include "CollectiveOps.h" + +namespace mlir::enzyme::distributed { + +// small helper +llvm::SmallVector<::mlir::Value> concatRanges(::mlir::ValueRange lhs, + ::mlir::ValueRange rhs) { + llvm::SmallVector<::mlir::Value> result; + result.reserve(lhs.size() + rhs.size()); + result.append(lhs.begin(), lhs.end()); + result.append(rhs.begin(), rhs.end()); + return result; +} + +LogicalResult DistributedCollectiveOp::verify() { + auto inputMeshFactors = axis::getProductProvenanceFactors(getInputMesh()); + if (failed(inputMeshFactors)) { + return emitOpError() + << "requires input_mesh to be produced by axis.product"; + } + + auto outputMeshFactors = axis::getProductProvenanceFactors(getOutputMesh()); + if (failed(outputMeshFactors)) { + return emitOpError() + << "requires output_mesh to be produced by axis.product"; + } + + if (!axis::areFactorIndexSpacesEqual(*inputMeshFactors, *outputMeshFactors)) { + return emitOpError() + << "requires input_mesh and output_mesh to have equal index space"; + } + + ArrayAttr reductionFunctions = getReductionFunctionsAttr(); + if (!reductionFunctions) { + return emitOpError() << "requires reduction_functions attribute"; + } + if (reductionFunctions.size() != getReductionGroups().size()) { + return emitOpError() << "requires reduction_functions size to match " + "reduction_groups size (" + << reductionFunctions.size() + << " != " << getReductionGroups().size() << ")"; + } + + for (auto [idx, reductionFunctionAttr] : + llvm::enumerate(reductionFunctions)) { + auto reductionFunction = + dyn_cast_or_null(reductionFunctionAttr); + if (!reductionFunction) { + return emitOpError() << "requires reduction_functions[" << idx + << "] to be a FlatSymbolRefAttr"; + } + + if (!lookupSymbolInEnclosingScopes(*this, reductionFunction)) { + return emitOpError() << "references unknown reduction function symbol " + << reductionFunction; + } + } + + auto reduction_group_factors = + axis::flattenGroupsToFactors(getReductionGroups()); + auto mapping_lhs_factors = axis::flattenGroupsToFactors(getMappingLhs()); + auto mapping_rhs_factors = axis::flattenGroupsToFactors(getMappingRhs()); + auto lhs_filtered = filterOutReplicationFactors(mapping_lhs_factors); + auto rhs_filtered = filterOutReplicationFactors(mapping_rhs_factors); + + // Create the set of axis we expect to see from the input, output types. + OpBuilder builder(getContext()); + builder.clearInsertionPoint(); + Location loc = getLoc(); + auto expected_input_tensor_axes = + axis::createAxesForRankedShape(getInputObject().getType(), builder, loc); + auto expected_output_tensor_axes = + axis::createAxesForRankedShape(getOutputTensorType(), builder, loc); + auto expected_input_factors = + axis::viewAxesAsFactors(expected_input_tensor_axes, builder, loc); + auto expected_output_factors = + axis::viewAxesAsFactors(expected_output_tensor_axes, builder, loc); + + /* + * Want: + * - reduction disjoint with lhs, rhs (rhs and lhs may overlap) + * - reduction + lhs_filtered = input_mesh + tensor axes + * - rhs_filtered = output_mesh + tensor axes + */ + auto lhs_space = concatRanges(reduction_group_factors, lhs_filtered); + auto expected_input_space = + concatRanges(*inputMeshFactors, expected_input_factors); + auto expected_output_space = + concatRanges(*outputMeshFactors, expected_output_factors); + if (!axis::areFactorsDisjoint(lhs_space)) { + return emitOpError() + << "requires reduction_groups + mapping_lhs to be disjoint"; + } + if (!axis::areFactorsDisjoint(rhs_filtered)) { + return emitOpError() << "requires mapping_rhs to be disjoint"; + } + if (!axis::areFactorIndexSpacesEqual(lhs_space, expected_input_space)) { + return emitOpError() + << "requires reduction_groups + mapping_lhs to match input_mesh " + "+ input_tensor axes"; + } + if (!axis::areFactorIndexSpacesEqual(rhs_filtered, expected_output_space)) { + return emitOpError() + << "requires mapping_rhs to match output_mesh + output_tensor axes"; + } + + if (!axis::areFactorGroupsDisjoint(getReductionGroups())) { + return emitOpError() << "requires reduction_groups to be pairwise disjoint"; + } + + if (getMappingLhs().size() != getMappingRhs().size()) { + return emitOpError() << "requires mapping_lhs and mapping_rhs to have the" + << " same length (" << getMappingLhs().size() + << " != " << getMappingRhs().size() << ")"; + } + + for (auto [idx, lhsMapping] : llvm::enumerate(getMappingLhs())) { + Value rhsMapping = getMappingRhs()[idx]; + + FailureOr lhsExtent = axis::getFactorGroupExtent( + cast>(lhsMapping)); + if (failed(lhsExtent)) { + return emitOpError() << "requires mapping_lhs[" << idx + << "] to be produced by axis.product"; + } + + FailureOr rhsExtent = axis::getFactorGroupExtent( + cast>(rhsMapping)); + if (failed(rhsExtent)) { + return emitOpError() << "requires mapping_rhs[" << idx + << "] to be produced by axis.product"; + } + + if (*lhsExtent != *rhsExtent) { + return emitOpError() << "requires mapping pair #" << idx + << " to have matching extent (" << *lhsExtent + << " != " << *rhsExtent << ")"; + } + } + + return success(); +} + +LogicalResult DistributedCollectiveOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + DistributedCollectiveOpAdaptor adaptor(operands, attributes, properties, + regions); + inferredReturnTypes.push_back( + AsynchHandleType::get(context, adaptor.getOutputTensorType())); + return success(); +} + +LogicalResult DistributedAwait::verify() { + auto handleType = dyn_cast(getAsyncHandle().getType()); + if (!handleType) { + return emitOpError() << "requires async_handle to be an AsynchHandleType"; + } + + Type expectedValueType = handleType.getValueType(); + Type actualValueType = getValue().getType(); + if (actualValueType != expectedValueType) { + return emitOpError() << "requires result type to match awaited handle " + << "value type " << expectedValueType << ", but got " + << actualValueType; + } + + return success(); +} + +LogicalResult DistributedAwait::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + DistributedAwaitAdaptor adaptor(operands, attributes, properties, regions); + auto handleType = + dyn_cast(adaptor.getAsyncHandle().getType()); + if (!handleType) { + return failure(); + } + + inferredReturnTypes.push_back(handleType.getValueType()); + return success(); +} + +} // namespace mlir::enzyme::distributed \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.h b/src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.h new file mode 100644 index 0000000000..8ebf2c4e99 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Distributed/CollectiveOps.h @@ -0,0 +1,119 @@ +#ifndef ENZYME_AD_JAX_DIALECT_DISTRIBUTED_COLLECTIVE_OPS_H +#define ENZYME_AD_JAX_DIALECT_DISTRIBUTED_COLLECTIVE_OPS_H + +#include "Dialect.h" +#include "Utilities.h" + +namespace mlir::enzyme::distributed { + +static ParseResult parseReductionGroups( + OpAsmParser &parser, + SmallVectorImpl &reductionGroups, + ArrayAttr &reductionFunctions) { + SmallVector parsedReductionFunctions; + + if (parser.parseLParen()) { + return failure(); + } + + if (succeeded(parser.parseOptionalRParen())) { + reductionFunctions = ArrayAttr::get(parser.getBuilder().getContext(), + parsedReductionFunctions); + return success(); + } + + while (true) { + OpAsmParser::UnresolvedOperand reductionGroup; + if (parser.parseOperand(reductionGroup)) { + return failure(); + } + reductionGroups.push_back(reductionGroup); + + Attribute reductionFunctionAttr; + if (parser.parseAttribute(reductionFunctionAttr)) { + return failure(); + } + auto reductionFunction = + dyn_cast_or_null(reductionFunctionAttr); + if (!reductionFunction) { + return parser.emitError(parser.getCurrentLocation()) + << "requires reduction function to be a flat symbol reference"; + } + parsedReductionFunctions.push_back(reductionFunction); + + if (succeeded(parser.parseOptionalComma())) { + continue; + } + if (parser.parseRParen()) { + return failure(); + } + break; + } + + reductionFunctions = ArrayAttr::get(parser.getBuilder().getContext(), + parsedReductionFunctions); + return success(); +} + +static void printReductionGroups(OpAsmPrinter &printer, Operation *op, + OperandRange reductionGroups, + ArrayAttr reductionFunctions) { + printer << '('; + for (auto [idx, reductionGroup] : llvm::enumerate(reductionGroups)) { + if (idx != 0) { + printer << ", "; + } + printer << reductionGroup << ' ' << reductionFunctions[idx]; + } + printer << ')'; +} + +static ParseResult +parseAxisMapping(OpAsmParser &parser, + SmallVectorImpl &mappingLhs, + SmallVectorImpl &mappingRhs) { + if (parser.parseLParen()) { + return failure(); + } + + if (succeeded(parser.parseOptionalRParen())) { + return success(); + } + + while (true) { + OpAsmParser::UnresolvedOperand lhs; + OpAsmParser::UnresolvedOperand rhs; + if (parser.parseOperand(lhs) || parser.parseArrow() || + parser.parseOperand(rhs)) { + return failure(); + } + mappingLhs.push_back(lhs); + mappingRhs.push_back(rhs); + + if (succeeded(parser.parseOptionalComma())) { + continue; + } + if (parser.parseRParen()) { + return failure(); + } + break; + } + + return success(); +} + +static void printAxisMapping(OpAsmPrinter &printer, Operation *op, + OperandRange mappingLhs, OperandRange mappingRhs) { + printer << '('; + for (auto [idx, lhs] : llvm::enumerate(mappingLhs)) { + if (idx != 0) { + printer << ", "; + } + printer << lhs << " -> " << mappingRhs[idx]; + } + printer << ')'; +} + +} // namespace mlir::enzyme::distributed + +#endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_COLLECTIVE_OPS_H \ 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 index 6dbb2c201e..348db14740 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp @@ -1,3 +1,4 @@ +#include "CollectiveOps.h" #include "Dialect.h" namespace mlir::enzyme::distributed { diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td index 6a8fea98c7..db2783a3ed 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -14,10 +14,6 @@ include "Interfaces.td" include "src/enzyme_ad/jax/Dialect/Axis/Dialect.td" include "src/enzyme_ad/jax/Dialect/Axis/Types.td" -defvar TensorType = HLO_Tensor; -defvar ShardedTensorType = HLO_Tensor; // Subject to change, for now shardy just adds attributes on top of a regular tensor. -defvar ShardingAttribute = Sdy_TensorSharding; - // Marker for distributed metadata ops. Includes builtin Pure. def PureDistributedMetadata : TraitList<[MetadataOpTrait]>; @@ -113,6 +109,17 @@ def DistributedYieldOp : DistributedOp<"DistributedYield", [Terminator]> { let assemblyFormat = "$returns type($returns) attr-dict"; } +def DistributedAwait : DistributedOp<"Await", [DeclareOpInterfaceMethods]> { + let description = [{ + Awaits a value that is being computed asynchronously. Extracts the value + from an AsynchHandleType. This is a no-op if the value is already available. + }]; + let arguments = (ins AsynchHandleType:$async_handle); + let results = (outs AnyType:$value); + let assemblyFormat = "$async_handle attr-dict `:` type($async_handle) `->` type($value)"; + let hasVerifier = 1; +} + def DistributedCallOp : DistributedOp<"DistributedCall", []> { let description = [{ Calls a distributed function. The execution context must match the enclosing context. @@ -126,4 +133,72 @@ def DistributedCallOp : DistributedOp<"DistributedCall", []> { let assemblyFormat = "$callee `replicate_over` $replicate_over $arguments attr-dict `:` type($replicate_over) `,` type($arguments) `->` type($returns)"; let hasVerifier = 1; } + +def DistributedCollectiveOp : DistributedOp<"Collective", [ + DeclareOpInterfaceMethods, + AttrSizedOperandSegments +]> { + let description = [{ + An operation representing reductions and resharding along multiple axis. Does not + cover collective permutes and point-to-point communication. + + A collective operation happens over the mesh and tile axis, and maps input mesh/tile + axis factors to output mesh/tile axis factors. Essentially, this is a + reduce-transpose-reshape operation. + + Reductions are specified as different groups of axis factors to reduce over. + Reductions take custom binary functions as input, which must be associative. + Each group of reduction axes may have a different reduction function function, + but these must be commutative with each other: for now, use two distinct collectives + when reductions cannot be fused. + + The transpose/reshape is specified as a mapping between non-reduced input/output axis factors. + This mapping space can also include replicate(N) factors on either side: mapping an input axis + to replicate means that you are asserting each device on that dimension contributes identical values, + while mapping replicate to an output axis means you are cloning the result over that axis. + + Collectives must be between the same mesh (verified via index space equality). + + Inputs: + - input_object: a shaped tensor (tensor or memref). + - input_mesh: the mesh axes (FactorGroup) that the input object is distributed over. + - output_mesh: the mesh axes (FactorGroup) that the output object is distributed over. + Must have equal index space to input_mesh. + - reduction_groups: a list of FactorGroups to reduce over (pairwise disjoint). + - reduction_functions: ArrayAttr of FlatSymbolRefAttr naming single-device reduction + functions. Size must match reduction_groups. + - mapping_lhs: input axis groups to map from (post-reduction). + - mapping_rhs: output axis groups to map to (same length, compatible index spaces). + + Output: + - An AsynchHandleType wrapping the output tensor type. + + Custom format: + - Reductions are specified as (%axisgroup @reduction_function), and are collected by + the parser into the reduction_groups list and reduction_functions attributes. + - Mappings are specified as (%axisgroup -> %axisgroup), and are collected by the parser into + the mapping_lhs and mapping_rhs lists. + }]; + let arguments = (ins + AnyTypeOf<[AnyTensor, AnyMemRef]>:$input_object, + FactorGroupType:$input_mesh, + FactorGroupType:$output_mesh, + Variadic:$reduction_groups, + ArrayAttr:$reduction_functions, + Variadic:$mapping_lhs, + Variadic:$mapping_rhs, + TypeAttrOf:$output_tensor_type + ); + let results = (outs AsynchHandleType:$async_handle); + let assemblyFormat = [{ + $input_object `:` type($input_object) `on` $input_mesh `:` type($input_mesh) + `to` $output_tensor_type `on` $output_mesh `:` type($output_mesh) + `reduces` custom($reduction_groups, $reduction_functions) + `:` type($reduction_groups) + `maps` custom($mapping_lhs, $mapping_rhs) + `:` `[` type($mapping_lhs) `]` `[` type($mapping_rhs) `]` + attr-dict + }]; + let hasVerifier = 1; +} #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_OPS_TD \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Types.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Types.cpp index 7e50e8745d..3cae0448de 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Types.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Types.cpp @@ -14,10 +14,15 @@ bool PhysicalCommAxisType::aliases(Value ax1, Value ax2) const { } bool LogicalMeshAxisType::aliases(Value ax1, Value ax2) const { - (void)ax1; - (void)ax2; - llvm::report_fatal_error( - "LogicalMeshAxisType::aliases is not implemented yet"); + // alias iff they are the same result of the same op + auto result1 = dyn_cast(ax1); + auto result2 = dyn_cast(ax2); + if (!result1 || !result2) { + llvm::report_fatal_error( + "LogicalMeshAxisType::aliases requires both axes to be OpResults"); + } + return result1.getOwner() == result2.getOwner() && + result1.getResultNumber() == result2.getResultNumber(); } // Replication axes are modeled as always disjoint. diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp index bc97e4529f..f07d30f376 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp @@ -2,8 +2,25 @@ namespace mlir::enzyme::distributed { -using axis::FactorGroupType; -using axis::getProductProvenanceFactors; +using namespace ::mlir::enzyme::axis; + +Operation *lookupSymbolInEnclosingScopes(Operation *from, + FlatSymbolRefAttr symRef) { + if (!from || !symRef) { + return nullptr; + } + + for (auto *scope = from; scope; scope = scope->getParentOp()) { + if (!scope->hasTrait()) { + continue; + } + if (auto *op = SymbolTable::lookupSymbolIn(scope, symRef)) { + return op; + } + } + + return nullptr; +} FailureOr> getEnclosingExecutionContext(Operation *op) { @@ -31,4 +48,18 @@ expandExecutionContextFactors(TypedValue context) { return copiedFactors; } +::llvm::SmallVector<::mlir::Value> +filterOutReplicationFactors(::mlir::ValueRange factors) { + llvm::SmallVector<::mlir::Value> filteredFactors; + for (auto factor : factors) { + // type of factor should wrap replication axis if it is a replication factor + auto factorType = + cast<::mlir::enzyme::axis::AxisFactorType>(factor.getType()); + ::mlir::Type axisType = factorType.getAxisType(); + if (!isa<::mlir::enzyme::distributed::ReplicationAxisType>(axisType)) { + filteredFactors.push_back(factor); + } + } + return filteredFactors; +} } // namespace mlir::enzyme::distributed \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h index 4a926dbdbe..8abf3efb7f 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h @@ -6,6 +6,11 @@ namespace mlir::enzyme::distributed { +// Walks parent operations and checks each symbol table scope for a flat symbol. +::mlir::Operation * +lookupSymbolInEnclosingScopes(::mlir::Operation *from, + ::mlir::FlatSymbolRefAttr symRef); + template ::mlir::FailureOr resolveSymbolOpFromAttr(::mlir::Operation *from, ::mlir::Attribute opAttr) { @@ -14,16 +19,11 @@ ::mlir::FailureOr resolveSymbolOpFromAttr(::mlir::Operation *from, return ::mlir::failure(); } - for (auto *scope = from; scope; scope = scope->getParentOp()) { - if (!scope->hasTrait<::mlir::OpTrait::SymbolTable>()) { - continue; - } - if (auto *op = ::mlir::SymbolTable::lookupSymbolIn(scope, symRef)) { - if (auto typedOp = llvm::dyn_cast(op)) { - return typedOp; - } - return ::mlir::failure(); + if (auto *op = lookupSymbolInEnclosingScopes(from, symRef)) { + if (auto typedOp = llvm::dyn_cast(op)) { + return typedOp; } + return ::mlir::failure(); } return ::mlir::failure(); @@ -39,6 +39,10 @@ ::mlir::FailureOr<::llvm::SmallVector<::mlir::Value>> expandExecutionContextFactors( ::mlir::TypedValue<::mlir::enzyme::axis::FactorGroupType> context); +// Creates a new range with all replication axes removed from the input range. +::llvm::SmallVector<::mlir::Value> +filterOutReplicationFactors(::mlir::ValueRange factors); + } // namespace mlir::enzyme::distributed #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_UTILITIES_H diff --git a/test/lit_tests/distributed/roundtrip.mlir b/test/lit_tests/distributed/roundtrip.mlir index d538782e23..085baacd90 100644 --- a/test/lit_tests/distributed/roundtrip.mlir +++ b/test/lit_tests/distributed/roundtrip.mlir @@ -36,6 +36,49 @@ module { // ----- +// Roundtrip distributed.Collective / distributed.Await. +module { + func.func @add(%lhs: f32, %rhs: f32) -> f32 { + %0 = arith.addf %lhs, %rhs : f32 + return %0 : f32 + } + + distributed.PhysicalMesh @mesh0 device_target "cpu" axes [!distributed.physical_comm_axis<2>, !distributed.physical_comm_axis<2>] + + distributed.MeshComputation @mc mesh @mesh0 { + %p0, %p1 = distributed.GetPhysicalMeshAxes @mesh0 : !distributed.physical_comm_axis<2>, !distributed.physical_comm_axis<2> + + %l0, %l1 = distributed.LogicalMeshAxes [2, 2] : !distributed.logical_mesh_axis<2>, !distributed.logical_mesh_axis<2> + %lf0 = axis.factor %l0 [2] : !distributed.logical_mesh_axis<2> + %lf1 = axis.factor %l1 [2] : !distributed.logical_mesh_axis<2> + %ta = axis.getaxis tensor<8xf32> 0 + %ta_out = axis.getaxis tensor<4xf32> 0 + %tf_out = axis.factor %ta_out [4] : !axis.shape_axis, 0> + %tf_to_mesh, %tf_remain = axis.factor %ta [4, 2] : !axis.shape_axis, 0> + %mesh_in = axis.product %lf0, %lf1 : !axis.axis_factor, 2, 1>, !axis.axis_factor, 2, 1> + %mesh_out = axis.product %lf0, %lf1 : !axis.axis_factor, 2, 1>, !axis.axis_factor, 2, 1> + %reduction = axis.product %lf1 : !axis.axis_factor, 2, 1> + %lhs_group_1 = axis.product %tf_to_mesh : !axis.axis_factor, 0>, 4, 2> + // rhs_group_1 = %mesh_out + %lhs_group_2 = axis.product %lf0, %tf_remain : !axis.axis_factor, 2, 1>, !axis.axis_factor, 0>, 2, 1> + %rhs_group_2 = axis.product %tf_out : !axis.axis_factor, 0>, 4, 1> + + distributed.Function @collective context %mesh_in : !axis.factor_group<4> arg_types [tensor<8xf32>] ret_types [tensor<4xf32>] { + ^bb0(%arg0: tensor<8xf32>): + %h = distributed.Collective %arg0 : tensor<8xf32> on %mesh_in : !axis.factor_group<4> to tensor<4xf32> on %mesh_out : !axis.factor_group<4> reduces (%reduction @add) : !axis.factor_group<2> maps (%lhs_group_1 -> %mesh_out, %lhs_group_2 -> %rhs_group_2) : [!axis.factor_group<4>, !axis.factor_group<4>] [!axis.factor_group<4>, !axis.factor_group<4>] + %v = distributed.Await %h : !distributed.asynch_handle> -> tensor<4xf32> + distributed.DistributedYield %v tensor<4xf32> + } + } +} + +// CHECK: func.func @add(%{{.*}}: f32, %{{.*}}: f32) -> f32 +// CHECK: distributed.PhysicalMesh @mesh0 device_target "cpu" axes [!distributed.physical_comm_axis<2>, !distributed.physical_comm_axis<2>] +// CHECK: %{{.*}} = distributed.Collective %{{.*}} : tensor<8xf32> on %{{.*}} : !axis.factor_group<4> to tensor<8xf32> on %{{.*}} : !axis.factor_group<4> reduces (%{{.*}} @add) : !axis.factor_group<2> maps (%{{.*}} -> %{{.*}}) : [!axis.factor_group<16>] [!axis.factor_group<16>] +// CHECK: %{{.*}} = distributed.Await %{{.*}} : !distributed.asynch_handle> -> tensor<8xf32> + +// ----- + // Roundtrip distributed.MeshComputation with a minimal metadata-only body. module { distributed.PhysicalMesh @mesh0 device_target "cpu" axes [!distributed.physical_comm_axis<4>] From e336862a7e0e3e842335a008e17282744ee0e49a Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 07/11] Fix factor space comparison --- src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp | 114 ++++++++++++------- 1 file changed, 76 insertions(+), 38 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp index 868696115a..0dd19bbbd9 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp @@ -212,63 +212,101 @@ bool areFactorsDisjoint(ValueRange factors) { return true; } -// Compares two factor lists as index-space descriptors, ignoring ordering. -// This is multiset equality over (extent, stride, provenance-axis equivalence) -// and is intentionally permutation-invariant. -bool areFactorIndexSpacesEqual(ValueRange lhsFactors, ValueRange rhsFactors) { - if (lhsFactors.size() != rhsFactors.size()) { - return false; +// From a list of factors known to be from the same axis, +// creates a list of pairs indicating the maximum factor ranges. +// Ranges are gauranteed to be return in major-first order. +llvm::SmallVector> build_max_factors(ValueRange factors) { + if (factors.empty()) { + return {}; + } + // convert into intervals + llvm::SmallVector> factor_pairs; + for (Value factor : factors) { + auto factorTyped = castTypedValue(factor, "AxisFactorType"); + auto factorType = factorTyped.getType(); + int extent = static_cast(factorType.getExtent()); + int stride = static_cast(factorType.getStride()); + factor_pairs.push_back({extent, stride}); + } + // sort intervals by stride + std::sort( + factor_pairs.begin(), factor_pairs.end(), + [](const auto &lhs, const auto &rhs) { return lhs.second > rhs.second; }); + + llvm::SmallVector> max_factors; + std::pair current_factor = factor_pairs[0]; + for (size_t i = 1; i < factor_pairs.size(); ++i) { + // if the stride of the current factor = stride * extent of the next factor, + // they can be combined. + const auto &next_factor = factor_pairs[i]; + if (current_factor.second == next_factor.first * next_factor.second) { + current_factor.first *= next_factor.first; + current_factor.second = next_factor.second; + } else { + max_factors.push_back(current_factor); + current_factor = next_factor; + } } + max_factors.push_back(current_factor); + return max_factors; +} - struct FactorInfo { +// Compares two factor lists as index-space descriptors, ignoring ordering. +// This is multiset equality over (extent, stride, provenance-axis +// equivalence) and is intentionally permutation-invariant. +bool areFactorIndexSpacesEqual(ValueRange lhsFactors, ValueRange rhsFactors) { + struct AxisFactors { Value provenance; - unsigned extent; - unsigned stride; + SmallVector lhsFactors; + SmallVector rhsFactors; }; - auto buildFactorInfo = [](ValueRange factors, - SmallVectorImpl &out) -> bool { - out.clear(); - out.reserve(factors.size()); + auto addFactorsToBuckets = [](ValueRange factors, bool isLhs, + SmallVectorImpl &grouped) { for (Value factor : factors) { auto factorTyped = castTypedValue(factor, "AxisFactorType"); - auto factorType = factorTyped.getType(); auto provenance = getFactorProvenanceAxis(factorTyped); if (failed(provenance)) { return false; } - out.push_back( - {*provenance, factorType.getExtent(), factorType.getStride()}); + + bool inserted = false; + for (AxisFactors &bucket : grouped) { + if (areAxesEquivalent(bucket.provenance, *provenance)) { + if (isLhs) { + bucket.lhsFactors.push_back(factor); + } else { + bucket.rhsFactors.push_back(factor); + } + inserted = true; + break; + } + } + if (!inserted) { + AxisFactors bucket; + bucket.provenance = *provenance; + if (isLhs) { + bucket.lhsFactors.push_back(factor); + } else { + bucket.rhsFactors.push_back(factor); + } + grouped.push_back(std::move(bucket)); + } } return true; }; - SmallVector lhsInfo; - SmallVector rhsInfo; - if (!buildFactorInfo(lhsFactors, lhsInfo) || - !buildFactorInfo(rhsFactors, rhsInfo)) { + SmallVector grouped; + if (!addFactorsToBuckets(lhsFactors, /*isLhs=*/true, grouped) || + !addFactorsToBuckets(rhsFactors, /*isLhs=*/false, grouped)) { return false; } - SmallVector rhsMatched(rhsInfo.size(), false); - for (const FactorInfo &lhs : lhsInfo) { - bool foundMatch = false; - for (auto [rhsIndex, rhs] : llvm::enumerate(rhsInfo)) { - if (rhsMatched[rhsIndex]) { - continue; - } - if (lhs.extent != rhs.extent || lhs.stride != rhs.stride) { - continue; - } - if (!areAxesEquivalent(lhs.provenance, rhs.provenance)) { - continue; - } - rhsMatched[rhsIndex] = true; - foundMatch = true; - break; - } - if (!foundMatch) { + for (AxisFactors &bucket : grouped) { + auto lhsMaxFactors = build_max_factors(ValueRange(bucket.lhsFactors)); + auto rhsMaxFactors = build_max_factors(ValueRange(bucket.rhsFactors)); + if (lhsMaxFactors != rhsMaxFactors) { return false; } } From 40966d5c109d22ca1e7f13f712f69ef27bea3591 Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Wed, 15 Jul 2026 13:33:39 -0500 Subject: [PATCH 08/11] Mesh getting and shardy pipeline --- .../jax/Dialect/Distributed/Utilities.cpp | 29 ++++++++++++ .../jax/Dialect/Distributed/Utilities.h | 4 ++ .../jax/Passes/Distributed/Passes.td | 2 +- .../Distributed/ShardyToDistributed.cpp | 44 ++++++------------- .../shardy_transformer_exported.mlir | 24 +++++----- 5 files changed, 59 insertions(+), 44 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp index f07d30f376..6107b5f007 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp @@ -22,6 +22,35 @@ Operation *lookupSymbolInEnclosingScopes(Operation *from, return nullptr; } +FailureOr findUniquePhysicalMesh(ModuleOp moduleOp) { + if (!moduleOp) { + return failure(); + } + + unsigned physicalMeshCount = 0; + PhysicalMeshOp physicalMesh; + for (PhysicalMeshOp meshOp : moduleOp.getOps()) { + ++physicalMeshCount; + if (physicalMeshCount == 1) { + physicalMesh = meshOp; + } + if (physicalMeshCount > 1) { + moduleOp.emitError() + << "expected exactly one distributed physical mesh in module, found " + << physicalMeshCount; + return failure(); + } + } + + if (physicalMeshCount == 0) { + moduleOp.emitError() + << "expected exactly one distributed physical mesh in module, found 0"; + return failure(); + } + + return physicalMesh; +} + FailureOr> getEnclosingExecutionContext(Operation *op) { auto parentFunction = op ? op->getParentOfType() diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h index 8abf3efb7f..37fc04f257 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h @@ -11,6 +11,10 @@ ::mlir::Operation * lookupSymbolInEnclosingScopes(::mlir::Operation *from, ::mlir::FlatSymbolRefAttr symRef); +// Finds the unique distributed physical mesh in the module. +::mlir::FailureOr<::mlir::enzyme::distributed::PhysicalMeshOp> +findUniquePhysicalMesh(::mlir::ModuleOp moduleOp); + template ::mlir::FailureOr resolveSymbolOpFromAttr(::mlir::Operation *from, ::mlir::Attribute opAttr) { diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.td b/src/enzyme_ad/jax/Passes/Distributed/Passes.td index 40854e15e3..7d462e17ff 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.td +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.td @@ -28,7 +28,7 @@ def ShardyToDistributedPass : Pass<"shardy-to-distributed", "::mlir::ModuleOp"> shardy passes: --sdy-propagation-pipeline --sdy-export-pipeline='enable-insert-explicit-collectives=true' - --sdy-convert-global-to-local + --sdy-convert-global-to-local='enable-rgv3=true' }]; } diff --git a/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp b/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp index cd17acf962..d2ca94a7f6 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp +++ b/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp @@ -1,4 +1,5 @@ #include "src/enzyme_ad/jax/Dialect/Distributed/Dialect.h" +#include "src/enzyme_ad/jax/Dialect/Distributed/Utilities.h" #include "src/enzyme_ad/jax/Passes/Distributed/FindShardyFunctionsAnalysis.h" #include "src/enzyme_ad/jax/Passes/Distributed/Passes.h" @@ -44,30 +45,11 @@ struct ShardyToDistributedPass } void runOnOperation() override { - ModuleOp module = getOperation(); - - unsigned physicalMeshCount = 0; - distributed::PhysicalMeshOp physicalMesh; - for (distributed::PhysicalMeshOp meshOp : - module.getOps()) { - ++physicalMeshCount; - if (physicalMeshCount == 1) { - physicalMesh = meshOp; - } - if (physicalMeshCount > 1) { - module.emitError() - << "expected exactly one distributed physical mesh in module, " - "found " - << physicalMeshCount; - signalPassFailure(); - return; - } - } + ModuleOp moduleOp = getOperation(); - if (physicalMeshCount == 0) { - module.emitError() - << "expected exactly one distributed physical mesh in module, " - "found 0"; + FailureOr physicalMesh = + distributed::findUniquePhysicalMesh(moduleOp); + if (failed(physicalMesh)) { signalPassFailure(); return; } @@ -79,7 +61,7 @@ struct ShardyToDistributedPass return; } if (analysis.getShardyFunctions().empty()) { - module.emitRemark() << "no shardy functions found, skipping pass"; + moduleOp.emitRemark() << "no shardy functions found, skipping pass"; return; } @@ -89,13 +71,13 @@ struct ShardyToDistributedPass } // Create a new mesh computation using the modules pysical mesh - OpBuilder builder(module.getContext()); + OpBuilder builder(moduleOp.getContext()); builder.setInsertionPointAfter( - physicalMesh); // graph region, doesn't matter + *physicalMesh); // graph region, doesn't matter distributed::MeshComputationOp meshComputation = builder.create( - module.getLoc(), builder.getStringAttr("mesh_computation"), - FlatSymbolRefAttr::get(physicalMesh.getSymNameAttr())); + moduleOp.getLoc(), builder.getStringAttr("mesh_computation"), + FlatSymbolRefAttr::get(physicalMesh->getSymNameAttr())); Region &meshComputationBody = meshComputation.getBody(); if (meshComputationBody.empty()) { meshComputationBody.emplaceBlock(); @@ -110,8 +92,8 @@ struct ShardyToDistributedPass for (sdy::MeshAxisAttr axis : commonMesh.getAxes()) { int64_t axisSize = axis.getSize(); if (axisSize <= 0 || axisSize > std::numeric_limits::max()) { - module.emitError() << "unsupported shardy mesh axis size " << axisSize - << " for axis " << axis.getName(); + moduleOp.emitError() << "unsupported shardy mesh axis size " << axisSize + << " for axis " << axis.getName(); signalPassFailure(); return; } @@ -121,7 +103,7 @@ struct ShardyToDistributedPass distributed::LogicalMeshAxesOp logicalMeshAxes = builder.create( - module.getLoc(), builder.getDenseI32ArrayAttr(logicalAxisExtents)); + moduleOp.getLoc(), builder.getDenseI32ArrayAttr(logicalAxisExtents)); llvm::DenseMap shardyToDistributedAxis; for (auto [idx, distributedAxis] : diff --git a/test/lit_tests/distributed/shardy_transformer_exported.mlir b/test/lit_tests/distributed/shardy_transformer_exported.mlir index 0dea1b7b49..196b465b64 100644 --- a/test/lit_tests/distributed/shardy_transformer_exported.mlir +++ b/test/lit_tests/distributed/shardy_transformer_exported.mlir @@ -5,28 +5,28 @@ module @shardy_transformer_block_pre_export { %cst_0 = stablehlo.constant dense<9.99999974E-6> : tensor %cst_1 = stablehlo.constant dense<1.250000e-01> : tensor %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> - %1 = "stablehlo.all_reduce"(%0) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + %1 = "stablehlo.all_reduce"(%0) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor }) : (tensor<32x256xf32>) -> tensor<32x256xf32> %2 = stablehlo.dot_general %arg0, %arg2, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> - %3 = "stablehlo.all_reduce"(%2) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + %3 = "stablehlo.all_reduce"(%2) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor }) : (tensor<32x256xf32>) -> tensor<32x256xf32> %4 = stablehlo.dot_general %arg0, %arg3, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> - %5 = "stablehlo.all_reduce"(%4) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + %5 = "stablehlo.all_reduce"(%4) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor }) : (tensor<32x256xf32>) -> tensor<32x256xf32> - %6 = "stablehlo.all_gather"(%1) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> + %6 = "stablehlo.all_gather"(%1) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> %7 = stablehlo.reshape %6 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<32x512xf32>) -> tensor<2x8x16x64xf32> - %8 = "stablehlo.all_gather"(%3) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> + %8 = "stablehlo.all_gather"(%3) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> %9 = stablehlo.reshape %8 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<32x512xf32>) -> tensor<2x8x16x64xf32> - %10 = "stablehlo.all_gather"(%5) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> + %10 = "stablehlo.all_gather"(%5) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> : (tensor<32x256xf32>) -> tensor<32x512xf32> %11 = stablehlo.reshape %10 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<32x512xf32>) -> tensor<2x8x16x64xf32> %12 = stablehlo.transpose %9, dims = [0, 1, 3, 2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<2x8x16x64xf32>) -> tensor<2x8x64x16xf32> %13 = stablehlo.dot_general %7, %12, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {}, {}, {}]>]>} : (tensor<2x8x16x64xf32>, tensor<2x8x64x16xf32>) -> tensor<2x8x16x16xf32> @@ -46,7 +46,7 @@ module @shardy_transformer_block_pre_export { %25 = stablehlo.reshape %24 : (tensor<1xi64>) -> tensor %26 = stablehlo.dynamic_slice %21, %c, %25, sizes = [32, 128] : (tensor<32x512xf32>, tensor, tensor) -> tensor<32x128xf32> %27 = stablehlo.dot_general %26, %arg4, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x256xf32>) -> tensor<32x256xf32> - %28 = "stablehlo.all_reduce"(%27) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + %28 = "stablehlo.all_reduce"(%27) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor @@ -61,7 +61,7 @@ module @shardy_transformer_block_pre_export { %34 = "stablehlo.collective_permute"(%33) <{channel_handle = #stablehlo.channel_handle, source_target_pairs = dense<[[0, 0], [1, 2], [2, 4], [3, 6], [4, 1], [5, 3], [6, 5], [7, 7], [8, 8], [9, 10], [10, 12], [11, 14], [12, 9], [13, 11], [14, 13], [15, 15], [16, 16], [17, 18], [18, 20], [19, 22], [20, 17], [21, 19], [22, 21], [23, 23], [24, 24], [25, 26], [26, 28], [27, 30], [28, 25], [29, 27], [30, 29], [31, 31]]> : tensor<32x2xi64>}> : (tensor<32x128xf32>) -> tensor<32x128xf32> %35 = stablehlo.add %arg0, %34 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> %36 = stablehlo.reduce(%35 init: %cst) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor<32x128xf32>, tensor) -> tensor<32xf32> - %37 = "stablehlo.all_reduce"(%36) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + %37 = "stablehlo.all_reduce"(%36) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor @@ -70,7 +70,7 @@ module @shardy_transformer_block_pre_export { %39 = stablehlo.subtract %35, %38 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> %40 = stablehlo.multiply %39, %39 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> %41 = stablehlo.reduce(%40 init: %cst) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}]>]>} : (tensor<32x128xf32>, tensor) -> tensor<32xf32> - %42 = "stablehlo.all_reduce"(%41) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + %42 = "stablehlo.all_reduce"(%41) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor @@ -99,7 +99,7 @@ module @shardy_transformer_block_pre_export { %62 = stablehlo.multiply %47, %54 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> %63 = stablehlo.add %62, %61 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> %64 = stablehlo.dot_general %63, %arg5, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor<32x128xf32>, tensor<128x1024xf32>) -> tensor<32x1024xf32> - %65 = "stablehlo.all_reduce"(%64) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23], [24, 26, 28, 30], [25, 27, 29, 31]]> : tensor<8x4xi64>, use_global_device_ids}> ({ + %65 = "stablehlo.all_reduce"(%64) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor @@ -107,14 +107,14 @@ module @shardy_transformer_block_pre_export { %66 = stablehlo.broadcast_in_dim %cst, dims = [] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : (tensor) -> tensor<32x1024xf32> %67 = stablehlo.maximum %65, %66 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"model"}]>]>} : tensor<32x1024xf32> %68 = stablehlo.dot_general %67, %arg6, contracting_dims = [1] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : (tensor<32x1024xf32>, tensor<1024x128xf32>) -> tensor<32x128xf32> - %69 = "stablehlo.all_reduce"(%68) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31]]> : tensor<16x2xi64>, use_global_device_ids}> ({ + %69 = "stablehlo.all_reduce"(%68) <{channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> ({ ^bb0(%arg9: tensor, %arg10: tensor): %73 = stablehlo.add %arg9, %arg10 : tensor stablehlo.return %73 : tensor }) : (tensor<32x128xf32>) -> tensor<32x128xf32> %70 = stablehlo.add %63, %69 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"data"}, {"tile"}]>]>} : tensor<32x128xf32> %71 = "stablehlo.collective_permute"(%70) <{channel_handle = #stablehlo.channel_handle, source_target_pairs = dense<[[0, 0], [1, 2], [2, 1], [3, 3], [4, 4], [5, 6], [6, 5], [7, 7], [8, 8], [9, 10], [10, 9], [11, 11], [12, 12], [13, 14], [14, 13], [15, 15], [16, 16], [17, 18], [18, 17], [19, 19], [20, 20], [21, 22], [22, 21], [23, 23], [24, 24], [25, 26], [26, 25], [27, 27], [28, 28], [29, 30], [30, 29], [31, 31]]> : tensor<32x2xi64>}> : (tensor<32x128xf32>) -> tensor<32x128xf32> - %72 = "stablehlo.all_gather"(%71) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7], [8, 10], [9, 11], [12, 14], [13, 15], [16, 18], [17, 19], [20, 22], [21, 23], [24, 26], [25, 27], [28, 30], [29, 31]]> : tensor<16x2xi64>, use_global_device_ids}> : (tensor<32x128xf32>) -> tensor<32x256xf32> + %72 = "stablehlo.all_gather"(%71) <{all_gather_dim = 1 : i64, channel_handle = #stablehlo.channel_handle, replica_groups = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> : (tensor<32x128xf32>) -> tensor<32x256xf32> return %72 : tensor<32x256xf32> } } From 5f324826e4d911de3fd9593ffbd897ae03b459dc Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Fri, 17 Jul 2026 10:48:58 -0500 Subject: [PATCH 09/11] Formatting fix --- src/enzyme_ad/jax/BUILD | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/enzyme_ad/jax/BUILD b/src/enzyme_ad/jax/BUILD index c820a6b179..25437c9308 100644 --- a/src/enzyme_ad/jax/BUILD +++ b/src/enzyme_ad/jax/BUILD @@ -693,8 +693,8 @@ td_library( "@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", + "@stablehlo//:stablehlo_ops_td_files", ], ) @@ -1391,15 +1391,15 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ - ":CheckedRewrite", ":AxisDialectIncGen", ":AxisOpsIncGen", ":AxisTypeInterfacesIncGen", ":AxisTypesIncGen", + ":CheckedRewrite", ":DistributedDialectIncGen", ":DistributedInterfacesIncGen", - ":DistributedPassesIncGen", ":DistributedOpsIncGen", + ":DistributedPassesIncGen", ":DistributedTypeInterfacesIncGen", ":DistributedTypesIncGen", ":EnzymeHLOPatternsIncGen", From 2fe07dd9bbaf643a438b457eff5b5c035ff0e31a Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Fri, 17 Jul 2026 11:13:39 -0500 Subject: [PATCH 10/11] Fix TD Formatting --- src/enzyme_ad/jax/Dialect/Axis/Dialect.td | 9 +++--- src/enzyme_ad/jax/Dialect/Axis/Interfaces.td | 3 +- src/enzyme_ad/jax/Dialect/Axis/Ops.td | 15 ++++----- src/enzyme_ad/jax/Dialect/Axis/Types.td | 15 ++++----- .../jax/Dialect/Distributed/Dialect.td | 10 +++--- .../jax/Dialect/Distributed/Interfaces.td | 3 +- src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 31 ++++++++++-------- .../jax/Dialect/Distributed/Traits.h | 4 +-- .../jax/Dialect/Distributed/Types.td | 31 ++++++++++-------- .../jax/Passes/Distributed/Passes.td | 32 ++++++++++--------- .../Distributed/ShardyToDistributed.cpp | 11 ++++--- 11 files changed, 85 insertions(+), 79 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Axis/Dialect.td b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td index 03e85ee521..331458cce4 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Dialect.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td @@ -21,10 +21,11 @@ class AxisDialectType traits = []> - : Op; + : Op ; -// Marker trait for static metadata computations. This implies Pure. -def MetadataTrait : NativeOpTrait<"enzyme::axis::MetadataTrait">; -def MetadataOpTrait : TraitList<[Pure, MetadataTrait]>; + // Marker trait for static metadata computations. This implies Pure. + def MetadataTrait : NativeOpTrait<"enzyme::axis::MetadataTrait">; + def MetadataOpTrait : TraitList<[Pure, MetadataTrait]>; #endif // ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_TD diff --git a/src/enzyme_ad/jax/Dialect/Axis/Interfaces.td b/src/enzyme_ad/jax/Dialect/Axis/Interfaces.td index 10b3b0c24d..d72a8000c8 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Interfaces.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Interfaces.td @@ -17,7 +17,8 @@ def AxisTypeInterface : TypeInterface<"AxisTypeInterface"> { return $_type.getExtent(); }]>, InterfaceMethod< - [{Returns true when two SSA values of this axis type are equivalent.}], + [{Returns true when two SSA values of this axis type are equivalent. +}], "bool", "aliases", (ins "::mlir::Value":$ax1, "::mlir::Value":$ax2) > diff --git a/src/enzyme_ad/jax/Dialect/Axis/Ops.td b/src/enzyme_ad/jax/Dialect/Axis/Ops.td index a7cd5f4a0d..c41e5a49ff 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Ops.td @@ -26,16 +26,15 @@ def AxisFactorOp : AxisOp<"factor", [MetadataOpTrait, DeclareOpInterfaceMethods< let description = [{ Factors one canonical axis into leftmost-major factors with static extents. }]; - let arguments = (ins - AnyType:$axis, - DenseI32ArrayAttr:$factor_extents - ); - let results = (outs Variadic:$axis_factors); - let assemblyFormat = "$axis $factor_extents attr-dict `:` type($axis)"; - let hasVerifier = 1; +let arguments = (ins AnyType : $axis, DenseI32ArrayAttr : $factor_extents); +let results = (outs Variadic : $axis_factors); +let assemblyFormat = "$axis $factor_extents attr-dict `:` type($axis)"; +let hasVerifier = 1; } -def AxisSegmentOp : AxisOp<"segment", [MetadataOpTrait, DeclareOpInterfaceMethods]> { +def AxisSegmentOp : AxisOp<"segment", [ + MetadataOpTrait, DeclareOpInterfaceMethods +]> { let description = [{ Splits one canonical axis into left-to-right contiguous linear segments. Segment result index is value-ordered: lower result indices denote lower diff --git a/src/enzyme_ad/jax/Dialect/Axis/Types.td b/src/enzyme_ad/jax/Dialect/Axis/Types.td index 3bb87c9efd..abce510af7 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Types.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Types.td @@ -6,12 +6,11 @@ include "Interfaces.td" class AxisType traits = []> : AxisDialectType] # - traits>; + [DeclareTypeInterfaceMethods] #traits>; // Canonical axis for ranked shape dimensions. def ShapeAxisType : AxisType<"ShapeAxis", "shape_axis"> { - let parameters = (ins "::mlir::Type":$shapeType, "unsigned":$axisIndex); + let parameters = (ins "::mlir::Type" : $shapeType, "unsigned" : $axisIndex); let assemblyFormat = "`<` $shapeType `,` $axisIndex `>`"; let extraClassDeclaration = [{ unsigned getExtent() const; @@ -32,10 +31,10 @@ def AxisSegmentType : AxisDialectType<"AxisSegment", "axis_segment"> { let assemblyFormat = "`<` $axisType `,` $extent `,` $offset `>`"; } -// Virtual axis reconstructed from a list of factors. -def FactorGroupType : AxisDialectType<"FactorGroup", "factor_group"> { - let parameters = (ins "unsigned":$extent); - let assemblyFormat = "`<` $extent `>`"; -} +// Virtual axis reconstructed from a list of factors.def FactorGroupType + : AxisDialectType<"FactorGroup", "factor_group"> { + let parameters = (ins "unsigned" : $extent); + let assemblyFormat = "`<` $extent `>`"; + } #endif // ENZYME_AD_JAX_DIALECT_AXIS_TYPES_TD diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td index ea967f42e1..17da124e7d 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td @@ -7,8 +7,8 @@ include "mlir/IR/Traits.td" include "mlir/IR/OpBase.td" def DistributedDialect : Dialect { - let name = "distributed"; - let description = [{}]; + let name = "distributed"; + let description = [{}]; let cppNamespace = "::mlir::enzyme::distributed"; let useDefaultTypePrinterParser = 1; } @@ -17,9 +17,7 @@ class DistributedType traits = [] let mnemonic = type_mnemonic; } -class DistributedOp traits = []> - : Op{ - } - +class DistributedOp trai ts = + [] > : Op {} #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_DIALECT_TD \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td b/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td index a6271b28b4..1dfb2ca721 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td @@ -3,8 +3,7 @@ include "mlir/IR/OpBase.td" -def PhysicalCommAxisOpInterface : - OpInterface<"PhysicalCommAxisOpInterface"> { +def PhysicalCommAxisOpInterface : OpInterface<"PhysicalCommAxisOpInterface"> { let description = [{ Interface for ops that define a physical communication axis. }]; diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td index db2783a3ed..fce7e264bb 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -17,8 +17,9 @@ include "src/enzyme_ad/jax/Dialect/Axis/Types.td" // Marker for distributed metadata ops. Includes builtin Pure. def PureDistributedMetadata : TraitList<[MetadataOpTrait]>; -def PhysicalMeshOp : DistributedOp<"PhysicalMesh", [PureDistributedMetadata, Symbol]> { - let description = [{ +def PhysicalMeshOp + : DistributedOp<"PhysicalMesh", [PureDistributedMetadata, Symbol]> { + let description = [{ Declares a mesh of device targets along with their local communication hierarchy. }]; let arguments = (ins @@ -40,7 +41,9 @@ def GetPhysicalMeshAxesOp : DistributedOp<"GetPhysicalMeshAxes", [PureDistribute let hasVerifier = 1; } -def LogicalMeshAxesOp : DistributedOp<"LogicalMeshAxes", [PureDistributedMetadata, DeclareOpInterfaceMethods]> { +def LogicalMeshAxesOp : Distr ibutedOp<"LogicalMeshAxes", [ + PureDistributedMetadata, DeclareOpInterfaceMethods + ]> { let description = [{ Declares a new logical mesh of axes, which may or may not correspond to a physical mesh. @@ -91,7 +94,7 @@ def DistributedFunctionOp : DistributedOp<"Function", [Symbol]> { ArrayAttr:$argument_types, ArrayAttr:$return_types ); - let regions = (region SizedRegion<1>:$body); + let regions = (region SizedRegion<1> : $body); let assemblyFormat = [{ $sym_name `context` $execution_context `:` type($execution_context) @@ -173,13 +176,15 @@ def DistributedCollectiveOp : DistributedOp<"Collective", [ Output: - An AsynchHandleType wrapping the output tensor type. - Custom format: - - Reductions are specified as (%axisgroup @reduction_function), and are collected by - the parser into the reduction_groups list and reduction_functions attributes. - - Mappings are specified as (%axisgroup -> %axisgroup), and are collected by the parser into - the mapping_lhs and mapping_rhs lists. - }]; - let arguments = (ins + Custom format + : -Reductions are specified as(% axisgroup @reduction_function), + and are collected by the parser into the reduction_groups list and + reduction_functions attributes.- + Mappings are specified as(% axisgroup->% axisgroup), + and are collected by the parser into the mapping_lhs and mapping_rhs + lists. + }]; + let arguments = (ins AnyTypeOf<[AnyTensor, AnyMemRef]>:$input_object, FactorGroupType:$input_mesh, FactorGroupType:$output_mesh, @@ -189,8 +194,8 @@ def DistributedCollectiveOp : DistributedOp<"Collective", [ Variadic:$mapping_rhs, TypeAttrOf:$output_tensor_type ); - let results = (outs AsynchHandleType:$async_handle); - let assemblyFormat = [{ + let results = (outs AsynchHandleType : $async_handle); + let assemblyFormat = [{ $input_object `:` type($input_object) `on` $input_mesh `:` type($input_mesh) `to` $output_tensor_type `on` $output_mesh `:` type($output_mesh) `reduces` custom($reduction_groups, $reduction_functions) diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Traits.h b/src/enzyme_ad/jax/Dialect/Distributed/Traits.h index 95a0abf782..1b21641230 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Traits.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Traits.h @@ -5,10 +5,8 @@ namespace mlir::OpTrait::enzyme::distributed { - template -class MetadataTrait : public OpTrait::TraitBase { -}; +class MetadataTrait : public OpTrait::TraitBase {}; } // namespace mlir::OpTrait::enzyme::distributed #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_TRAITS_H \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Types.td b/src/enzyme_ad/jax/Dialect/Distributed/Types.td index 1b7ab395cb..8f11ce5e45 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Types.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Types.td @@ -5,33 +5,36 @@ include "Dialect.td" include "Interfaces.td" include "src/enzyme_ad/jax/Dialect/Axis/Interfaces.td" -class DistributedAxisType traits = []> - : DistributedType] # - traits>; +class DistributedAxisType traits = []> + : DistributedType] #traits>; // A natural axis defined by the hardware topology. -def PhysicalCommAxisType : DistributedAxisType<"PhysicalCommAxis", "physical_comm_axis"> { - let parameters = (ins "unsigned":$extent); - let assemblyFormat = "`<` $extent `>`"; +def PhysicalCommAxisType + : DistributedAxisType<"PhysicalCommAxis", "physical_comm_axis"> { + let parameters = (ins "unsigned" : $extent); + let assemblyFormat = "`<` $extent `>`"; } // Axis of a logical mesh where hardware details are abstracted away. -def LogicalMeshAxisType : DistributedAxisType<"LogicalMeshAxis", "logical_mesh_axis"> { - let parameters = (ins "unsigned":$extent); - let assemblyFormat = "`<` $extent `>`"; +def LogicalMeshAxisType + : DistributedAxisType<"LogicalMeshAxis", "logical_mesh_axis"> { + let parameters = (ins "unsigned" : $extent); + let assemblyFormat = "`<` $extent `>`"; } // A standalone axis type used to represent replication. -def ReplicationAxisType : DistributedAxisType<"ReplicationAxis", "replication_axis"> { - let parameters = (ins "unsigned":$extent); - let assemblyFormat = "`<` $extent `>`"; +def ReplicationAxisType + : DistributedAxisType<"ReplicationAxis", "replication_axis"> { + let parameters = (ins "unsigned" : $extent); + let assemblyFormat = "`<` $extent `>`"; } // Wraps a value type that may be synchronized at a later point. def AsynchHandleType : DistributedType<"AsynchHandle", "asynch_handle"> { let summary = "Handle to a value that can be synchronized later"; - let parameters = (ins "::mlir::Type":$valueType); + let parameters = (ins "::mlir::Type" : $valueType); let assemblyFormat = "`<` $valueType `>`"; } diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.td b/src/enzyme_ad/jax/Passes/Distributed/Passes.td index 7d462e17ff..edf0340b2d 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.td +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.td @@ -4,9 +4,9 @@ include "mlir/Pass/PassBase.td" def PrintShardyFunctionNamesPass - : Pass<"print-shardy-function-names", "::mlir::ModuleOp"> { - let summary = "Print names of functions that carry Shardy IR"; - let description = [{ + : Pass<"print-shardy-function-names", "::mlir::ModuleOp"> { + let summary = "Print names of functions that carry Shardy IR"; + let description = [{ Scans func.func operations and prints the symbol names of functions that either carry sdy-prefixed attributes on the function/signature or contain sdy dialect operations/attributes in their bodies. @@ -17,19 +17,21 @@ def ShardyToDistributedPass : Pass<"shardy-to-distributed", "::mlir::ModuleOp"> let summary = "Convert Shardy functions to Distributed mesh computations."; let description = [{ Converts from Shardy to Distributed IR. This pass requires all Shardy functions - to be available in the module: calling external functions is not supported. - Also requires exactly one physical mesh to be declared in the module. - Shardy functions on multiple meshes will be supported in the future, but - for now this pass will fail if multiple meshes are declared. + to be available in the module + : calling external functions is not supported + .Also requires exactly one physical mesh to be declared in the + module.Shardy functions on multiple meshes will be supported in + the future, + but for now this pass will fail if multiple meshes are declared. - Note: this pass expects local tensor sizes as input and all - layout mismatches to have explicit layout conversion ops, either - reshard or reduce. This can be achieved with some collection of existing - shardy passes: - --sdy-propagation-pipeline - --sdy-export-pipeline='enable-insert-explicit-collectives=true' - --sdy-convert-global-to-local='enable-rgv3=true' - }]; + Note : this pass expects local tensor sizes as input and all layout + mismatches to have explicit layout conversion ops, + either reshard or reduce.This can be achieved with some collection of + existing shardy passes : --sdy - + propagation - pipeline-- sdy - export - + pipeline = 'enable-insert-explicit-collectives=true' --sdy - convert - + global - to - local = 'enable-rgv3=true' +}]; } #endif // ENZYME_AD_JAX_DISTRIBUTED_PASSES diff --git a/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp b/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp index d2ca94a7f6..1a9a9ccc77 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp +++ b/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp @@ -71,13 +71,13 @@ struct ShardyToDistributedPass } // Create a new mesh computation using the modules pysical mesh - OpBuilder builder(moduleOp.getContext()); + OpBuilder builder(moduleOp.getContext()); builder.setInsertionPointAfter( - *physicalMesh); // graph region, doesn't matter + *physicalMesh); // graph region, doesn't matter distributed::MeshComputationOp meshComputation = builder.create( - moduleOp.getLoc(), builder.getStringAttr("mesh_computation"), - FlatSymbolRefAttr::get(physicalMesh->getSymNameAttr())); + moduleOp.getLoc(), builder.getStringAttr("mesh_computation"), + FlatSymbolRefAttr::get(physicalMesh->getSymNameAttr())); Region &meshComputationBody = meshComputation.getBody(); if (meshComputationBody.empty()) { meshComputationBody.emplaceBlock(); @@ -103,7 +103,8 @@ struct ShardyToDistributedPass distributed::LogicalMeshAxesOp logicalMeshAxes = builder.create( - moduleOp.getLoc(), builder.getDenseI32ArrayAttr(logicalAxisExtents)); + moduleOp.getLoc(), + builder.getDenseI32ArrayAttr(logicalAxisExtents)); llvm::DenseMap shardyToDistributedAxis; for (auto [idx, distributedAxis] : From f1ddd0e61d698875ded4a0ce094cd440076ecafa Mon Sep 17 00:00:00 2001 From: Egan Johnson Date: Fri, 17 Jul 2026 12:02:00 -0500 Subject: [PATCH 11/11] Fix broken TD files --- src/enzyme_ad/jax/Dialect/Axis/Dialect.td | 14 +- src/enzyme_ad/jax/Dialect/Axis/Types.td | 26 +-- .../jax/Dialect/Distributed/Dialect.td | 13 +- src/enzyme_ad/jax/Dialect/Distributed/Ops.td | 194 +++++++++--------- 4 files changed, 120 insertions(+), 127 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Axis/Dialect.td b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td index 331458cce4..a09fea314d 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Dialect.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td @@ -14,18 +14,16 @@ def AxisDialect : Dialect { let useDefaultTypePrinterParser = 1; } -class AxisDialectType traits = []> - : TypeDef { +class AxisDialectType traits = [ +]> : TypeDef { let mnemonic = type_mnemonic; } class AxisOp traits = []> - : Op ; + : Op; - // Marker trait for static metadata computations. This implies Pure. - def MetadataTrait : NativeOpTrait<"enzyme::axis::MetadataTrait">; - def MetadataOpTrait : TraitList<[Pure, MetadataTrait]>; +// Marker trait for static metadata computations. This implies Pure. +def MetadataTrait : NativeOpTrait<"enzyme::axis::MetadataTrait">; +def MetadataOpTrait : TraitList<[Pure, MetadataTrait]>; #endif // ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_TD diff --git a/src/enzyme_ad/jax/Dialect/Axis/Types.td b/src/enzyme_ad/jax/Dialect/Axis/Types.td index abce510af7..d728091858 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Types.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Types.td @@ -12,29 +12,31 @@ class AxisType traits = []> def ShapeAxisType : AxisType<"ShapeAxis", "shape_axis"> { let parameters = (ins "::mlir::Type" : $shapeType, "unsigned" : $axisIndex); let assemblyFormat = "`<` $shapeType `,` $axisIndex `>`"; - let extraClassDeclaration = [{ - unsigned getExtent() const; - }]; + let extraClassDeclaration = [{ unsigned getExtent() const; }]; } // Factor of a canonical axis with static extent and stride metadata. def AxisFactorType : AxisDialectType<"AxisFactor", "axis_factor"> { - let parameters = - (ins "::mlir::Type":$axisType, "unsigned":$extent, "unsigned":$stride); + let parameters = (ins "::mlir::Type" + : $axisType, "unsigned" + : $extent, "unsigned" + : $stride); let assemblyFormat = "`<` $axisType `,` $extent `,` $stride `>`"; } // Contiguous linear segment of a canonical axis with static extent and offset. def AxisSegmentType : AxisDialectType<"AxisSegment", "axis_segment"> { - let parameters = - (ins "::mlir::Type":$axisType, "unsigned":$extent, "unsigned":$offset); + let parameters = (ins "::mlir::Type" + : $axisType, "unsigned" + : $extent, "unsigned" + : $offset); let assemblyFormat = "`<` $axisType `,` $extent `,` $offset `>`"; } -// Virtual axis reconstructed from a list of factors.def FactorGroupType - : AxisDialectType<"FactorGroup", "factor_group"> { - let parameters = (ins "unsigned" : $extent); - let assemblyFormat = "`<` $extent `>`"; - } +// Virtual axis reconstructed from a list of factors. +def FactorGroupType : AxisDialectType<"FactorGroup", "factor_group"> { + let parameters = (ins "unsigned" : $extent); + let assemblyFormat = "`<` $extent `>`"; +} #endif // ENZYME_AD_JAX_DIALECT_AXIS_TYPES_TD diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td index 17da124e7d..8050e7c6e6 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td @@ -9,15 +9,16 @@ include "mlir/IR/OpBase.td" def DistributedDialect : Dialect { let name = "distributed"; let description = [{}]; - let cppNamespace = "::mlir::enzyme::distributed"; - let useDefaultTypePrinterParser = 1; + let cppNamespace = "::mlir::enzyme::distributed"; + let useDefaultTypePrinterParser = 1; } -class DistributedType traits = []> : TypeDef{ - let mnemonic = type_mnemonic; +class DistributedType traits = [ +]> : TypeDef { + let mnemonic = type_mnemonic; } -class DistributedOp trai ts = - [] > : Op {} +class DistributedOp traits = []> + : Op {} #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_DIALECT_TD \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td index fce7e264bb..1fa33b3253 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -19,129 +19,123 @@ def PureDistributedMetadata : TraitList<[MetadataOpTrait]>; def PhysicalMeshOp : DistributedOp<"PhysicalMesh", [PureDistributedMetadata, Symbol]> { - let description = [{ - Declares a mesh of device targets along with their local communication hierarchy. - }]; - let arguments = (ins - SymbolNameAttr:$sym_name, - StrAttr:$device_target, - ArrayAttr:$axes // array of physical axis types. - ); - let assemblyFormat = "$sym_name `device_target` $device_target `axes` $axes attr-dict"; - let hasVerifier = 1; + let description = [{Declares a mesh of device targets along with their local + communication hierarchy.}]; + let arguments = (ins SymbolNameAttr + : $sym_name, StrAttr + : $device_target, + ArrayAttr + : $axes // array of physical axis types. + ); + let assemblyFormat = + "$sym_name `device_target` $device_target `axes` $axes attr-dict"; + let hasVerifier = 1; } -def GetPhysicalMeshAxesOp : DistributedOp<"GetPhysicalMeshAxes", [PureDistributedMetadata]> { - let description = [{ - Resolves a module-level physical mesh symbol into one SSA axis value per mesh axis. - }]; - let arguments = (ins FlatSymbolRefAttr:$physical_mesh); - let results = (outs Variadic:$axes); - let assemblyFormat = "$physical_mesh attr-dict `:` type($axes)"; - let hasVerifier = 1; +def GetPhysicalMeshAxesOp + : DistributedOp<"GetPhysicalMeshAxes", [PureDistributedMetadata]> { + let description = + [{Resolves a module - + level physical mesh symbol into one SSA axis value per mesh axis.}]; + let arguments = (ins FlatSymbolRefAttr : $physical_mesh); + let results = (outs Variadic : $axes); + let assemblyFormat = "$physical_mesh attr-dict `:` type($axes)"; + let hasVerifier = 1; } -def LogicalMeshAxesOp : Distr ibutedOp<"LogicalMeshAxes", [ - PureDistributedMetadata, DeclareOpInterfaceMethods - ]> { - let description = [{ - Declares a new logical mesh of axes, which may or may not correspond to a - physical mesh. - }]; - let arguments = (ins - DenseI32ArrayAttr:$axis_extents - ); - let results = (outs Variadic:$axes); - let assemblyFormat = "$axis_extents attr-dict `:` type($axes)"; - let hasVerifier = 1; +def LogicalMeshAxesOp : DistributedOp<"LogicalMeshAxes", [ + PureDistributedMetadata, DeclareOpInterfaceMethods +]> { + let description = [{Declares a new logical mesh of axes, + which may or may not correspond to a physical mesh.}]; + let arguments = (ins DenseI32ArrayAttr : $axis_extents); + let results = (outs Variadic : $axes); + let assemblyFormat = "$axis_extents attr-dict `:` type($axes)"; + let hasVerifier = 1; } -def MeshComputationOp : DistributedOp<"MeshComputation", [ - HasOnlyGraphRegion, - NoTerminator, - Symbol, - SymbolTable -]> { - let description = [{ +def MeshComputationOp + : DistributedOp<"MeshComputation", + [HasOnlyGraphRegion, NoTerminator, Symbol, SymbolTable]> { + let description = [{ Represents a closed distributed program. Contains a mix non-computational SSA values (operations that represent a static computation, like axis algebra) and distributed functions. Distributed functions can call into normal single-device MLIR functions outside of this region, but all functions with internal distributed logic (communication) must be fully contained in this region, allowing us to reason about the whole call graph for optimization decisions. - }]; - let arguments = (ins - SymbolNameAttr:$sym_name, - FlatSymbolRefAttr:$physical_mesh - ); - let regions = (region SizedRegion<1>:$body); - let assemblyFormat = [{ - $sym_name `mesh` $physical_mesh $body attr-dict - }]; - let hasVerifier = 1; + }]; + let arguments = (ins SymbolNameAttr + : $sym_name, FlatSymbolRefAttr + : $physical_mesh); + let regions = (region SizedRegion<1> : $body); + let assemblyFormat = "$sym_name `mesh` $physical_mesh $body attr-dict"; + let hasVerifier = 1; } def DistributedFunctionOp : DistributedOp<"Function", [Symbol]> { - let description = [{ - Represents a distributed function. Functions have an execution context indicating what subspace of - the logical/physical mesh they are executed on, and can only be called from a matching enclosing - context. All arguments and return values are local values, rather than global tensors with - a sharding annotation. - }]; - let arguments = (ins - SymbolNameAttr:$sym_name, - FactorGroupType:$execution_context, - ArrayAttr:$argument_types, - ArrayAttr:$return_types - ); - let regions = (region SizedRegion<1> : $body); - let assemblyFormat = [{ - $sym_name `context` $execution_context - `:` type($execution_context) - `arg_types` $argument_types `ret_types` $return_types - $body attr-dict - }]; - let hasVerifier = 1; + let description = + [{Represents a distributed function.Functions have an execution context + indicating what subspace of the logical / + physical mesh they are executed on, + and can only be called from a matching enclosing context.All arguments + and return values are local values, + rather than global tensors with a sharding annotation.}]; + let arguments = (ins SymbolNameAttr + : $sym_name, FactorGroupType + : $execution_context, ArrayAttr + : $argument_types, ArrayAttr + : $return_types); + let regions = (region SizedRegion<1> : $body); + let assemblyFormat = + "$sym_name `context` $execution_context `:` type($execution_context) " + "`arg_types` $argument_types `ret_types` $return_types $body attr-dict"; + let hasVerifier = 1; } def DistributedYieldOp : DistributedOp<"DistributedYield", [Terminator]> { - let description = [{ + let description = [{ Terminator for distributed functions. Yields local values back to the caller. - }]; - let arguments = (ins Variadic:$returns); - let assemblyFormat = "$returns type($returns) attr-dict"; + }]; + let arguments = (ins Variadic : $returns); + let assemblyFormat = "$returns type($returns) attr-dict"; } -def DistributedAwait : DistributedOp<"Await", [DeclareOpInterfaceMethods]> { - let description = [{ - Awaits a value that is being computed asynchronously. Extracts the value - from an AsynchHandleType. This is a no-op if the value is already available. - }]; - let arguments = (ins AsynchHandleType:$async_handle); - let results = (outs AnyType:$value); - let assemblyFormat = "$async_handle attr-dict `:` type($async_handle) `->` type($value)"; - let hasVerifier = 1; +def DistributedAwait + : DistributedOp<"Await", + [DeclareOpInterfaceMethods]> { + let description = [{ + Awaits a value that is being computed asynchronously + .Extracts the value from an AsynchHandleType.This is a no - + op if the value is already available. + }]; + let arguments = (ins AsynchHandleType : $async_handle); + let results = (outs AnyType : $value); + let assemblyFormat = + "$async_handle attr-dict `:` type($async_handle) `->` type($value)"; + let hasVerifier = 1; } def DistributedCallOp : DistributedOp<"DistributedCall", []> { - let description = [{ - Calls a distributed function. The execution context must match the enclosing context. - }]; - let arguments = (ins + let description = + [{Calls a distributed + function.The execution context must match the enclosing context.}]; + let arguments = (ins FlatSymbolRefAttr:$callee, FactorGroupType:$replicate_over, Variadic:$arguments ); - let results = (outs Variadic:$returns); - let assemblyFormat = "$callee `replicate_over` $replicate_over $arguments attr-dict `:` type($replicate_over) `,` type($arguments) `->` type($returns)"; - let hasVerifier = 1; + let results = (outs Variadic : $returns); + let assemblyFormat = + "$callee `replicate_over` $replicate_over $arguments attr-dict `:` " + "type($replicate_over) `,` type($arguments) `->` type($returns)"; + let hasVerifier = 1; } def DistributedCollectiveOp : DistributedOp<"Collective", [ - DeclareOpInterfaceMethods, - AttrSizedOperandSegments + DeclareOpInterfaceMethods, AttrSizedOperandSegments ]> { - let description = [{ + let description = [{ An operation representing reductions and resharding along multiple axis. Does not cover collective permutes and point-to-point communication. @@ -195,15 +189,13 @@ def DistributedCollectiveOp : DistributedOp<"Collective", [ TypeAttrOf:$output_tensor_type ); let results = (outs AsynchHandleType : $async_handle); - let assemblyFormat = [{ - $input_object `:` type($input_object) `on` $input_mesh `:` type($input_mesh) - `to` $output_tensor_type `on` $output_mesh `:` type($output_mesh) - `reduces` custom($reduction_groups, $reduction_functions) - `:` type($reduction_groups) - `maps` custom($mapping_lhs, $mapping_rhs) - `:` `[` type($mapping_lhs) `]` `[` type($mapping_rhs) `]` - attr-dict - }]; - let hasVerifier = 1; + let assemblyFormat = + "$input_object `:` type($input_object) `on` $input_mesh `:` " + "type($input_mesh) `to` $output_tensor_type `on` $output_mesh `:` " + "type($output_mesh) `reduces` custom($reduction_groups, " + "$reduction_functions) `:` type($reduction_groups) `maps` " + "custom($mapping_lhs, $mapping_rhs) `:` `[` " + "type($mapping_lhs) `]` `[` type($mapping_rhs) `]` attr-dict"; + let hasVerifier = 1; } #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_OPS_TD \ No newline at end of file