diff --git a/src/enzyme_ad/jax/BUILD b/src/enzyme_ad/jax/BUILD index db469a1209..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", ], ) @@ -954,7 +954,7 @@ gentbl_cc_library( td_library( name = "DistributedPassesTdFiles", - srcs = [], + srcs = ["Passes/Distributed/Passes.td"], deps = [ "@llvm-project//mlir:PassBaseTdFiles", ], @@ -1391,15 +1391,15 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ - ":CheckedRewrite", ":AxisDialectIncGen", ":AxisOpsIncGen", ":AxisTypeInterfacesIncGen", ":AxisTypesIncGen", + ":CheckedRewrite", ":DistributedDialectIncGen", ":DistributedInterfacesIncGen", - ":DistributedPassesIncGen", ":DistributedOpsIncGen", + ":DistributedPassesIncGen", ":DistributedTypeInterfacesIncGen", ":DistributedTypesIncGen", ":EnzymeHLOPatternsIncGen", diff --git a/src/enzyme_ad/jax/Dialect/Axis/Dialect.td b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td index 03e85ee521..a09fea314d 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Dialect.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td @@ -14,9 +14,8 @@ def AxisDialect : Dialect { let useDefaultTypePrinterParser = 1; } -class AxisDialectType traits = []> - : TypeDef { +class AxisDialectType traits = [ +]> : TypeDef { let mnemonic = type_mnemonic; } 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.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..c41e5a49ff 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,25 +22,40 @@ 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. }]; +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 +]> { + let description = [{ + 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:$factor_extents + DenseI32ArrayAttr:$segment_extents ); - let results = (outs Variadic:$axis_factors); - let assemblyFormat = "$axis $factor_extents attr-dict `:` type($axis)"; + let results = (outs Variadic:$axis_segments); + let assemblyFormat = "$axis $segment_extents attr-dict `:` type($axis)"; let hasVerifier = 1; } -def AxisGroupOp : AxisOp<"group", [Pure, DeclareOpInterfaceMethods]> { +def AxisProductOp : AxisOp<"product", [MetadataOpTrait, DeclareOpInterfaceMethods]> { let description = [{ - Recombines a leftmost-major factor list into one virtual factor group. + 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..d728091858 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Types.td +++ b/src/enzyme_ad/jax/Dialect/Axis/Types.td @@ -6,28 +6,36 @@ 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; - }]; + 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 assemblyFormat = "`<` $axisType `,` $extent `,` $offset `>`"; +} + // Virtual axis reconstructed from a list of factors. def FactorGroupType : AxisDialectType<"FactorGroup", "factor_group"> { - let parameters = (ins "unsigned":$extent); + let parameters = (ins "unsigned" : $extent); let assemblyFormat = "`<` $extent `>`"; } diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp index 4a98fd7e7e..0dd19bbbd9 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp @@ -1,7 +1,27 @@ #include "Utilities.h" +#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) { @@ -21,14 +41,44 @@ 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 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(lhsTyped); + assert(succeeded(lhsProvenance) && "factor must have a provenance axis"); + if (failed(lhsProvenance)) { + return false; + } + lhsAxis = *lhsProvenance; + } + + Value rhsAxis = rhsProvenanceAxis; + if (!rhsAxis) { + auto rhsProvenance = getFactorProvenanceAxis(rhsTyped); + 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); @@ -46,17 +96,18 @@ static bool arePairwiseFactorsDisjoint(const AxisFactorType &f1, } // 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(TypedValue segment) { + return static_cast(segment.getType().getExtent()); } // Returns the defining op for a canonical axis SSA value. @@ -77,14 +128,43 @@ 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(); +} + +// 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. @@ -97,7 +177,6 @@ bool areFactorsDisjoint(ValueRange factors) { "factor disjointness uses quadratic pairwise checks"); struct FactorInfo { - AxisFactorType factorType; Value provenance; }; @@ -105,29 +184,211 @@ 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({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; +} + +// 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; +} + +// 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; + SmallVector lhsFactors; + SmallVector rhsFactors; + }; + + auto addFactorsToBuckets = [](ValueRange factors, bool isLhs, + SmallVectorImpl &grouped) { + for (Value factor : factors) { + auto factorTyped = + castTypedValue(factor, "AxisFactorType"); + auto provenance = getFactorProvenanceAxis(factorTyped); + if (failed(provenance)) { + return false; + } + + 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 grouped; + if (!addFactorsToBuckets(lhsFactors, /*isLhs=*/true, grouped) || + !addFactorsToBuckets(rhsFactors, /*isLhs=*/false, grouped)) { + return false; + } + + 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; + } + } + + return true; +} + +// Checks segment pairwise non-overlap metadata. +bool arePairwiseSegmentsDisjoint(Value lhsSegment, Value rhsSegment, + Value lhsProvenanceAxis, + Value rhsProvenanceAxis) { + 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(lhsTyped); + assert(succeeded(lhsProvenance) && "segment must have a provenance axis"); + if (failed(lhsProvenance)) { + return false; + } + lhsAxis = *lhsProvenance; + } + + Value rhsAxis = rhsProvenanceAxis; + if (!rhsAxis) { + auto rhsProvenance = getSegmentProvenanceAxis(rhsTyped); + 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 segmentTyped = + castTypedValue(segment, "AxisSegmentType"); + auto segmentType = segmentTyped.getType(); + assert(segmentType.getExtent() > 0 && "segment extent must be positive"); + + auto provenance = getSegmentProvenanceAxis(segmentTyped); + 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; } } @@ -146,23 +407,115 @@ 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) { + auto factorTyped = castTypedValue(factor, "AxisFactorType"); + + auto provenance = getFactorProvenanceAxis(factorTyped); + assert(succeeded(provenance) && "factor must have a provenance axis"); + if (failed(provenance) || *provenance != axis) { return false; } - auto provenance = - getFactorProvenanceAxis(cast>(factor)); - assert(succeeded(provenance) && "factor must have a provenance axis"); + product *= static_cast(getFactorExtent(factorTyped)); + } + + return product == + static_cast(getAxisExtent( + castTypedValue(axis, "AxisTypeInterface"))); +} + +// 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 segmentTyped = + castTypedValue(segment, "AxisSegmentType"); + auto segmentType = segmentTyped.getType(); + + auto provenance = getSegmentProvenanceAxis(segmentTyped); + assert(succeeded(provenance) && "segment must have a provenance axis"); if (failed(provenance) || *provenance != axis) { return false; } - product *= static_cast(getFactorExtent(factor)); + 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( + 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)); +} - return product == static_cast(getAxisExtent(axis)); +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 6c1c96b764..09993d7aec 100644 --- a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h @@ -6,10 +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::TypedValue segment); // Returns the defining op for a canonical axis SSA value. ::mlir::FailureOr<::mlir::Operation *> getAxisProvenanceOp(::mlir::Value axis); @@ -18,16 +21,65 @@ ::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); + +// 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, + ::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); + +// 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); +// 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/AxisOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp index f497e333f3..cc0e82ea5e 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp @@ -3,25 +3,120 @@ 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(); +} + +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(); } 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/Dialect.td b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td index ea967f42e1..8050e7c6e6 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.td @@ -7,19 +7,18 @@ include "mlir/IR/Traits.td" include "mlir/IR/OpBase.td" def DistributedDialect : Dialect { - let name = "distributed"; - let description = [{}]; - let cppNamespace = "::mlir::enzyme::distributed"; - let useDefaultTypePrinterParser = 1; + let name = "distributed"; + let description = [{}]; + 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 traits = []> - : 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/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/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/MeshComputationOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp index d6600ee47d..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 { @@ -5,11 +6,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..1fa33b3253 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Ops.td @@ -12,54 +12,190 @@ include "Dialect.td" include "Types.td" include "Interfaces.td" include "src/enzyme_ad/jax/Dialect/Axis/Dialect.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; +include "src/enzyme_ad/jax/Dialect/Axis/Types.td" // Marker for distributed metadata ops. Includes builtin Pure. def PureDistributedMetadata : TraitList<[MetadataOpTrait]>; -def GetPhysicalAxisOp : DistributedOp<"GetPhysicalAxis", [PureDistributedMetadata]> { - let description = [{ - Resolves a module-level physical axis symbol into an SSA axis value. - }]; - let arguments = (ins FlatSymbolRefAttr:$physical_axis); - let results = (outs PhysicalCommAxisType:$axis); - let assemblyFormat = "$physical_axis attr-dict `:` type($axis)"; - let hasVerifier = 1; +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; +} + +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 MeshComputationOp : DistributedOp<"MeshComputation", [ - HasOnlyGraphRegion, - NoTerminator, - Symbol, - SymbolTable +def LogicalMeshAxesOp : DistributedOp<"LogicalMeshAxes", [ + PureDistributedMetadata, DeclareOpInterfaceMethods ]> { - let description = [{ + 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 = [{ 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:$computation_name, - FlatSymbolRefAttr:$physical_mesh - ); - let regions = (region SizedRegion<1>:$body); - let assemblyFormat = [{ - $computation_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; } 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 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; +} + +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/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.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/Types.td b/src/enzyme_ad/jax/Dialect/Distributed/Types.td index 0997ab68e9..8f11ce5e45 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Types.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Types.td @@ -5,44 +5,37 @@ include "Dialect.td" include "Interfaces.td" include "src/enzyme_ad/jax/Dialect/Axis/Interfaces.td" -class AxisType traits = []> - : DistributedType] # - traits>; +class DistributedAxisType traits = []> + : DistributedType] #traits>; // A natural axis defined by the hardware topology. -def PhysicalCommAxisType : AxisType<"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 : AxisType<"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 : AxisType<"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 `>`"; } - -// 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"; - let parameters = (ins "::mlir::Type":$valueType); + let parameters = (ins "::mlir::Type" : $valueType); 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.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp index 7aa04ef732..6107b5f007 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.cpp @@ -1 +1,94 @@ -#include "Utilities.h" \ No newline at end of file +#include "Utilities.h" + +namespace mlir::enzyme::distributed { + +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 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() + : 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; +} + +::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 39673deca2..37fc04f257 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Utilities.h @@ -6,6 +6,15 @@ 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); + +// 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) { @@ -13,22 +22,30 @@ ::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) { + + if (auto *op = lookupSymbolInEnclosingScopes(from, symRef)) { + if (auto typedOp = llvm::dyn_cast(op)) { + return typedOp; + } return ::mlir::failure(); } - return typedOp; + + return ::mlir::failure(); } -using ::mlir::enzyme::axis::areFactorsComplete; -using ::mlir::enzyme::axis::areFactorsDisjoint; -using ::mlir::enzyme::axis::getAxisExtent; -using ::mlir::enzyme::axis::getFactorExtent; -using ::mlir::enzyme::axis::getFactorProvenanceAxis; +// 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); + +// 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 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 deleted file mode 100644 index 05c82777d3..0000000000 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.cpp +++ /dev/null @@ -1,7 +0,0 @@ -#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..5e6b13040e 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.h +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.h @@ -7,7 +7,11 @@ namespace mlir { namespace enzyme { namespace distributed { -void registerdistributedPasses(); +#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" } // namespace distributed } // namespace enzyme diff --git a/src/enzyme_ad/jax/Passes/Distributed/Passes.td b/src/enzyme_ad/jax/Passes/Distributed/Passes.td index 96c8254b61..edf0340b2d 100644 --- a/src/enzyme_ad/jax/Passes/Distributed/Passes.td +++ b/src/enzyme_ad/jax/Passes/Distributed/Passes.td @@ -3,8 +3,35 @@ 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. + }]; +} + +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 = '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 new file mode 100644 index 0000000000..1a9a9ccc77 --- /dev/null +++ b/src/enzyme_ad/jax/Passes/Distributed/ShardyToDistributed.cpp @@ -0,0 +1,123 @@ +#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" + +#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 moduleOp = getOperation(); + + FailureOr physicalMesh = + distributed::findUniquePhysicalMesh(moduleOp); + if (failed(physicalMesh)) { + signalPassFailure(); + return; + } + + const FindShardyFunctionsAnalysis &analysis = + getAnalysis(); + if (!analysis.isValid()) { + signalPassFailure(); + return; + } + if (analysis.getShardyFunctions().empty()) { + moduleOp.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(moduleOp.getContext()); + builder.setInsertionPointAfter( + *physicalMesh); // graph region, doesn't matter + distributed::MeshComputationOp meshComputation = + builder.create( + moduleOp.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()) { + moduleOp.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( + moduleOp.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/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 } 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/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 diff --git a/test/lit_tests/distributed/roundtrip.mlir b/test/lit_tests/distributed/roundtrip.mlir new file mode 100644 index 0000000000..085baacd90 --- /dev/null +++ b/test/lit_tests/distributed/roundtrip.mlir @@ -0,0 +1,94 @@ +// RUN: enzymexlamlir-opt --split-input-file %s | FileCheck %s + +// 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> + %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.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 + +// ----- + +// 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>] + + 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> 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..196b465b64 --- /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 = #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 = #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 = #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 = #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 = #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 = #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> + %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 = #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> + %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 = #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<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 = #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<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 = #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<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 = #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 = #stablehlo.replica_group_mesh_axes]>, use_global_device_ids}> : (tensor<32x128xf32>) -> tensor<32x256xf32> + return %72 : tensor<32x256xf32> + } +} + 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> -}