up

Author: metascroy

backends/mlx/runtime/MLXInterpreter.h

@@ -98,6 +98,11 @@ inline std::vector<int> infer_shape_with_minus_one(
 
 inline void exec_noop(const NoopNode&, ExecutionState&, StreamOrDevice) {}
 
+inline void
+exec_id_copy(const IdCopyNode& n, ExecutionState& st, StreamOrDevice) {
+  st.set_tensor(n.out, st.const_tensor_ref(n.x));
+}
+
 inline void
 exec_addmm(const AddmmNode& n, ExecutionState& st, StreamOrDevice s) {
   const auto& mat1 = st.const_tensor_ref(n.mat1);
@@ -154,6 +159,9 @@ class Interpreter {
       case OpCode::NOOP:
         ops::exec_noop(std::get<NoopNode>(instr.node), st, s);
         break;
+      case OpCode::ID_COPY:
+        ops::exec_id_copy(std::get<IdCopyNode>(instr.node), st, s);
+        break;
       case OpCode::ADDMM:
         ops::exec_addmm(std::get<AddmmNode>(instr.node), st, s);
         break;

backends/mlx/serialization/schema.fbs

@@ -72,6 +72,11 @@ table IntOrVidOrTid {
 
 table NoopNode {}
 
+table IdCopyNode {
+    x: Tid (required);
+    out: Tid (required);
+}
+
 table AddmmNode {
     mat1: Tid (required);      // First matrix
     mat2: Tid (required);      // Second matrix
@@ -89,6 +94,7 @@ table AddmmNode {
 // Reordering or removing members changes numeric type IDs and breaks existing .pte files.
 union OpNode {
     NoopNode,
+    IdCopyNode,
     AddmmNode
    // BC: Add new op nodes here (append only)
 }