Add Op index in IR for internal debugging.

LiteRT-PiperOrigin-RevId: 828803023

Author: YunanAZ

litert/core/model/model.cc

@@ -477,6 +477,10 @@ bool IsConstant(const LiteRtTensorT& tensor) {
       << "Constant tensors should not be defined by an op";
   return is_const;
 }
+// TODO: reuse this in cc api.
+bool IsSubgraphInput(const LiteRtTensorT& tensor) {
+  return !IsConstant(tensor) && tensor.DefiningOp() == nullptr;
+}
 
 void AttachInput(LiteRtTensor tensor, LiteRtOpT& op) {
   op.Inputs().push_back(tensor);
@@ -532,7 +536,7 @@ bool IsCompiledOp(const LiteRtModelT& graph, LiteRtOpT& op) {
   // If there hasn't been a round of serialization,
   // since dispatches were added, they won't be in the code
   // table.
-  // TODO: Fix this once the code table is updateded
+  // TODO: Fix this once the code table is updated
   // dynamically.
   return litert::internal::GetTflOpCodeInd(op) ==
              litert::internal::kDispatchOpCodeTflInd ||

litert/core/model/model.h

@@ -498,6 +498,10 @@ class LiteRtOpT {
   // Get any custom options attached to this op. Empty if there are none.
   litert::BufferRef<uint8_t> CustomOptions() const { return custom_options_; }
 
+  // Op index. (For internal debugging only)
+  void SetOpIndex(uint32_t op_index) { op_index_ = op_index; }
+  uint32_t OpIndex() const { return op_index_; }
+
   // Attach custom opaque optins to this op.
   template <class... Args>
   void SetCustomOptions(Args&&... args) {
@@ -573,6 +577,8 @@ class LiteRtOpT {
 
   std::string custom_code_;
 
+  uint32_t op_index_;
+
   // TFLITE
   int32_t tfl_op_code_ind_ = litert::internal::kDispatchOpCodeTflInd;
   ::litert::internal::TflOptions tfl_option_;
@@ -1215,6 +1221,9 @@ UseIndices FindUseInds(const LiteRtTensorT& tensor, const LiteRtOpT& op);
 // Is this tensor a constant tensor?
 bool IsConstant(const LiteRtTensorT& tensor);
 
+// Is this tensor a subgraph input tensor?
+bool IsSubgraphInput(const LiteRtTensorT& tensor);
+
 // MUTATORS
 
 // Attaches the pre-allocated tensor to be an input of given op.

litert/core/model/model_load.cc

@@ -15,6 +15,7 @@
 #include "litert/core/model/model_load.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -25,7 +26,6 @@
 #include "absl/strings/string_view.h"  // from @com_google_absl
 #include "litert/c/internal/litert_logging.h"
 #include "litert/c/litert_common.h"
-#include "litert/c/litert_model.h"
 #include "litert/c/litert_op_code.h"
 #include "litert/cc/litert_buffer_ref.h"
 #include "litert/cc/litert_expected.h"
@@ -90,7 +90,8 @@ class FlatbufferContext {
 };
 
 LiteRtStatus UnpackOp(FlatbufferContext& context, LiteRtSubgraphT& parent,
-                      const TflPackedOp& tfl_op, LiteRtOpT& litert_op) {
+                      const TflPackedOp& tfl_op, LiteRtOpT& litert_op,
+                      size_t op_index) {
   // I/O TENSORS
 
   if (tfl_op.intermediates() && tfl_op.intermediates()->size() != 0) {
@@ -145,6 +146,7 @@ LiteRtStatus UnpackOp(FlatbufferContext& context, LiteRtSubgraphT& parent,
   // OP CODE
 
   LITERT_RETURN_IF_ERROR(context.SetOpCode(litert_op, tfl_op.opcode_index()));
+  litert_op.SetOpIndex(op_index);
 
   return kLiteRtStatusOk;
 }
@@ -275,7 +277,7 @@ LiteRtStatus UnpackSubgraph(FlatbufferContext& context,
   for (auto i = 0; i < num_ops; ++i) {
     const auto* tfl_op = tfl_subgraph.operators()->Get(i);
     LITERT_RETURN_IF_ERROR(UnpackOp(context, litert_subgraph, *tfl_op,
-                                    litert_subgraph.EmplaceOp()));
+                                    litert_subgraph.EmplaceOp(), i));
   }
 
   // Update subgraph I/O.