diff --git a/ethosu/regor/compiler/scheduler_packing.cpp b/ethosu/regor/compiler/scheduler_packing.cpp
index 745e02524d1161256f19d7b751420c15d598436b..7af4acf2f10edc81f9d1567700167323124ea18c 100644
--- a/ethosu/regor/compiler/scheduler_packing.cpp
+++ b/ethosu/regor/compiler/scheduler_packing.cpp
@@ -140,7 +140,16 @@ void SchedulerPacking::FilterOperations(const std::vector<Operation *> &executio
 
         if ( ShouldDecompose(_arch, schedOp.get()) )
         {
+            auto srcKey = schedOp->_srcKey;
             auto schedOps = DecomposeSchedulerOperation(std::move(schedOp));
+            // Track source keys
+            for ( auto &newOp : schedOps )
+            {
+                if ( !newOp->_srcKey )
+                {
+                    newOp->_srcKey = srcKey;
+                }
+            }
             _schedList.insert(
                 _schedList.end(), std::make_move_iterator(schedOps.begin()), std::make_move_iterator(schedOps.end()));
         }