[Feat] Track which cares are used during inference

examples/cybersecurity/dump_luts.py

@@ -0,0 +1,119 @@
+#  Copyright (C) 2021 Xilinx, Inc
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import os
+from argparse import ArgumentParser
+
+import torch
+from torch.utils.data import DataLoader
+
+from logicnets.nn import    generate_truth_tables, \
+                            lut_inference, \
+                            save_luts, \
+                            module_list_to_verilog_module
+
+from train import configs, model_config, dataset_config, other_options, test
+from dataset import get_preqnt_dataset
+from models import UnswNb15NeqModel, UnswNb15LutModel
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Generate histograms of states used throughout LogicNets")
+    parser.add_argument('--arch', type=str, choices=configs.keys(), default="jsc-s",
+        help="Specific the neural network model to use (default: %(default)s)")
+    parser.add_argument('--batch-size', type=int, default=None, metavar='N',
+        help="Batch size for evaluation (default: %(default)s)")
+    parser.add_argument('--input-bitwidth', type=int, default=None,
+        help="Bitwidth to use at the input (default: %(default)s)")
+    parser.add_argument('--hidden-bitwidth', type=int, default=None,
+        help="Bitwidth to use for activations in hidden layers (default: %(default)s)")
+    parser.add_argument('--output-bitwidth', type=int, default=None,
+        help="Bitwidth to use at the output (default: %(default)s)")
+    parser.add_argument('--input-fanin', type=int, default=None,
+        help="Fanin to use at the input (default: %(default)s)")
+    parser.add_argument('--hidden-fanin', type=int, default=None,
+        help="Fanin to use for the hidden layers (default: %(default)s)")
+    parser.add_argument('--output-fanin', type=int, default=None,
+        help="Fanin to use at the output (default: %(default)s)")
+    parser.add_argument('--hidden-layers', nargs='+', type=int, default=None,
+        help="A list of hidden layer neuron sizes (default: %(default)s)")
+    parser.add_argument('--dataset-file', type=str, default='data/unsw_nb15_binarized.npz',
+        help="The file to use as the dataset input (default: %(default)s)")
+    parser.add_argument('--log-dir', type=str, default='./log',
+        help="A location to store the calculated histograms (default: %(default)s)")
+    parser.add_argument('--checkpoint', type=str, required=True,
+        help="The checkpoint file which contains the model weights")
+    args = parser.parse_args()
+    defaults = configs[args.arch]
+    options = vars(args)
+    del options['arch']
+    config = {}
+    for k in options.keys():
+        config[k] = options[k] if options[k] is not None else defaults[k] # Override defaults, if specified.
+
+    if not os.path.exists(config['log_dir']):
+        os.makedirs(config['log_dir'])
+
+    # Split up configuration options to be more understandable
+    model_cfg = {}
+    for k in model_config.keys():
+        model_cfg[k] = config[k]
+    dataset_cfg = {}
+    for k in dataset_config.keys():
+        dataset_cfg[k] = config[k]
+    options_cfg = {}
+    for k in other_options.keys():
+        if k == 'cuda':
+            continue
+        options_cfg[k] = config[k]
+
+    # Fetch the test set
+    dataset = {}
+    dataset['train'] = get_preqnt_dataset(dataset_cfg['dataset_file'], split='train')
+    train_loader = DataLoader(dataset["train"], batch_size=config['batch_size'], shuffle=False)
+
+    # Instantiate the PyTorch model
+    x, y = dataset['train'][0]
+    dataset_length = len(dataset['train'])
+    model_cfg['input_length'] = len(x)
+    model_cfg['output_length'] = 1
+    model = UnswNb15NeqModel(model_cfg)
+
+    # Load the model weights
+    checkpoint = torch.load(options_cfg['checkpoint'], map_location='cpu')
+    model.load_state_dict(checkpoint['model_dict'])
+
+    # Test the PyTorch model
+    print("Running inference of baseline model on training set (%d examples)..." % (dataset_length))
+    model.eval()
+    baseline_accuracy = test(model, train_loader, cuda=False)
+    print("Baseline accuracy: %f" % (baseline_accuracy))
+
+    # Instantiate LUT-based model
+    lut_model = UnswNb15LutModel(model_cfg)
+    lut_model.load_state_dict(checkpoint['model_dict'])
+
+    # Generate the truth tables in the LUT module
+    print("Converting to NEQs to LUTs...")
+    generate_truth_tables(lut_model, verbose=True)
+
+    # Test the LUT-based model
+    print("Running inference of LUT-based model training set (%d examples)..." % (dataset_length))
+    lut_inference(lut_model, track_used_luts=True)
+    lut_model.eval()
+    lut_accuracy = test(lut_model, train_loader, cuda=False)
+    print("LUT-Based Model accuracy: %f" % (lut_accuracy))
+    print("Saving LUTs to %s... " % (options_cfg["log_dir"] + "/luts.pth"))
+    save_luts(lut_model, options_cfg["log_dir"] + "/luts.pth")
+    print("Done!")
+

