Eganj distributed pull part2

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",

src/enzyme_ad/jax/Dialect/Axis/Dialect.td

@@ -14,9 +14,8 @@ def AxisDialect : Dialect {
   let useDefaultTypePrinterParser = 1;
 }
 
-class AxisDialectType<string name, string type_mnemonic,
-                      list<Trait> traits = []>
-    : TypeDef<AxisDialect, name, traits> {
+class AxisDialectType<string name, string type_mnemonic, list<Trait> traits = [
+]> : TypeDef<AxisDialect, name, traits> {
   let mnemonic = type_mnemonic;
 }
 

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)
     >

src/enzyme_ad/jax/Dialect/Axis/Ops.cpp

@@ -22,6 +22,24 @@ static LogicalResult verifyFactorExtents(ArrayRef<int32_t> extents,
   return success();
 }
 
+static LogicalResult verifySegmentExtents(ArrayRef<int32_t> 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<uint64_t>(extent);
+  }
+  if (sum != sourceExtent) {
+    return op->emitOpError() << "requires sum(segment_extents) == axis extent "
+                                "("
+                             << sum << " != " << sourceExtent << ")";
+  }
+  return success();
+}
+
 static void computeMajorToMinorStrides(ArrayRef<int32_t> extents,
                                        SmallVectorImpl<unsigned> &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<AxisTypeInterface>(getAxis().getType());
+  if (!axisIface) {
+    return emitOpError() << "requires axis operand type to implement "
+                            "AxisTypeInterface";
+  }
+
+  ArrayRef<int32_t> 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<AxisSegmentType>(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<unsigned>(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<unsigned>(segmentExtents[idx]);
+  }
+
+  return success();
+}
+
+LogicalResult AxisSegmentOp::inferReturnTypes(
+    MLIRContext *context, std::optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  AxisSegmentOpAdaptor adaptor(operands, attributes, properties, regions);
+  ArrayRef<int32_t> 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<unsigned>(segmentExtent),
+        runningOffset));
+    runningOffset += static_cast<unsigned>(segmentExtent);
+  }
+  return success();
+}
+
+LogicalResult AxisProductOp::verify() {
   uint64_t extentProduct = 1;
   for (Value factor : getFactors()) {
     if (!isa<OpResult>(factor)) {
@@ -169,18 +255,18 @@ LogicalResult AxisGroupOp::verify() {
     extentProduct *= static_cast<uint64_t>(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> location, ValueRange operands,
     DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
     SmallVectorImpl<Type> &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<AxisFactorType>(factor.getType());

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<InferTypeOpInterface>]> {
+def AxisGetAxisOp : AxisOp<"getaxis", [MetadataOpTrait, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let description = [{
     Materializes a canonical ShapeAxis value for one ranked, static shape index.
   }];
@@ -22,25 +22,40 @@ def AxisGetAxisOp : AxisOp<"getaxis", [Pure, DeclareOpInterfaceMethods<InferType
   let hasVerifier = 1;
 }
 
-def AxisFactorOp : AxisOp<"factor", [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+def AxisFactorOp : AxisOp<"factor", [MetadataOpTrait, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   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<AxisFactorType> : $axis_factors);
+let assemblyFormat = "$axis $factor_extents attr-dict `:` type($axis)";
+let hasVerifier = 1;
+}
+
+def AxisSegmentOp : AxisOp<"segment", [
+  MetadataOpTrait, DeclareOpInterfaceMethods<InferTypeOpInterface>
+]> {
+  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<AxisFactorType>:$axis_factors);
-  let assemblyFormat = "$axis $factor_extents attr-dict `:` type($axis)";
+  let results = (outs Variadic<AxisSegmentType>:$axis_segments);
+  let assemblyFormat = "$axis $segment_extents attr-dict `:` type($axis)";
   let hasVerifier = 1;
 }
 
-def AxisGroupOp : AxisOp<"group", [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+def AxisProductOp : AxisOp<"product", [MetadataOpTrait, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   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<AxisFactorType>:$factors);
-  let results = (outs FactorGroupType:$group);
+  let results = (outs FactorGroupType:$product);
   let assemblyFormat = "$factors attr-dict `:` type($factors)";
   let hasVerifier = 1;
 }

src/enzyme_ad/jax/Dialect/Axis/Types.td

@@ -6,28 +6,36 @@ include "Interfaces.td"
 
 class AxisType<string name, string type_mnemonic, list<Trait> traits = []>
     : AxisDialectType<name, type_mnemonic,
-                      [DeclareTypeInterfaceMethods<AxisTypeInterface>] #
-                          traits>;
+                      [DeclareTypeInterfaceMethods<AxisTypeInterface>] #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 `>`";
 }