ProfileData::merge bytecode

Summary:
Implement merge of bytecode profile results.

`buck2 profile analysis --mode=bytecode --recursive` works now.

Bytecode profile is not very helpful for typical buck users, but it was easier to implement, and we need to start somewhere.

Reviewed By: bobyangyf

Differential Revision: D38681490

fbshipit-source-id: 832234c172121b31470cb6d84d0cad4407a934fb

Author: stepancheg

starlark/src/eval/runtime/profile/bc.rs

@@ -20,6 +20,7 @@
 use std::collections::HashMap;
 use std::iter::Sum;
 use std::mem;
+use std::ops::AddAssign;
 
 use gazebo::prelude::*;
 
@@ -45,24 +46,55 @@ impl<'a> Sum<&'a BcInstrStat> for BcInstrStat {
     }
 }
 
+impl<'a> AddAssign<&'a BcInstrStat> for BcInstrStat {
+    fn add_assign(&mut self, other: &'a BcInstrStat) {
+        let BcInstrStat { count } = other;
+        self.count += count;
+    }
+}
+
 #[derive(Default, Clone, Copy, Dupe, Debug)]
 struct BcInstrPairsStat {
     count: u64,
     // We are not measuring time here, because even time for single opcode
     // is not very accurate or helpful, and time for pairs is even less helpful.
 }
 
+impl<'a> AddAssign<&'a BcInstrPairsStat> for BcInstrPairsStat {
+    fn add_assign(&mut self, other: &'a BcInstrPairsStat) {
+        let BcInstrPairsStat { count } = other;
+        self.count += count;
+    }
+}
+
 #[derive(Clone, Debug)]
 pub(crate) struct BcProfileData {
     by_instr: [BcInstrStat; BcOpcode::COUNT],
 }
 
+impl<'a> AddAssign<&'a BcProfileData> for BcProfileData {
+    fn add_assign(&mut self, rhs: &'a BcProfileData) {
+        for (lhs, rhs) in self.by_instr.iter_mut().zip(rhs.by_instr.iter()) {
+            *lhs += rhs;
+        }
+    }
+}
+
 #[derive(Default, Clone, Debug)]
 pub(crate) struct BcPairsProfileData {
     last: Option<BcOpcode>,
     by_instr: HashMap<[BcOpcode; 2], BcInstrPairsStat>,
 }
 
+impl<'a> AddAssign<&'a BcPairsProfileData> for BcPairsProfileData {
+    fn add_assign(&mut self, rhs: &'a BcPairsProfileData) {
+        self.last = None;
+        for (pair, stat) in &rhs.by_instr {
+            *self.by_instr.entry(*pair).or_default() += stat;
+        }
+    }
+}
+
 // Derive doesn't work here.
 impl Default for BcProfileData {
     fn default() -> Self {
@@ -104,6 +136,14 @@ impl BcProfileData {
         }
         csv.finish()
     }
+
+    pub(crate) fn merge<'a>(iter: impl IntoIterator<Item = &'a BcProfileData>) -> BcProfileData {
+        let mut sum = BcProfileData::default();
+        for profile in iter {
+            sum += profile;
+        }
+        sum
+    }
 }
 
 impl BcPairsProfileData {
@@ -138,6 +178,16 @@ impl BcPairsProfileData {
         }
         csv.finish()
     }
+
+    pub(crate) fn merge<'a>(
+        iter: impl IntoIterator<Item = &'a BcPairsProfileData>,
+    ) -> BcPairsProfileData {
+        let mut sum = BcPairsProfileData::default();
+        for profile in iter {
+            sum += profile;
+        }
+        sum
+    }
 }
 
 enum BcProfileDataMode {
@@ -210,6 +260,8 @@ mod tests {
     use crate::environment::Globals;
     use crate::environment::Module;
     use crate::eval::bc::opcode::BcOpcode;
+    use crate::eval::runtime::profile::bc::BcPairsProfileData;
+    use crate::eval::runtime::profile::bc::BcProfileData;
     use crate::eval::Evaluator;
     use crate::eval::ProfileMode;
     use crate::syntax::AstModule;
@@ -261,4 +313,18 @@ mod tests {
             csv
         );
     }
+
+    #[test]
+    fn test_bc_profile_data_merge() {
+        let bc = BcProfileData::default();
+        // Smoke test.
+        BcProfileData::merge([&bc, &bc, &bc]);
+    }
+
+    #[test]
+    fn test_bc_pairs_profile_data_merge() {
+        let bc = BcPairsProfileData::default();
+        // Smoke test.
+        BcPairsProfileData::merge([&bc, &bc, &bc]);
+    }
 }

starlark/src/eval/runtime/profile/data.rs

@@ -108,6 +108,59 @@ impl ProfileData {
                 return Err(ProfileDataError::DifferentProfileModes.into());
             }
         }
-        Err(ProfileDataError::MergeNotImplemented(profile_mode.dupe()).into())
+        let profile = match &profile_mode {
+            ProfileMode::Bytecode => {
+                let profiles = profiles.try_map(|p| match &p.profile {
+                    ProfileDataImpl::Bc(bc) => Ok(&**bc),
+                    _ => Err(ProfileDataError::ProfileDataNotConsistent),
+                })?;
+                let profile = BcProfileData::merge(profiles);
+                ProfileDataImpl::Bc(box profile)
+            }
+            ProfileMode::BytecodePairs => {
+                let profiles = profiles.try_map(|p| match &p.profile {
+                    ProfileDataImpl::BcPairs(bc_pairs) => Ok(bc_pairs),
+                    _ => Err(ProfileDataError::ProfileDataNotConsistent),
+                })?;
+                let profile = BcPairsProfileData::merge(profiles);
+                ProfileDataImpl::BcPairs(profile)
+            }
+            profile_mode => {
+                return Err(ProfileDataError::MergeNotImplemented(profile_mode.dupe()).into());
+            }
+        };
+        Ok(ProfileData {
+            profile_mode,
+            profile,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::eval::runtime::profile::bc::BcPairsProfileData;
+    use crate::eval::runtime::profile::bc::BcProfileData;
+    use crate::eval::runtime::profile::data::ProfileDataImpl;
+    use crate::eval::ProfileData;
+    use crate::eval::ProfileMode;
+
+    #[test]
+    fn merge_bc() {
+        let profile = ProfileData {
+            profile_mode: ProfileMode::Bytecode,
+            profile: ProfileDataImpl::Bc(box BcProfileData::default()),
+        };
+        // Smoke.
+        ProfileData::merge([&profile, &profile]).unwrap();
+    }
+
+    #[test]
+    fn merge_bc_pairs() {
+        let profile = ProfileData {
+            profile_mode: ProfileMode::BytecodePairs,
+            profile: ProfileDataImpl::BcPairs(BcPairsProfileData::default()),
+        };
+        // Smoke.
+        ProfileData::merge([&profile, &profile]).unwrap();
     }
 }