examples/cybersecurity/models.py

@@ -63,13 +63,15 @@ def __init__(self, model_config):
         self.verilog_dir = None
         self.top_module_filename = None
         self.dut = None
+        self.verify = True
         self.logfile = None
 
-    def verilog_inference(self, verilog_dir, top_module_filename, logfile: bool = False, add_registers: bool = False):
+    def verilog_inference(self, verilog_dir, top_module_filename, logfile: bool = False, add_registers: bool = False, verify: bool = True):
         self.verilog_dir = realpath(verilog_dir)
         self.top_module_filename = top_module_filename
-        self.dut = PyVerilator.build(f"{self.verilog_dir}/{self.top_module_filename}", verilog_path=[self.verilog_dir], build_dir=f"{self.verilog_dir}/verilator")
+        self.dut = PyVerilator.build(f"{self.verilog_dir}/{self.top_module_filename}", verilog_path=[self.verilog_dir], build_dir=f"{self.verilog_dir}/verilator", command_args=("--x-assign","0",))
         self.is_verilog_inference = True
+        self.verify = verify
         self.logfile = logfile
         if add_registers:
             self.latency = len(self.num_neurons)
@@ -95,21 +97,22 @@ def verilog_forward(self, x):
         self.dut.io.clk = 0
         for i in range(x.shape[0]):
             x_i = x[i,:]
-            y_i = self.pytorch_forward(x[i:i+1,:])[0]
             xv_i = list(map(lambda z: input_quant.get_bin_str(z), x_i))
-            ys_i = list(map(lambda z: output_quant.get_bin_str(z), y_i))
             xvc_i = reduce(lambda a,b: a+b, xv_i[::-1])
-            ysc_i = reduce(lambda a,b: a+b, ys_i[::-1])
             self.dut["M0"] = int(xvc_i, 2)
             for j in range(self.latency + 1):
                 #print(self.dut.io.M5)
                 res = self.dut[f"M{num_layers}"]
                 result = f"{res:0{int(total_output_bits)}b}"
                 self.dut.io.clk = 1
                 self.dut.io.clk = 0
-            expected = f"{int(ysc_i,2):0{int(total_output_bits)}b}"
             result = f"{res:0{int(total_output_bits)}b}"
-            assert(expected == result)
+            if self.verify:
+                y_i = self.pytorch_forward(x[i:i+1,:])[0]
+                ys_i = list(map(lambda z: output_quant.get_bin_str(z), y_i))
+                ysc_i = reduce(lambda a,b: a+b, ys_i[::-1])
+                expected = f"{int(ysc_i,2):0{int(total_output_bits)}b}"
+                assert(expected == result)
             res_split = [result[i:i+output_bitwidth] for i in range(0, len(result), output_bitwidth)][::-1]
             yv_i = torch.Tensor(list(map(lambda z: int(z, 2), res_split)))
             y[i,:] = yv_i

examples/cybersecurity/neq2lut.py

@@ -20,7 +20,8 @@
 
 from logicnets.nn import    generate_truth_tables, \
                             lut_inference, \
-                            module_list_to_verilog_module
+                            module_list_to_verilog_module, \
+                            load_histograms
 from logicnets.synthesis import synthesize_and_get_resource_counts
 from logicnets.util import proc_postsynth_file
 
@@ -34,6 +35,8 @@
     "checkpoint": None,
     "generate_bench": False,
     "add_registers": False,
+    "histograms": None,
+    "freq_thresh": None,
     "simulate_pre_synthesis_verilog": False,
     "simulate_post_synthesis_verilog": False,
 }
@@ -68,6 +71,10 @@
         help="A location to store the log output of the training run and the output model (default: %(default)s)")
     parser.add_argument('--checkpoint', type=str, required=True,
         help="The checkpoint file which contains the model weights")
+    parser.add_argument('--histograms', type=str, default=None,
+        help="The checkpoint histograms of LUT usage (default: %(default)s)")
+    parser.add_argument('--freq-thresh', type=int, default=None,
+        help="Threshold to use to include this truth table into the model (default: %(default)s)")
     parser.add_argument('--generate-bench', action='store_true', default=False,
         help="Generate the truth table in BENCH format as well as verilog (default: %(default)s)")
     parser.add_argument('--dump-io', action='store_true', default=False,
@@ -141,9 +148,12 @@
                     'test_accuracy': lut_accuracy}
 
     torch.save(modelSave, options_cfg["log_dir"] + "/lut_based_model.pth")
