Working commit: single token type

Author: EganJ

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

@@ -13,7 +13,7 @@ def TokenReaderOpInterface : OpInterface<"TokenReaderOpInterface"> {
         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 SSA values of tokens read from this op.", "::llvm::ArrayRef<::mlir::TypedValue<::mlir::enzyme::distributed::TokenType>>", "getReadTokens">,
         InterfaceMethod<"Returns the types of tokens read from this op, parallel to getReadTokens.", "::llvm::ArrayRef<::mlir::Type>", "getReadTokenTypes">
     ];
 }
@@ -25,7 +25,7 @@ def TokenWriterOpInterface : OpInterface<"TokenWriterOpInterface"> {
         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 SSA values of tokens written from this op.", "::llvm::ArrayRef<::mlir::TypedValue<::mlir::enzyme::distributed::TokenType>>", "getWriteTokens">,
         InterfaceMethod<"Returns the types of tokens written from this op, parallel to getWriteTokens.", "::llvm::ArrayRef<::mlir::Type>", "getWriteTokenTypes">
     ];
 }

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

@@ -156,15 +156,15 @@ 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::TypedValue<TokenType>> SendOp::getWriteTokens() {
+  return llvm::SmallVector<mlir::TypedValue<TokenType>, 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::TypedValue<TokenType>> RecvOp::getReadTokens() {
+  return llvm::SmallVector<mlir::TypedValue<TokenType>, 1>{getToken()};
 }
 llvm::ArrayRef<mlir::Type> RecvOp::getReadTokenTypes() {
   return llvm::SmallVector<mlir::Type, 1>{getValue().getType()};

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

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

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

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