diff --git a/src/enzyme_ad/jax/BUILD b/src/enzyme_ad/jax/BUILD index 56fab16aa5..8229dc50b0 100644 --- a/src/enzyme_ad/jax/BUILD +++ b/src/enzyme_ad/jax/BUILD @@ -577,6 +577,105 @@ td_library( ], ) +td_library( + name = "AxisDialectFiles", + srcs = [ + "Dialect/Axis/Dialect.td", + "Dialect/Axis/Interfaces.td", + "Dialect/Axis/Ops.td", + "Dialect/Axis/Types.td", + ], + includes = ["."], + deps = [ + "@llvm-project//mlir:BuiltinDialectTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "AxisDialectIncGen", + tbl_outs = [ + ( + [ + "-gen-dialect-decls", + "-dialect=axis", + ], + "Dialect/Axis/AxisDialect.h.inc", + ), + ( + [ + "-gen-dialect-defs", + "-dialect=axis", + ], + "Dialect/Axis/AxisDialect.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/Axis/Dialect.td", + deps = [ + ":AxisDialectFiles", + ], +) + +gentbl_cc_library( + name = "AxisOpsIncGen", + tbl_outs = [ + ( + ["-gen-op-decls"], + "Dialect/Axis/AxisOps.h.inc", + ), + ( + ["-gen-op-defs"], + "Dialect/Axis/AxisOps.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/Axis/Ops.td", + deps = [ + ":AxisDialectFiles", + "@llvm-project//mlir:InferTypeOpInterface", + ], +) + +gentbl_cc_library( + name = "AxisTypesIncGen", + tbl_outs = [ + ( + ["--gen-typedef-decls"], + "Dialect/Axis/AxisTypes.h.inc", + ), + ( + ["--gen-typedef-defs"], + "Dialect/Axis/AxisTypes.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/Axis/Types.td", + deps = [ + ":AxisDialectFiles", + ], +) + +gentbl_cc_library( + name = "AxisTypeInterfacesIncGen", + tbl_outs = [ + ( + ["--gen-type-interface-decls"], + "Dialect/Axis/AxisTypeInterfaces.h.inc", + ), + ( + ["--gen-type-interface-defs"], + "Dialect/Axis/AxisTypeInterfaces.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/Axis/Interfaces.td", + deps = [ + ":AxisDialectFiles", + ], +) + td_library( name = "DistributedDialectFiles", srcs = [ @@ -585,9 +684,17 @@ td_library( "Dialect/Distributed/Ops.td", "Dialect/Distributed/Types.td", ], + includes = [ + ".", + "../../../external/shardy", + ], deps = [ + ":AxisDialectFiles", "@llvm-project//mlir:BuiltinDialectTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", "@llvm-project//mlir:OpBaseTdFiles", + "@shardy//shardy/dialect/sdy/ir:sdy_td_files", + "@stablehlo//:stablehlo_ops_td_files", ], ) @@ -632,6 +739,8 @@ gentbl_cc_library( td_file = "Dialect/Distributed/Ops.td", deps = [ ":DistributedDialectFiles", + "@llvm-project//mlir:InferTypeOpInterface", + "@shardy//shardy/dialect/sdy/ir:sdy_td_files", ], ) @@ -673,6 +782,25 @@ gentbl_cc_library( ], ) +gentbl_cc_library( + name = "DistributedTypeInterfacesIncGen", + tbl_outs = [ + ( + ["--gen-type-interface-decls"], + "Dialect/Distributed/DistributedTypeInterfaces.h.inc", + ), + ( + ["--gen-type-interface-defs"], + "Dialect/Distributed/DistributedTypeInterfaces.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Dialect/Distributed/Interfaces.td", + deps = [ + ":DistributedDialectFiles", + ], +) + td_library( name = "TesseraDialectTdFiles", srcs = [ @@ -824,6 +952,33 @@ gentbl_cc_library( ], ) +td_library( + name = "DistributedPassesTdFiles", + srcs = [], + deps = [ + "@llvm-project//mlir:PassBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "DistributedPassesIncGen", + tbl_outs = [ + ( + [ + "-gen-pass-decls", + "-name=distributed", + ], + "Passes/Distributed/Passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Passes/Distributed/Passes.td", + deps = [ + ":DistributedPassesTdFiles", + "@llvm-project//mlir:FunctionInterfacesTdFiles", + ], +) + cc_library( name = "CheckedRewrite", hdrs = ["CheckedRewrite.h"], @@ -1195,8 +1350,10 @@ cc_library( "Analysis/*.cpp", "Implementations/*.cpp", "Passes/*.cpp", + "Passes/Distributed/*.cpp", "Passes/Tessera/*.cpp", "Dialect/*.cpp", + "Dialect/Axis/*.cpp", "Dialect/Distributed/*.cpp", "Dialect/Tessera/*.cpp", "Dialect/Perfify/*.cpp", @@ -1213,8 +1370,10 @@ cc_library( "Analysis/*.h", "Implementations/*.h", "Passes/*.h", + "Passes/Distributed/*.h", "Passes/Tessera/*.h", "Dialect/*.h", + "Dialect/Axis/*.h", "Dialect/Distributed/*.h", "Dialect/Tessera/*.h", "Dialect/Perfify/*.h", @@ -1232,9 +1391,16 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ + ":AxisDialectIncGen", + ":AxisOpsIncGen", + ":AxisTypeInterfacesIncGen", + ":AxisTypesIncGen", ":CheckedRewrite", ":DistributedDialectIncGen", + ":DistributedInterfacesIncGen", ":DistributedOpsIncGen", + ":DistributedPassesIncGen", + ":DistributedTypeInterfacesIncGen", ":DistributedTypesIncGen", ":EnzymeHLOPatternsIncGen", ":EnzymeXLAAttrsIncGen", diff --git a/src/enzyme_ad/jax/Dialect/Axis/Dialect.cpp b/src/enzyme_ad/jax/Dialect/Axis/Dialect.cpp new file mode 100644 index 0000000000..191971d79a --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Dialect.cpp @@ -0,0 +1,22 @@ +#include "Dialect.h" + +#include "llvm/ADT/TypeSwitch.h" + +// Include the .cpp.inc files +#include "src/enzyme_ad/jax/Dialect/Axis/AxisDialect.cpp.inc" + +#define GET_TYPEDEF_CLASSES +#include "src/enzyme_ad/jax/Dialect/Axis/AxisTypes.cpp.inc" + +#include "src/enzyme_ad/jax/Dialect/Axis/AxisTypeInterfaces.cpp.inc" + +void mlir::enzyme::axis::AxisDialect::initialize() { + addTypes< +#define GET_TYPEDEF_LIST +#include "src/enzyme_ad/jax/Dialect/Axis/AxisTypes.cpp.inc" + >(); + addOperations< +#define GET_OP_LIST +#include "src/enzyme_ad/jax/Dialect/Axis/AxisOps.cpp.inc" + >(); +} diff --git a/src/enzyme_ad/jax/Dialect/Axis/Dialect.h b/src/enzyme_ad/jax/Dialect/Axis/Dialect.h new file mode 100644 index 0000000000..65ed073519 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Dialect.h @@ -0,0 +1,27 @@ +#ifndef ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_H +#define ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_H + +#include "mlir/Bytecode/BytecodeOpInterface.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Types.h" +#include "mlir/Interfaces/InferTypeOpInterface.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Support/LLVM.h" + +#include "Traits.h" + +// Include the dialect +#include "src/enzyme_ad/jax/Dialect/Axis/AxisDialect.h.inc" +// Type interfaces +#include "src/enzyme_ad/jax/Dialect/Axis/AxisTypeInterfaces.h.inc" +// Types +#define GET_TYPEDEF_CLASSES +#include "src/enzyme_ad/jax/Dialect/Axis/AxisTypes.h.inc" +// Ops +#define GET_OP_CLASSES +#include "src/enzyme_ad/jax/Dialect/Axis/AxisOps.h.inc" + +#endif // ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_H diff --git a/src/enzyme_ad/jax/Dialect/Axis/Dialect.td b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td new file mode 100644 index 0000000000..a09fea314d --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Dialect.td @@ -0,0 +1,29 @@ +#ifndef ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_TD +#define ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_TD + +include "mlir/IR/AttrTypeBase.td" +include "mlir/IR/DialectBase.td" +include "mlir/IR/OpBase.td" +include "mlir/IR/Traits.td" +include "mlir/Interfaces/SideEffectInterfaces.td" + +def AxisDialect : Dialect { + let name = "axis"; + let description = [{SSA algebra over canonical axes and axis factors.}]; + let cppNamespace = "::mlir::enzyme::axis"; + let useDefaultTypePrinterParser = 1; +} + +class AxisDialectType traits = [ +]> : TypeDef { + let mnemonic = type_mnemonic; +} + +class AxisOp traits = []> + : Op; + +// Marker trait for static metadata computations. This implies Pure. +def MetadataTrait : NativeOpTrait<"enzyme::axis::MetadataTrait">; +def MetadataOpTrait : TraitList<[Pure, MetadataTrait]>; + +#endif // ENZYME_AD_JAX_DIALECT_AXIS_DIALECT_TD diff --git a/src/enzyme_ad/jax/Dialect/Axis/Interfaces.td b/src/enzyme_ad/jax/Dialect/Axis/Interfaces.td new file mode 100644 index 0000000000..d72a8000c8 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Interfaces.td @@ -0,0 +1,28 @@ +#ifndef ENZYME_AD_JAX_DIALECT_AXIS_INTERFACES_TD +#define ENZYME_AD_JAX_DIALECT_AXIS_INTERFACES_TD + +include "mlir/IR/OpBase.td" + +def AxisTypeInterface : TypeInterface<"AxisTypeInterface"> { + let cppNamespace = "::mlir::enzyme::axis"; + let description = [{ + Interface for canonical axis types with statically decidable equivalence. + }]; + + let methods = [ + InterfaceMethod< + [{Returns the compile-time extent of this canonical axis.}], + "unsigned", "extent", (ins), "", + [{ + return $_type.getExtent(); + }]>, + InterfaceMethod< + [{Returns true when two SSA values of this axis type are equivalent. +}], + "bool", "aliases", + (ins "::mlir::Value":$ax1, "::mlir::Value":$ax2) + > + ]; +} + +#endif // ENZYME_AD_JAX_DIALECT_AXIS_INTERFACES_TD diff --git a/src/enzyme_ad/jax/Dialect/Axis/Ops.cpp b/src/enzyme_ad/jax/Dialect/Axis/Ops.cpp new file mode 100644 index 0000000000..0df959618f --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Ops.cpp @@ -0,0 +1,204 @@ +#include "Dialect.h" +#include "Utilities.h" + +namespace mlir::enzyme::axis { + +namespace { + +static LogicalResult verifyFactorExtents(ArrayRef extents, + unsigned sourceExtent, Operation *op) { + uint64_t product = 1; + for (int32_t extent : extents) { + if (extent <= 0) { + return op->emitOpError() << "requires all factor extents to be > 0"; + } + product *= static_cast(extent); + } + if (product != sourceExtent) { + return op->emitOpError() << "requires product(factor_extents) == axis " + "extent (" + << product << " != " << sourceExtent << ")"; + } + return success(); +} + +static void computeMajorToMinorStrides(ArrayRef extents, + SmallVectorImpl &strides) { + // Leftmost factors are most major. + strides.resize(extents.size()); + unsigned running = 1; + for (int i = static_cast(extents.size()) - 1; i >= 0; --i) { + strides[i] = running; + running *= static_cast(extents[i]); + } +} + +} // namespace + +LogicalResult AxisGetAxisOp::verify() { + auto shaped = dyn_cast(getShapeType()); + if (!shaped) { + return emitOpError() << "requires shape_type to be a shaped type"; + } + if (!shaped.hasRank()) { + return emitOpError() << "requires shape_type to be ranked"; + } + + int64_t axisIndex = getAxisIndex(); + if (axisIndex < 0 || axisIndex >= shaped.getRank()) { + return emitOpError() << "requires axis_index in [0, rank), got " + << axisIndex << " for rank " << shaped.getRank(); + } + if (shaped.isDynamicDim(static_cast(axisIndex))) { + return emitOpError() << "requires static shape dimension at axis_index " + << axisIndex; + } + + auto expectedType = ShapeAxisType::get(getContext(), getShapeType(), + static_cast(axisIndex)); + if (getAxis().getType() != expectedType) { + return emitOpError() + << "requires result type to match shape_type and axis_index"; + } + + return success(); +} + +LogicalResult AxisGetAxisOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + AxisGetAxisOpAdaptor adaptor(operands, attributes, properties, regions); + int64_t axisIndex = adaptor.getAxisIndex(); + if (axisIndex < 0) { + if (location) { + mlir::emitError(*location) << "requires non-negative axis_index"; + } + return failure(); + } + + inferredReturnTypes.push_back(ShapeAxisType::get( + context, adaptor.getShapeType(), static_cast(axisIndex))); + return success(); +} + +LogicalResult AxisFactorOp::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 factorExtents = getFactorExtents(); + + if (failed(verifyFactorExtents(factorExtents, axisIface.extent(), + getOperation()))) { + return failure(); + } + + if (getAxisFactors().size() != factorExtents.size()) { + return emitOpError() << "requires number of results to match number of " + "factor extents"; + } + + SmallVector expectedStrides; + computeMajorToMinorStrides(factorExtents, expectedStrides); + + for (auto [idx, axisFactorVal] : llvm::enumerate(getAxisFactors())) { + auto axisFactorType = dyn_cast(axisFactorVal.getType()); + if (!axisFactorType) { + return emitOpError() << "requires all results to have AxisFactorType"; + } + if (axisFactorType.getAxisType() != getAxis().getType()) { + return emitOpError() << "requires result #" << idx + << " to have base axis type equal to operand type"; + } + if (axisFactorType.getExtent() != + static_cast(factorExtents[idx])) { + return emitOpError() << "requires result #" << idx + << " extent to match factor extent"; + } + if (axisFactorType.getStride() != expectedStrides[idx]) { + return emitOpError() << "requires result #" << idx + << " stride to follow leftmost-major convention"; + } + } + + return success(); +} + +LogicalResult AxisFactorOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + AxisFactorOpAdaptor adaptor(operands, attributes, properties, regions); + ArrayRef factorExtents = adaptor.getFactorExtents(); + + SmallVector strides; + computeMajorToMinorStrides(factorExtents, strides); + + inferredReturnTypes.reserve(factorExtents.size()); + Type axisType = adaptor.getAxis().getType(); + for (auto [idx, factorExtent] : llvm::enumerate(factorExtents)) { + inferredReturnTypes.push_back(AxisFactorType::get( + context, axisType, static_cast(factorExtent), strides[idx])); + } + return success(); +} + +LogicalResult AxisGroupOp::verify() { + uint64_t extentProduct = 1; + for (Value factor : getFactors()) { + if (!isa(factor)) { + return emitOpError() << "requires factor operands to be op results"; + } + if (!factor.getDefiningOp()) { + return emitOpError() + << "requires factor operands to be produced by axis.factor"; + } + + auto factorType = dyn_cast(factor.getType()); + if (!factorType) { + return emitOpError() << "requires all operands to be AxisFactorType"; + } + extentProduct *= static_cast(factorType.getExtent()); + } + + if (getGroup().getType().getExtent() != extentProduct) { + return emitOpError() + << "requires group extent to equal product of factor extents"; + } + return success(); +} + +LogicalResult AxisGroupOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, PropertyRef properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + AxisGroupOpAdaptor adaptor(operands, attributes, properties, regions); + uint64_t extentProduct = 1; + for (Value factor : adaptor.getFactors()) { + auto factorType = dyn_cast(factor.getType()); + if (!factorType) { + if (location) + mlir::emitError(*location) + << "requires all operands to be AxisFactorType"; + return failure(); + } + extentProduct *= static_cast(factorType.getExtent()); + } + + inferredReturnTypes.push_back( + FactorGroupType::get(context, static_cast(extentProduct))); + return success(); +} + +} // namespace mlir::enzyme::axis + +#define GET_OP_CLASSES +#include "src/enzyme_ad/jax/Dialect/Axis/AxisOps.cpp.inc" diff --git a/src/enzyme_ad/jax/Dialect/Axis/Ops.td b/src/enzyme_ad/jax/Dialect/Axis/Ops.td new file mode 100644 index 0000000000..f3f571e957 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Ops.td @@ -0,0 +1,46 @@ +#ifndef ENZYME_AD_JAX_DIALECT_AXIS_OPS_TD +#define ENZYME_AD_JAX_DIALECT_AXIS_OPS_TD + +include "mlir/IR/BuiltinAttributes.td" +include "mlir/Interfaces/InferTypeOpInterface.td" +include "mlir/Interfaces/SideEffectInterfaces.td" + +include "Dialect.td" +include "Interfaces.td" +include "Types.td" + +def AxisGetAxisOp : AxisOp<"getaxis", [Pure, DeclareOpInterfaceMethods]> { + let description = [{ + Materializes a canonical ShapeAxis value for one ranked, static shape index. + }]; + let arguments = (ins + TypeAttr:$shape_type, + I32Attr:$axis_index + ); + let results = (outs ShapeAxisType:$axis); + let assemblyFormat = "$shape_type $axis_index attr-dict"; + let hasVerifier = 1; +} + +def AxisFactorOp : AxisOp<"factor", [Pure, 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 AxisGroupOp + : AxisOp<"group", [Pure, DeclareOpInterfaceMethods]> { + let description = [{ + Recombines a leftmost-major factor list into one virtual factor group. + }]; + let arguments = (ins Variadic:$factors); + let results = (outs FactorGroupType:$group); + let assemblyFormat = "$factors attr-dict `:` type($factors)"; + let hasVerifier = 1; +} + +#endif // ENZYME_AD_JAX_DIALECT_AXIS_OPS_TD diff --git a/src/enzyme_ad/jax/Dialect/Axis/Traits.h b/src/enzyme_ad/jax/Dialect/Axis/Traits.h new file mode 100644 index 0000000000..3cf22ab572 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Traits.h @@ -0,0 +1,13 @@ +#ifndef ENZYME_AD_JAX_DIALECT_AXIS_TRAITS_H +#define ENZYME_AD_JAX_DIALECT_AXIS_TRAITS_H + +#include "mlir/IR/OpDefinition.h" + +namespace mlir::OpTrait::enzyme::axis { + +template +class MetadataTrait : public OpTrait::TraitBase {}; + +} // namespace mlir::OpTrait::enzyme::axis + +#endif // ENZYME_AD_JAX_DIALECT_AXIS_TRAITS_H diff --git a/src/enzyme_ad/jax/Dialect/Axis/Types.cpp b/src/enzyme_ad/jax/Dialect/Axis/Types.cpp new file mode 100644 index 0000000000..69c8a1eeb2 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Types.cpp @@ -0,0 +1,28 @@ +#include "Dialect.h" + +#include "llvm/Support/ErrorHandling.h" + +namespace mlir::enzyme::axis { + +unsigned ShapeAxisType::getExtent() const { + auto shaped = ::mlir::dyn_cast<::mlir::ShapedType>(getShapeType()); + assert(shaped && "ShapeAxisType expects a shaped type"); + assert(shaped.hasRank() && "ShapeAxisType expects a ranked shaped type"); + assert(getAxisIndex() < shaped.getRank() && "axis index out of range"); + assert(!shaped.isDynamicDim(getAxisIndex()) && + "dynamic shape dimensions are not yet supported"); + return static_cast(shaped.getDimSize(getAxisIndex())); +} + +// Shape-axis equivalence is currently based on typed index-space coordinates. +bool ShapeAxisType::aliases(Value ax1, Value ax2) const { + (void)ax1; + auto other = dyn_cast(ax2.getType()); + assert(other && "ShapeAxisType aliasing requires shape axis values"); + // TODO: include object/value provenance once index-space identity + // is represented beyond the current type payload. + return getAxisIndex() == other.getAxisIndex() && + getShapeType() == other.getShapeType(); +} + +} // namespace mlir::enzyme::axis diff --git a/src/enzyme_ad/jax/Dialect/Axis/Types.td b/src/enzyme_ad/jax/Dialect/Axis/Types.td new file mode 100644 index 0000000000..95b642ac33 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Types.td @@ -0,0 +1,33 @@ +#ifndef ENZYME_AD_JAX_DIALECT_AXIS_TYPES_TD +#define ENZYME_AD_JAX_DIALECT_AXIS_TYPES_TD + +include "Dialect.td" +include "Interfaces.td" + +class AxisType traits = []> + : AxisDialectType] #traits>; + +// Canonical axis for ranked shape dimensions. +def ShapeAxisType : AxisType<"ShapeAxis", "shape_axis"> { + let parameters = (ins "::mlir::Type" : $shapeType, "unsigned" : $axisIndex); + let assemblyFormat = "`<` $shapeType `,` $axisIndex `>`"; + 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 assemblyFormat = "`<` $axisType `,` $extent `,` $stride `>`"; +} + +// Virtual axis reconstructed from a list of factors. +def FactorGroupType : AxisDialectType<"FactorGroup", "factor_group"> { + let parameters = (ins "unsigned":$extent); + let assemblyFormat = "`<` $extent `>`"; +} + +#endif // ENZYME_AD_JAX_DIALECT_AXIS_TYPES_TD diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp new file mode 100644 index 0000000000..4a98fd7e7e --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.cpp @@ -0,0 +1,168 @@ +#include "Utilities.h" + +namespace mlir::enzyme::axis { + +// Dispatches alias checks for canonical axes. Canonical axes are +// either equivalent or wholly disjoint. +static bool areAxesEquivalent(Value lhs, Value rhs) { + if (!isa(lhs.getType()) || + !isa(rhs.getType())) { + return false; + } + if (lhs.getType().getTypeID() != rhs.getType().getTypeID()) { + return false; + } + auto lhsAxisIface = dyn_cast(lhs.getType()); + assert(lhsAxisIface && "axis value type must implement AxisTypeInterface"); + if (!lhsAxisIface) { + return false; + } + return lhsAxisIface.aliases(lhs, 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(); + if (majorStride < minorStride) { + std::swap(majorStride, minorStride); + std::swap(majorExtent, minorExtent); + } + + (void)majorExtent; + unsigned minorSpan = minorStride * minorExtent; + if (majorStride < minorSpan) { + return false; + } + if ((majorStride % minorSpan) != 0) { + return false; + } + return true; +} + +// 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()); +} + +// 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()); +} + +// Returns the defining op for a canonical axis SSA value. +FailureOr getAxisProvenanceOp(Value axis) { + auto result = dyn_cast(axis); + if (!result) { + return failure(); + } + return result.getOwner(); +} + +// Returns the defining source axis for a factor value. +FailureOr getFactorProvenanceAxis(TypedValue factor) { + if (auto axisFactor = factor.getDefiningOp()) { + return axisFactor.getAxis(); + } + + return failure(); +} + +// Returns the factor list used to build a factor-group SSA value. +FailureOr +getGroupProvenanceFactors(TypedValue factorGroup) { + auto groupOp = factorGroup.getDefiningOp(); + if (!groupOp) { + return failure(); + } + return groupOp.getFactors(); +} + +// Checks factor compatibility and pairwise non-overlap metadata. +bool areFactorsDisjoint(ValueRange factors) { + if (factors.empty()) { + return true; + } + + assert(factors.size() < 100 && + "factor disjointness uses quadratic pairwise checks"); + + struct FactorInfo { + AxisFactorType factorType; + Value provenance; + }; + + // Cache provenance once so pairwise checks remain pure and cheap. + 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; + } + assert(factorType.getExtent() > 0 && "factor extent must be positive"); + assert(factorType.getStride() > 0 && "factor stride must be positive"); + + auto provenance = + getFactorProvenanceAxis(cast>(factor)); + assert(succeeded(provenance) && "factor must have a provenance axis"); + cachedFactors.push_back({factorType, *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]; + + if (!areAxesEquivalent(lhsAxis, rhsAxis)) { + continue; + } + if (!arePairwiseFactorsDisjoint(lhsType, rhsType)) { + return false; + } + } + } + + return true; +} + +// Checks that factors reconstruct the full source axis extent. +bool areFactorsComplete(Value axis, ValueRange factors) { + if (factors.empty() || !areFactorsDisjoint(factors)) { + return false; + } + + // Given disjointness, we are complete iff all factors belong to the target + // axis and their extents cover the whole source-axis extent. + uint64_t product = 1; + for (Value factor : factors) { + auto factorType = dyn_cast(factor.getType()); + assert(factorType && "factor value must have AxisFactorType"); + if (!factorType) { + return false; + } + + auto provenance = + getFactorProvenanceAxis(cast>(factor)); + assert(succeeded(provenance) && "factor must have a provenance axis"); + if (failed(provenance) || *provenance != axis) { + return false; + } + + product *= static_cast(getFactorExtent(factor)); + } + + return product == static_cast(getAxisExtent(axis)); +} + +} // namespace mlir::enzyme::axis diff --git a/src/enzyme_ad/jax/Dialect/Axis/Utilities.h b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h new file mode 100644 index 0000000000..6c1c96b764 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Axis/Utilities.h @@ -0,0 +1,33 @@ +#ifndef ENZYME_AD_JAX_DIALECT_AXIS_UTILITIES_H +#define ENZYME_AD_JAX_DIALECT_AXIS_UTILITIES_H + +#include "Dialect.h" + +namespace mlir::enzyme::axis { + +// Returns the static extent for any canonical axis SSA value. +int getAxisExtent(::mlir::Value axis); + +// Returns the static extent for any factor SSA value. +int getFactorExtent(::mlir::Value factor); + +// Returns the defining op for a canonical axis SSA value. +::mlir::FailureOr<::mlir::Operation *> getAxisProvenanceOp(::mlir::Value axis); + +// Resolves the source canonical axis used to produce a factor value. +::mlir::FailureOr<::mlir::Value> +getFactorProvenanceAxis(::mlir::TypedValue factor); + +// Returns the factor list used to build a factor-group SSA value. +::mlir::FailureOr<::mlir::ValueRange> +getGroupProvenanceFactors(::mlir::TypedValue factorGroup); + +// Checks that factors are pairwise non-overlapping for one source axis. +bool areFactorsDisjoint(::mlir::ValueRange factors); + +// Checks that factors cover an axis exactly and therefore are disjoint. +bool areFactorsComplete(::mlir::Value axis, ::mlir::ValueRange factors); + +} // 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 new file mode 100644 index 0000000000..f497e333f3 --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Distributed/AxisOps.cpp @@ -0,0 +1,28 @@ +#include "Dialect.h" +#include "Utilities.h" + +namespace mlir::enzyme::distributed { + +LogicalResult GetPhysicalAxisOp::verify() { + FailureOr physicalAxis = + resolveSymbolOpFromAttr( + *this, getPhysicalAxisAttr()); + if (failed(physicalAxis)) { + return emitOpError() << "references unknown physical axis symbol " + << getPhysicalAxisAttr(); + } + + 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(); +} + +} // namespace mlir::enzyme::distributed diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp index 343d327678..969b513e6e 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.cpp @@ -8,6 +8,8 @@ #define GET_TYPEDEF_CLASSES #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedTypes.cpp.inc" +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedInterfaces.cpp.inc" + // Initialize the dialect void mlir::enzyme::distributed::DistributedDialect::initialize() { addTypes< diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h index e8277364c6..4183533952 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h +++ b/src/enzyme_ad/jax/Dialect/Distributed/Dialect.h @@ -2,6 +2,7 @@ #define ENZYME_AD_JAX_DIALECT_DISTRIBUTED_DIALECT_H #include "mlir/Bytecode/BytecodeOpInterface.h" +#include "mlir/Dialect/Quant/IR/QuantTypes.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Dialect.h" #include "mlir/IR/DialectImplementation.h" @@ -9,14 +10,39 @@ #include "mlir/IR/Region.h" #include "mlir/IR/SymbolTable.h" #include "mlir/IR/Types.h" +#include "mlir/Interfaces/InferTypeOpInterface.h" +#include "mlir/Support/LLVM.h" +#include "stablehlo/dialect/Base.h" + +#include "shardy/dialect/sdy/ir/dialect.h" +#include "src/enzyme_ad/jax/Dialect/Axis/Dialect.h" + +// Need to declare the TypedOpResult struct before including the generated +// interface declarations +namespace mlir::enzyme::distributed { +template struct TypedOpResult { + ::mlir::OpResult value; + TypedOpResult(mlir::Value value) : value(llvm::cast(value)) { + assert(isa(value.getType()) && "value must have the correct type"); + } + TypedOpResult(::mlir::OpResult value) : value(value) { + assert(isa(value.getType()) && "value must have the correct type"); + } + operator mlir::OpResult() const { return value; } + + mlir::OpResult asOpResult() const { return value; } +}; +} // namespace mlir::enzyme::distributed // Include the dialect #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedDialect.h.inc" -// Traits and interfaces +// Traits #include "Traits.h" // Types #define GET_TYPEDEF_CLASSES #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedTypes.h.inc" +// Interfaces +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedInterfaces.h.inc" // Operations #define GET_OP_CLASSES #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedOps.h.inc" diff --git a/src/enzyme_ad/jax/Dialect/Distributed/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/Interfaces.td b/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td index 8fd74fd4e2..1dfb2ca721 100644 --- a/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td +++ b/src/enzyme_ad/jax/Dialect/Distributed/Interfaces.td @@ -3,7 +3,65 @@ include "mlir/IR/OpBase.td" -def DeviceDefTrait : NativeOpTrait<"enzyme::distributed::DeviceDefTrait">; -def ChannelDefTrait : NativeOpTrait<"enzyme::distributed::ChannelDefTrait">; +def PhysicalCommAxisOpInterface : OpInterface<"PhysicalCommAxisOpInterface"> { + let description = [{ + Interface for ops that define a physical communication axis. + }]; + + let methods = [ + InterfaceMethod< + [{Returns the extent of this physical axis in number of devices.}], + "int64_t", "getPhysicalAxisExtent", (ins)> + ]; +} + +def HandleDeclarationOpInterface : + OpInterface<"HandleDeclarationOpInterface"> { + let methods = [ + InterfaceMethod< + [{Returns the SSA value representing an event handle}], + "mlir::Value", "getHandle", (ins) + > + ]; +} + +def EventOpInterface : + OpInterface<"EventOpInterface"> { + let description = [{ + An op representing part or all of an event. Not all ops in a lane + need to be events, but any op spanning multiple lanes must be. + Events are ordered by a happens-before relationship, determined by: + - Explicit happensAfter(): a list of handles whose events must complete before this event can start. + - Implicit happens after: No interface needed, this is determined by the order of operations within a + basic block and by SSA dependencies. + - simultaneousWith(): a handle that this event happens simultaneously with. Multiple events + (including those on the same lane) may be simultaneous with the same handle, forming an + equivalence relation. From a happens-before view, these events are merged into the same + event and share predecessors and successors. Simultaneous events must not have a happens-before + relationship between each other, with one exception: concurrent ops in the same lane whose only + dependency is from implicit op order, with no intervening non-simultaneous ops. + - concurrentWith(Op other): A "must be concurrent" check allowing an underapproximation of + concurrency. Currently this is only used when checking the legality of simultaneousWith(). + }]; + let methods = [ + InterfaceMethod< + [{Returns the SSA value handles whose events must complete before this event can start}], + "::llvm::SmallVector<::mlir::Value>", "happensAfter", (ins) + >, + InterfaceMethod< + [{Returns the SSA value handles whose events must happen after this event.}], + "::llvm::SmallVector<::mlir::Value>", "happensBefore", (ins) + >, + InterfaceMethod< + [{Returns the SSA value handles whose events this event happens simultaneously with}], + "::llvm::SmallVector<::mlir::Value>", "simultaneousWith", (ins) + >, + InterfaceMethod< + [{An underapproximating relation allowing concurrency between two ops. }], + "bool", "concurrentWith", (ins "mlir::Operation *":$other) + > + ]; +} + #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_INTERFACES \ No newline at end of file diff --git a/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp new file mode 100644 index 0000000000..d6600ee47d --- /dev/null +++ b/src/enzyme_ad/jax/Dialect/Distributed/MeshComputationOps.cpp @@ -0,0 +1,21 @@ +#include "Dialect.h" + +namespace mlir::enzyme::distributed { + +LogicalResult MeshComputationOp::verify() { + Block &bodyBlock = getBody().front(); + for (Operation &op : bodyBlock.getOperations()) { + if (!op.hasTrait()) { + return emitOpError() + << "only static metadata ops are allowed in the mesh body; " + << "operation '" << op.getName() << "' is not marked with " + << "metadata trait"; + } + } + return success(); +} + +} // namespace mlir::enzyme::distributed + +#define GET_OP_CLASSES +#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedOps.cpp.inc" diff --git a/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp b/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp deleted file mode 100644 index 21d62cecf6..0000000000 --- a/src/enzyme_ad/jax/Dialect/Distributed/Ops.cpp +++ /dev/null @@ -1,137 +0,0 @@ -#include "mlir/IR/Builders.h" -#include "llvm/ADT/TypeSwitch.h" - -#include "Dialect.h" - -using mlir::OpTrait::enzyme::distributed::ChannelDefTrait; -using mlir::OpTrait::enzyme::distributed::DeviceDefTrait; -namespace mlir::enzyme::distributed { - -/** - * Returns success if the symbol reference attribute refers to an operation of - * type T in the given symbol table. - */ -template -llvm::LogicalResult checkSymbolIsA(::mlir::SymbolTableCollection &symbol_table, - mlir::Operation *owning_op, - mlir::SymbolRefAttr attr) { - auto loc = owning_op->getLoc(); - Operation *symOp = symbol_table.lookupNearestSymbolFrom(owning_op, attr); - if (dyn_cast_or_null(symOp) == nullptr) { - mlir::emitError(loc) << "invalid symbol reference or symbol type: " << attr; - return mlir::failure(); - } - return mlir::success(); -} - -/** - * Returns success if the attribute is a symbol reference to an operation with - * the given trait in the provided symbol table. - */ -template