+    if options_cfg["histograms"] is not None:
+        luts = torch.load(options_cfg["histograms"])
+        load_histograms(lut_model, luts)
 
     print("Generating verilog in %s..." % (options_cfg["log_dir"]))
-    module_list_to_verilog_module(lut_model.module_list, "logicnet", options_cfg["log_dir"], generate_bench=options_cfg["generate_bench"], add_registers=options_cfg["add_registers"])
+    module_list_to_verilog_module(lut_model.module_list, "logicnet", options_cfg["log_dir"], generate_bench=options_cfg["generate_bench"], add_registers=options_cfg["add_registers"], freq_thresh=options_cfg["freq_thresh"])
     print("Top level entity stored at: %s/logicnet.v ..." % (options_cfg["log_dir"]))
 
     if args.dump_io:
@@ -154,9 +164,10 @@
     else:
         io_filename = None
 
+
     if args.simulate_pre_synthesis_verilog:
         print("Running inference simulation of Verilog-based model...")
-        lut_model.verilog_inference(options_cfg["log_dir"], "logicnet.v", logfile=io_filename, add_registers=options_cfg["add_registers"])
+        lut_model.verilog_inference(options_cfg["log_dir"], "logicnet.v", logfile=io_filename, add_registers=options_cfg["add_registers"], verify=options_cfg["freq_thresh"] is None or options_cfg["freq_thresh"] == 0)
         verilog_accuracy = test(lut_model, test_loader, cuda=False)
         print("Verilog-Based Model accuracy: %f" % (verilog_accuracy))
 
@@ -166,7 +177,7 @@
     if args.simulate_post_synthesis_verilog:
         print("Running post-synthesis inference simulation of Verilog-based model...")
         proc_postsynth_file(options_cfg["log_dir"])
-        lut_model.verilog_inference(options_cfg["log_dir"]+"/post_synth", "logicnet_post_synth.v", io_filename, add_registers=options_cfg["add_registers"])
+        lut_model.verilog_inference(options_cfg["log_dir"]+"/post_synth", "logicnet_post_synth.v", io_filename, add_registers=options_cfg["add_registers"], verify=options_cfg["freq_thresh"] is None or options_cfg["freq_thresh"] == 0)
         post_synth_accuracy = test(lut_model, test_loader, cuda=False)
         print("Post-synthesis Verilog-Based Model accuracy: %f" % (post_synth_accuracy))
 

examples/cybersecurity/train.py

@@ -44,6 +44,8 @@
         "learning_rate": 1e-1,
         "seed": 25,
         "checkpoint": None,
+        "histograms": None,
+        "freq_thresh": None,
     },
     "nid-s-comp": {
         "hidden_layers": [49, 7],
@@ -59,6 +61,8 @@
         "learning_rate": 1e-1,
         "seed": 81,
         "checkpoint": None,
+        "histograms": None,
+        "freq_thresh": None,
     },
     "nid-m": {
         "hidden_layers": [593, 256, 128, 128],
@@ -74,6 +78,8 @@
         "learning_rate": 1e-1,
         "seed": 20,
         "checkpoint": None,
+        "histograms": None,
+        "freq_thresh": None,
     },
     "nid-m-comp": {
         "hidden_layers": [593, 256, 49, 7],
@@ -89,6 +95,8 @@
         "learning_rate": 1e-1,
         "seed": 40,
         "checkpoint": None,
+        "histograms": None,
+        "freq_thresh": None,
     },
     "nid-l": {
         "hidden_layers": [593, 100, 100, 100],
@@ -104,6 +112,8 @@
         "learning_rate": 1e-1,
         "seed": 2,
         "checkpoint": None,
+        "histograms": None,
+        "freq_thresh": None,
     },
     "nid-l-comp": {
         "hidden_layers": [593, 100, 25, 5],
@@ -119,6 +129,8 @@
         "learning_rate": 1e-1,
         "seed": 83,
         "checkpoint": None,
+        "histograms": None,
+        "freq_thresh": None,
     },
 }
 

examples/jet_substructure/config/yaml_IP_OP_config.yml

@@ -45,5 +45,5 @@ L1Reg: 0.0001
 NormalizeInputs: 1 
 InputType: Dense
 ApplyPca: false
-PcaDimensions: 10
+PcaDimensions: 12