From 1f5ff96c1088fd43b0a46de33c9a2b47d66ba166 Mon Sep 17 00:00:00 2001
From: Tarry Singh <tarry.singh@deepkapha.com>
Date: Mon, 9 Sep 2024 15:00:28 +0000
Subject: [PATCH] zml/ops: add wiring for operand output alias in
 zml.ops.triton

---
 zml/ops.zig | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/zml/ops.zig b/zml/ops.zig
index ad85403..6ab3bad 100644
--- a/zml/ops.zig
+++ b/zml/ops.zig
@@ -819,6 +819,7 @@ pub const TritonOps = struct {
     grid: [3]i32,
     num_stages: i32,
     num_warps: i32,
+    output_operand_aliases: []const i64 = &.{},
 };
 
 /// Generate an MLIR call to the given member function with the given tensors.
@@ -877,6 +878,7 @@ pub fn triton(inputs: anytype, outputs: anytype, opts: TritonOps) [outputs.len]T
             .api_version = .typed_ffi,
             .operand_layouts = &operands_layouts,
             .result_layouts = &results_layouts,
+            .output_operand_aliases = opts.output_operand_aliases,
         },
         &res_types,
         ctx.mlirCtx().location(@src()),