Distributed dialect- change subbranches to regions

Author: EganJ

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

@@ -111,20 +111,44 @@ GroupSplitOp::verifySymbolUses(::mlir::SymbolTableCollection &symbol_table) {
                                        getDeviceGroupAttr());
 }
 
-LogicalResult
-SplitBranchOp::verifySymbolUses(::mlir::SymbolTableCollection &symbol_table) {
-  // Split branches have programs for individual devices or channels
-  Operation *dev_or_chan =
-      symbol_table.lookupNearestSymbolFrom(*this, getDeviceOrChannelAttr());
-  if (!dev_or_chan || !(dev_or_chan->hasTrait<DeviceDefTrait>() ||
-                        dev_or_chan->hasTrait<ChannelDefTrait>())) {
-    mlir::emitError(getLoc())
-        << "branches must reference a valid device or channel";
-    return mlir::failure();
+// Printer/parser for GroupsplitOp branches
+mlir::ParseResult parseSplitBranches(
+    OpAsmParser &parser, mlir::ArrayAttr &branchAssignments,
+    llvm::SmallVector<std::unique_ptr<::mlir::Region>, 2> &branchesRegions) {
+  // Expect 0 or more `branch` $symbol_name $symbol_region
+  // While next token is `branch`:
+  llvm::SmallVector<mlir::Attribute, 2> assignment_symbols;
+  while (parser.parseOptionalKeyword("branch").succeeded()) {
+    // Parse symbol name
+    mlir::SymbolRefAttr sym;
+    auto sym_parse_failed = parser.parseAttribute<mlir::SymbolRefAttr>(sym);
+    if (sym_parse_failed)
+      return mlir::failure();
+    assignment_symbols.push_back(sym);
+
+    // Put placeholder region in list and parse into it
+    branchesRegions.push_back(std::make_unique<mlir::Region>());
+    auto parse_region_failed = parser.parseRegion(*branchesRegions.back());
+    if (parse_region_failed)
+      return mlir::failure();
   }
+
+  branchAssignments = mlir::ArrayAttr::get(parser.getBuilder().getContext(),
+                                           assignment_symbols);
   return mlir::success();
 }
 
+void printSplitBranches(OpAsmPrinter &printer, const GroupSplitOp &op,
+                        const mlir::ArrayAttr branchAssignments,
+                        const llvm::MutableArrayRef<mlir::Region> branches) {
+  // Print each branch as `branch` $symbol_name $symbol_region
+  for (size_t i = 0; i < branches.size(); i++) {
+    printer << " branch ";
+    printer.printAttribute(branchAssignments[i]);
+    printer.printRegion(branches[i]);
+  }
+}
+
 LogicalResult
 DefineTokenOp::verifySymbolUses(::mlir::SymbolTableCollection &symbol_table) {
   // Tokens need to indicate which channel they communicate over

src/enzyme_ad/jax/Dialect/Distributed/Ops.td

@@ -59,25 +59,16 @@ def MeshForOp : DistributedOp<"MeshFor", [DeclareOpInterfaceMethods<SymbolUserOp
     let assemblyFormat = "$mesh $body attr-dict";
 }
 
-def GroupSplitOp : DistributedOp<"GroupSplit", [DeclareOpInterfaceMethods<SymbolUserOpInterface>, NoTerminator, SingleBlock]>{
+def GroupSplitOp : DistributedOp<"GroupSplit", [DeclareOpInterfaceMethods<SymbolUserOpInterface>, NoTerminator]>{
     let arguments = (ins
-        SymbolRefAttr:$device_group // TODO: verify it's a group
+        SymbolRefAttr:$device_group,
+        ArrayAttr:$branch_assignments // Symbols mapping to each branch region
     );
-    let regions = (region SizedRegion<1>:$declarations); // Takes as args the devices and channels in the group
+    // TODO check that declarations only declares tokens.
+    let regions = (region VariadicRegion<SizedRegion<1>>:$branches);
     let results = (outs ); // TODO
     // let hasVerifier = 1; // TODO
-    // let hasCanonicalizer = 1; // TODO: token declarations up front, followed by device and channel branches in order of listing in the group
-    let assemblyFormat = "$device_group $declarations  attr-dict";
-}
-
-def SplitBranchOp : DistributedOp<"SplitBranch", [DeclareOpInterfaceMethods<SymbolUserOpInterface>, NoTerminator, SingleBlock]>{
-    let arguments = (ins
-        SymbolRefAttr:$device_or_channel // TODO: verify it's a device or channel
-        );
-    let regions = (region MaxSizedRegion<1>:$body); // Takes as args the device or channel
-    let results = (outs ); // TODO
-    // let hasVerifier = 1; // TODO: parent is a groupsplitop
-    let assemblyFormat = "$device_or_channel $body attr-dict";
+    let assemblyFormat = "$device_group custom<SplitBranches>($branch_assignments, $branches) attr-dict";
 }
 
 def DefineTokenOp : DistributedOp<"DefineToken", [DeclareOpInterfaceMethods<SymbolUserOpInterface>]>{

test/lit_tests/distributed/roundtrip.mlir

@@ -6,37 +6,31 @@ distributed.Channel @chan1 [@myCpu, @gpuMesh] [@gpuMesh, @myCpu]
 distributed.DeviceGroup @gpusWithHost [@myGpu, @myCpu] [@chan1]
 
 func.func @foo() {
-    distributed.GroupSplit @gpusWithHost {
-        %tok = distributed.DefineToken @chan1
-        distributed.SplitBranch @chan1 { }
-        distributed.SplitBranch @myCpu {}
-        distributed.SplitBranch @gpuMesh {
-            distributed.MeshFor @gpuMesh {
+    distributed.GroupSplit @gpusWithHost 
+       branch @myGpu {
+               distributed.MeshFor @gpuMesh {
+               }
+       }
+       branch @myCpu {
+          distributed.DefineToken @chan1
+       }
 
-            }
-        }
-    }
     func.return
 }
 
-// CHECK: module {
-// CHECK-NEXT:  distributed.LeafDevice @myGpu
-// CHECK-NEXT:  distributed.DeviceMesh @gpuMesh @myGpu [2, 2]
-// CHECK-NEXT:  distributed.LeafDevice @myCpu
-// CHECK-NEXT:  distributed.Channel @chan1 [@myCpu, @gpuMesh] [@gpuMesh, @myCpu]
-// CHECK-NEXT:  distributed.DeviceGroup @gpusWithHost [@myGpu, @myCpu] [@chan1]
-// CHECK-NEXT:  func.func @foo() {
-// CHECK-NEXT:    distributed.GroupSplit @gpusWithHost {
-// CHECK-NEXT:      %0 = distributed.DefineToken @chan1
-// CHECK-NEXT:      distributed.SplitBranch @chan1 {
-// CHECK-NEXT:      }
-// CHECK-NEXT:      distributed.SplitBranch @myCpu {
-// CHECK-NEXT:      }
-// CHECK-NEXT:      distributed.SplitBranch @gpuMesh {
-// CHECK-NEXT:        distributed.MeshFor @gpuMesh {
-// CHECK-NEXT:        }
-// CHECK-NEXT:      }
-// CHECK-NEXT:    }
-// CHECK-NEXT:    return
-// CHECK-NEXT:  }
-// CHECK-NEXT: }
+//CHECK: module {
+//CHECK-NEXT:   distributed.LeafDevice @myGpu
+//CHECK-NEXT:   distributed.DeviceMesh @gpuMesh @myGpu [2, 2]
+//CHECK-NEXT:   distributed.LeafDevice @myCpu
+//CHECK-NEXT:   distributed.Channel @chan1 [@myCpu, @gpuMesh] [@gpuMesh, @myCpu]
+//CHECK-NEXT:   distributed.DeviceGroup @gpusWithHost [@myGpu, @myCpu] [@chan1]
+//CHECK-NEXT:   func.func @foo() {
+//CHECK-NEXT:     distributed.GroupSplit @gpusWithHost  branch @myGpu{
+//CHECK-NEXT:       distributed.MeshFor @gpuMesh {
+//CHECK-NEXT:       }
+//CHECK-NEXT:     } branch @myCpu{
+//CHECK-NEXT:       %0 = distributed.DefineToken @chan1
+//CHECK-NEXT:     }
+//CHECK-NEXT:     return
+//CHECK-NEXT:   }
+//CHECK-NEXT: }