Vizier Algorithms integrated

arch_gym/envs/CFUPlaygroundEnv.py

@@ -31,18 +31,18 @@ def __init__(self, target_vals, max_steps, workload, reward_type = 'both', log_t
         self.observation_space = spaces.Box(low=0, high=1e12, shape=self.observation_shape)
 
         # define the action space
-        self.action_space = spaces.Tuple((
-            spaces.Discrete(2),     
-            spaces.Discrete(2),     
-            spaces.Discrete(11),
-            spaces.Discrete(2),    
-            spaces.Discrete(11),
-            spaces.Discrete(2),
-            spaces.Discrete(4),
-            spaces.Discrete(2),
-            spaces.Discrete(2),     
-            spaces.Discrete(2)
-        ))
+        self.action_space = spaces.Dict({
+            "Bypass": spaces.Discrete(2),
+            "CFU_enable": spaces.Discrete(2),
+            "Data_cache_size": spaces.Discrete(11),
+            "Hardware_Divider": spaces.Discrete(2),
+            "Instruction_cache_size": spaces.Discrete(11),
+            "Hardware_Multiplier": spaces.Discrete(2),
+            "Branch_predictor_type": spaces.Discrete(4),
+            "Safe_mode_enable": spaces.Discrete(2),
+            "Single_Cycle_Shifter": spaces.Discrete(2),     
+            "Single_Cycle_Multiplier": spaces.Discrete(2)
+        })
 
     def reset(self):
         self.no_steps=0
@@ -100,20 +100,17 @@ def calculate_reward(self):
     def runCFUPlaygroundEnv(self, action):
 
         # update action string to pass to subprocess        
-        self.action = str(action[0])         # Bypass
-        #self.action += ',' + str(action[1])        # CFU_enable 
-        self.action += ',0'                         # CFU_enable (currently set to false)
-        self.action += ',' + ('0' if action[2] == 0 else str(1<<(4+action[2])))
-                                                    # Data cache size
-        self.action += ',' + str(action[3])         # Hardware Divider
-        self.action += ',' + ('0' if action[4] == 0 else str(1<<(4+action[4])))
-                                                    # Instruction cache size
-        self.action += ',' + str(action[5])         # Hardware Multiplier
-        self.action += ',' + self.Branch_predict_types[action[6]]
-                                                    # Branch predictor
-        self.action += ',' + str(action[7])         # Safe mode
-        self.action += ',' + str(action[8])         # Single Cycle Shifter
-        self.action += ',' + str(action[9])         # Single Cycle Multiplier
+        self.action = str(action["Bypass"])
+        #self.action += ',' + str(action["CFU_enable"])
+        self.action += ',0'      # CFU_enable (currently set to false)
+        self.action += ',' + ('0' if action["Data_cache_size"] == 0 else str(1<<(4+action["Data_cache_size"])))
+        self.action += ',' + str(action["Hardware_Divider"])
+        self.action += ',' + ('0' if action["Instruction_cache_size"] == 0 else str(1<<(4+action["Instruction_cache_size"])))
+        self.action += ',' + str(action["Hardware_Multiplier"])
+        self.action += ',' + self.Branch_predict_types[action["Branch_predictor_type"]]
+        self.action += ',' + str(action["Safe_mode_enable"])
+        self.action += ',' + str(action["Single_Cycle_Shifter"])
+        self.action += ',' + str(action["Single_Cycle_Multiplier"])
         self.action += ',' + self.target
         self.action += ',' + self.workload
 

install_sim.sh

@@ -6,7 +6,26 @@ if [ $1  == 'cfu' ]; then
 
     cd sims/CFU-Playground/CFU-Playground
 
+    ./scripts/setup_vexriscv_build.sh
     ./scripts/setup
     make install-sf
-    ./scripts/setup_vexriscv_build.sh
+fi
+
+#install vizier in arch-gym conda environment
+#Assumes user is in the arch-gym conda environment
+
+if [ $1 == 'viz' ]; then
+
+    git clone https://github.com/ShvetankPrakash/vizier.git
+    cd vizier
+
+    sudo apt-get install -y libprotobuf-dev
+
+    pip install -r requirements.txt --use-deprecated=legacy-resolver
+    pip install -e .
+
+    ./build_protos.sh
+
+    pip install -r requirements-algorithms.txt
+    pip install -r requirements-benchmarks.txt
 fi

sims/CFU-Playground/train_EMUKIT_GP_CFUPlayground.py

@@ -0,0 +1,176 @@
+from concurrent import futures
+import grpc
+import portpicker
+import sys
+import os
+
+sys.path.append('../../arch_gym/envs')
+import CFUPlayground_wrapper
+from absl import flags, app, logging
+
+
+import envlogger
+from envlogger.testing import catch_env   
+
+import numpy as np
+import pandas as pd
+from vizier._src.algorithms.designers import emukit
+from vizier._src.algorithms.designers.emukit import EmukitDesigner
+from vizier.service import clients
+from vizier.service import pyvizier as vz
+from vizier.service import vizier_server
+from vizier.service import vizier_service_pb2_grpc
+
+
+flags.DEFINE_string('workload', 'micro_speech', 'workload the processor is being optimized for')
+flags.DEFINE_integer('num_steps', 1, 'Number of training steps.')
+# flags.DEFINE_integer('num_episodes', 1, 'Number of training episodes.')
+flags.DEFINE_string('traject_dir', 'EMUKIT_trajectories','Directory to save the dataset.')
+flags.DEFINE_bool('use_envlogger',True, 'Use envlogger to log the data.')  
+flags.DEFINE_string('summary_dir', '.', 'Directory to save the summary.')
+flags.DEFINE_string('reward_formulation', 'both', 'Which reward formulation to use?')
+flags.DEFINE_integer('num_random_sample', 100, 'hyperparameter for emukit')
+
+FLAGS = flags.FLAGS
+
+
+envdm = catch_env.Catch()
+
+
+def log_fitness_to_csv(filename, fitness_dict):
+    """Logs fitness history to csv file
+
+    Args:
+        filename (str): path to the csv file
+        fitness_dict (dict): dictionary containing the fitness history
+    """
+    df = pd.DataFrame([fitness_dict['reward']])
+    csvfile = os.path.join(filename, "fitness.csv")
+    df.to_csv(csvfile, index=False, header=False, mode='a')
+
+    # append to csv
+    df = pd.DataFrame([fitness_dict])
+    csvfile = os.path.join(filename, "trajectory.csv")
+    df.to_csv(csvfile, index=False, header=False, mode='a')
+
+
+def wrap_in_envlogger(env, envlogger_dir):
+    """Wraps the environment in envlogger
+
+    Args:
+        env (gym.Env): gym environment
+        envlogger_dir (str): path to the directory where the data will be logged
+    """
+    metadata = {
+        'agent_type': 'EMUKIT_GP_EI',
+        'num_steps': FLAGS.num_steps,
+        'env_type': type(env).__name__,
+    }
+    if FLAGS.use_envlogger:
+        logging.info('Wrapping environment with EnvironmentLogger...')
+        env = envlogger.EnvLogger(env,
+                                  data_directory=envlogger_dir,
+                                  max_episodes_per_file=1000,
+                                  metadata=metadata)
+        logging.info('Done wrapping environment with EnvironmentLogger.')
+        return env
+    else:
+        return env
+
+
+def main(_):
+
+    env = CFUPlayground_wrapper.make_cfuplaygroundEnv(target_vals = [1000, 1000],rl_form='EMUKIT', reward_type = FLAGS.reward_formulation, max_steps = FLAGS.num_steps, workload = FLAGS.workload)
+
+    fitness_hist = {}
+    problem = vz.ProblemStatement()
+    problem.search_space.select_root().add_int_param(name='Bypass', min_value = 0, max_value = 1)
+    problem.search_space.select_root().add_int_param(name='CFU_enable', min_value = 0, max_value = 1)
+    problem.search_space.select_root().add_int_param(name='Data_cache_size', min_value = 0, max_value = 10)
+    problem.search_space.select_root().add_int_param(name='Hardware_Divider', min_value = 0, max_value = 1)
+    problem.search_space.select_root().add_int_param(name='Instruction_cache_size', min_value = 0, max_value = 10)
+    problem.search_space.select_root().add_int_param(name='Hardware_Multiplier', min_value = 0, max_value = 1)
+    problem.search_space.select_root().add_int_param(name='Branch_predictor_type', min_value = 0, max_value = 3)
+    problem.search_space.select_root().add_int_param(name='Safe_mode_enable', min_value = 0, max_value = 1)
+    problem.search_space.select_root().add_int_param(name='Single_Cycle_Shifter', min_value = 0, max_value = 1)
+    problem.search_space.select_root().add_int_param(name='Single_Cycle_Multiplier', min_value = 0, max_value = 1)
+
+    problem.metric_information.append(
+        vz.MetricInformation(
+            name='Reward', goal=vz.ObjectiveMetricGoal.MAXIMIZE))
+
+
+    study_config = vz.StudyConfig.from_problem(problem)
+    # study_config.algorithm = vz.Algorithm.EMUKIT_GP_EI
+    emukit_designer = emukit.EmukitDesigner(problem, num_random_samples= FLAGS.num_random_sample)
+
+
+
+
+    port = portpicker.pick_unused_port()
+    address = f'localhost:{port}'
+
+    # Setup server.
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=100))
+
+    # Setup Vizier Service.
+    servicer = vizier_server.VizierService()
+    vizier_service_pb2_grpc.add_VizierServiceServicer_to_server(servicer, server)
+    server.add_secure_port(address, grpc.local_server_credentials())
+
+    # Start the server.
+    server.start()
+
+    clients.environment_variables.service_endpoint = address  # Server address.
+    study = clients.Study.from_study_config(
+        study_config, owner='owner', study_id='example_study_id')
+
+     # experiment name 
+    exp_name = FLAGS.workload+ "_num_steps_" + str(FLAGS.num_steps) + "_reward_type+" + FLAGS.reward_formulation
+
+    # append logs to base path
+    log_path = os.path.join(FLAGS.summary_dir, 'EMUKIT_logs', FLAGS.reward_formulation, exp_name)
+
+    # get the current working directory and append the exp name
+    traject_dir = os.path.join(FLAGS.summary_dir, FLAGS.traject_dir, FLAGS.reward_formulation, exp_name)
+
+    # check if log_path exists else create it
+    if not os.path.exists(log_path):
+        os.makedirs(log_path)
+
+    if FLAGS.use_envlogger:
+        if not os.path.exists(traject_dir):
+            os.makedirs(traject_dir)
+    env = wrap_in_envlogger(env, traject_dir)
+
+    count = 0
+    env.reset()
+    suggestions = emukit_designer.suggest(count=flags.FLAGS.num_steps)
+    for suggestion in suggestions:
+        count += 1
+
+        action = {"Bypass": int(str(suggestion.parameters['Bypass'])),
+            "CFU_enable": int(str(suggestion.parameters['CFU_enable'])),
+            "Data_cache_size": int(str(suggestion.parameters['Data_cache_size'])),
+            "Hardware_Divider": int(str(suggestion.parameters['Hardware_Divider'])),
+            "Instruction_cache_size": int(str(suggestion.parameters['Instruction_cache_size'])),
+            "Hardware_Multiplier": int(str(suggestion.parameters['Hardware_Multiplier'])),
+            "Branch_predictor_type": int(str(suggestion.parameters['Branch_predictor_type'])),
+            "Safe_mode_enable": int(str(suggestion.parameters['Safe_mode_enable'])),
+            "Single_Cycle_Shifter": int(str(suggestion.parameters['Single_Cycle_Shifter'])),     
+            "Single_Cycle_Multiplier": int(str(suggestion.parameters['Single_Cycle_Multiplier']))}
+
+        done, reward, info, obs = (env.step(action))
+        fitness_hist['reward'] = reward
+        fitness_hist['action'] = action
+        fitness_hist['obs'] = obs
+        if count == FLAGS.num_steps:
+            done = True
+        log_fitness_to_csv(log_path, fitness_hist)
+        print("Observation: ",obs)
+        final_measurement = vz.Measurement({'Reward': reward})
+        suggestion = suggestion.to_trial()
+        suggestion.complete(final_measurement)
+
+if __name__ == '__main__':
+    app.run(main)