Add stablehlo.async_start and stablehlo.async_done.

See openxla/stablehlo#2897 for context.

This CL introduces a `stablehlo.future` type and `stablehlo.async_start` and
`stablehlo.async_done`. It does not add any translation for them. That will
come in a later change. I also did not make `async_start` and `async_done`
variadic for now, even though some collectives are variadic. We can add support
for that later if we need it.

PiperOrigin-RevId: 874740969

Author: mwhittaker

xla/hlo/translate/mhlo_to_hlo/gen_hlo_op_writer.td

@@ -55,6 +55,8 @@ defvar CustomHloConverterOps = [
   StableHLO_AllGatherOp,
   StableHLO_AllReduceOp,
   StableHLO_AllToAllOp,
+  StableHLO_AsyncDoneOp,
+  StableHLO_AsyncStartOp,
   StableHLO_BatchNormGradOp,
   StableHLO_BatchNormTrainingOp,
   StableHLO_BitcastConvertOp,

xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc

@@ -3064,6 +3064,16 @@ LogicalResult ExportXlaOp(UniformDequantizeOp op, OpLoweringContext ctx) {
   return failure();
 }
 
+LogicalResult ExportXlaOp(AsyncStartOp op, OpLoweringContext ctx) {
+  // TODO(mwhittaker): Implement.
+  return failure();
+}
+
+LogicalResult ExportXlaOp(AsyncDoneOp op, OpLoweringContext ctx) {
+  // TODO(mwhittaker): Implement.
+  return failure();
+}
+
 }  // namespace
 }  // namespace stablehlo
 

xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h

@@ -178,6 +178,21 @@ MAP_HLO_TO_HLO_TYPE_REWRITE(AsyncDoneOp)
 
 #undef MAP_HLO_TO_HLO_TYPE_REWRITE
 
+// TODO(mwhittaker): Remove this when we translate async ops between StableHLO
+// and MHLO.
+#define MAP_STABLEHLO_TO_STABLEHLO_TYPE_REWRITE(OpName) \
+  template <>                                           \
+  struct HloToStablehloOpImpl<stablehlo::OpName> {      \
+    using Type = stablehlo::OpName;                     \
+  };                                                    \
+  template <>                                           \
+  struct StablehloToHloOpImpl<stablehlo::OpName> {      \
+    using Type = stablehlo::OpName;                     \
+  };
+MAP_STABLEHLO_TO_STABLEHLO_TYPE_REWRITE(AsyncStartOp)
+MAP_STABLEHLO_TO_STABLEHLO_TYPE_REWRITE(AsyncDoneOp)
+#undef MAP_STABLEHLO_TO_STABLEHLO_TYPE_REWRITE
+
 }  // namespace stablehlo
 }  // namespace mlir