Send, Recv ops and interfaces

Author: EganJ

src/enzyme_ad/jax/BUILD

@@ -569,11 +569,11 @@ gentbl_cc_library(
     name = "DistributedInterfacesIncGen",
     tbl_outs = [
         (
-            ["--gen-interface-decls"],
+            ["--gen-op-interface-decls"],
             "Dialect/Distributed/DistributedInterfaces.h.inc",
         ),
         (
-            ["--gen-interface-defs"],
+            ["--gen-op-interface-defs"],
             "Dialect/Distributed/DistributedInterfaces.cpp.inc",
         ),
     ],
@@ -744,6 +744,7 @@ cc_library(
     deps = [
         ":CheckedRewrite",
         ":DistributedDialectIncGen",
+        ":DistributedInterfacesIncGen",
         ":DistributedOpsIncGen",
         ":DistributedTypesIncGen",
         ":EnzymeHLOPatternsIncGen",

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<

src/enzyme_ad/jax/Dialect/Distributed/Dialect.h

@@ -10,14 +10,14 @@
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/Types.h"
 
-// Include the dialect
-#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedDialect.h.inc"
-// Traits and interfaces
 #include "Traits.h"
-// Types
+#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedDialect.h.inc"
+
 #define GET_TYPEDEF_CLASSES
 #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedTypes.h.inc"
-// Operations
+
+#include "src/enzyme_ad/jax/Dialect/Distributed/DistributedInterfaces.h.inc"
+
 #define GET_OP_CLASSES
 #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedOps.h.inc"
 

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

@@ -6,4 +6,29 @@ include "mlir/IR/OpBase.td"
 def DeviceDefTrait : NativeOpTrait<"enzyme::distributed::DeviceDefTrait">;
 def ChannelDefTrait : NativeOpTrait<"enzyme::distributed::ChannelDefTrait">;
 
+def TokenReaderOpInterface : OpInterface<"TokenReaderOpInterface"> {
+    let cppNamespace = "::mlir::enzyme::distributed";
+    let description = [{
+        An interface to determine which ops can read from a channel and what type they expect.
+        Ops may read from multiple channels.
+    }];
+    let methods = [
+        InterfaceMethod<"Returns the SSA values of tokens read from this op.", "::llvm::ArrayRef<::mlir::TypedValue<::mlir::enzyme::distributed::ReadTokenType>>", "getReadTokens">,
+        InterfaceMethod<"Returns the types of tokens read from this op, parallel to getReadTokens.", "::llvm::ArrayRef<::mlir::Type>", "getReadTokenTypes">
+    ];
+}
+
+def TokenWriterOpInterface : OpInterface<"TokenWriterOpInterface"> {
+    let cppNamespace = "::mlir::enzyme::distributed";
+    let description = [{
+        An interface to determine which ops can write to a channel and what type they provide.
+        Ops may write to multiple channels.
+    }];
+    let methods = [
+        InterfaceMethod<"Returns the SSA values of tokens written from this op.", "::llvm::ArrayRef<::mlir::TypedValue<::mlir::enzyme::distributed::WriteTokenType>>", "getWriteTokens">,
+        InterfaceMethod<"Returns the types of tokens written from this op, parallel to getWriteTokens.", "::llvm::ArrayRef<::mlir::Type>", "getWriteTokenTypes">
+    ];
+}
+
+
 #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_INTERFACES

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

@@ -156,6 +156,20 @@ DefineTokenOp::verifySymbolUses(::mlir::SymbolTableCollection &symbol_table) {
                                               getChannelAttr());
 }
 
+llvm::ArrayRef<mlir::TypedValue<WriteTokenType>> SendOp::getWriteTokens() {
+  return llvm::SmallVector<mlir::TypedValue<WriteTokenType>, 1>{getToken()};
+}
+llvm::ArrayRef<mlir::Type> SendOp::getWriteTokenTypes() {
+  return llvm::SmallVector<mlir::Type, 1>{getValue().getType()};
+}
+
+llvm::ArrayRef<mlir::TypedValue<ReadTokenType>> RecvOp::getReadTokens() {
+  return llvm::SmallVector<mlir::TypedValue<ReadTokenType>, 1>{getToken()};
+}
+llvm::ArrayRef<mlir::Type> RecvOp::getReadTokenTypes() {
+  return llvm::SmallVector<mlir::Type, 1>{getValue().getType()};
+}
+
 } // namespace mlir::enzyme::distributed
 #define GET_OP_CLASSES
 #include "src/enzyme_ad/jax/Dialect/Distributed/DistributedOps.cpp.inc"

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

@@ -75,9 +75,24 @@ def DefineTokenOp : DistributedOp<"DefineToken", [DeclareOpInterfaceMethods<Symb
     let arguments = (ins
             SymbolRefAttr:$channel
         );
-    let results = (outs TokenType:$token_out);
+    let results = (outs ReadTokenType:$read_token, WriteTokenType:$write_token);
     // let hasVerifier = 1; // TODO: verify writers and readers are connected to the channel
     let assemblyFormat = "$channel attr-dict";
 }
 
+def SendOp : DistributedOp<"Send", [DeclareOpInterfaceMethods<TokenWriterOpInterface>]>{
+    let arguments = (ins
+        WriteTokenType:$token,
+        // value to send
+        AnyType:$value);
+    let assemblyFormat = "$token type($value) $value attr-dict";
+}
+
+def RecvOp : DistributedOp<"Recv", [DeclareOpInterfaceMethods<TokenReaderOpInterface>]>{
+    let arguments = (ins
+        ReadTokenType:$token);
+    let results = (outs AnyType:$value);
+    let assemblyFormat = "$token type($value) attr-dict";
+}
+
 #endif // ENZYME_AD_JAX_DIALECT_DISTRIBUTED_OPS_TD

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

@@ -3,6 +3,7 @@
 
 include "Dialect.td"
 
-def TokenType : DistributedType<"Token", "token">;
+def ReadTokenType : DistributedType<"ReadToken", "read_token">;
+def WriteTokenType : DistributedType<"WriteToken", "write_token">;
 
 #endif // ENZYME_DISTRIBUTED_DIALECT_TYPES_H