-
Notifications
You must be signed in to change notification settings - Fork 27
/
Copy pathneq2lut.py
181 lines (159 loc) · 9.12 KB
/
neq2lut.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
# 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, \
module_list_to_verilog_module
from train import configs, model_config, dataset_config, test
from dataset import JetSubstructureDataset
from models import JetSubstructureNeqModel, JetSubstructureLutModel
from logicnets.synthesis import synthesize_and_get_resource_counts
from logicnets.util import proc_postsynth_file
other_options = {
"cuda": None,
"log_dir": None,
"checkpoint": None,
"generate_bench": False,
"add_registers": False,
"simulate_pre_synthesis_verilog": False,
"simulate_post_synthesis_verilog": False,
}
if __name__ == "__main__":
parser = ArgumentParser(description="Synthesize convert a PyTorch trained model into verilog")
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/processed-pythia82-lhc13-all-pt1-50k-r1_h022_e0175_t220_nonu_truth.z',
help="The file to use as the dataset input (default: %(default)s)")
parser.add_argument('--clock-period', type=float, default=1.0,
help="Target clock frequency to use during Vivado synthesis (default: %(default)s)")
parser.add_argument('--dataset-config', type=str, default='config/yaml_IP_OP_config.yml',
help="The file to use to configure the input dataset (default: %(default)s)")
parser.add_argument('--dataset-split', type=str, default='test', choices=['train', 'test'],
help="Dataset to use for evaluation (default: %(default)s)")
parser.add_argument('--log-dir', type=str, default='./log',
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('--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,
help="Dump I/O to the verilog LUT to a text file in the log directory (default: %(default)s)")
parser.add_argument('--add-registers', action='store_true', default=False,
help="Add registers between each layer in generated verilog (default: %(default)s)")
parser.add_argument('--simulate-pre-synthesis-verilog', action='store_true', default=False,
help="Simulate the verilog generated by LogicNets (default: %(default)s)")
parser.add_argument('--simulate-post-synthesis-verilog', action='store_true', default=False,
help="Simulate the post-synthesis verilog produced by vivado (default: %(default)s)")
parser.add_argument('--fpga-part', type=str, default="xcu280-fsvh2892-2L-e",
help="Part to use for Vivado project (default: %(default)s)")
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[args.dataset_split] = JetSubstructureDataset(dataset_cfg['dataset_file'], dataset_cfg['dataset_config'], split=args.dataset_split)
test_loader = DataLoader(dataset[args.dataset_split], batch_size=config['batch_size'], shuffle=False)
# Instantiate the PyTorch model
x, y = dataset[args.dataset_split][0]
model_cfg['input_length'] = len(x)
model_cfg['output_length'] = len(y)
model = JetSubstructureNeqModel(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 on baseline model...")
model.eval()
baseline_accuracy, baseline_avg_roc_auc = test(model, test_loader, cuda=False)
print("Baseline accuracy: %f" % (baseline_accuracy))
print("Baseline AVG ROC AUC: %f" % (baseline_avg_roc_auc))
# Instantiate LUT-based model
lut_model = JetSubstructureLutModel(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 on LUT-based model...")
lut_inference(lut_model)
lut_model.eval()
lut_accuracy, lut_avg_roc_auc = test(lut_model, test_loader, cuda=False)
print("LUT-Based Model accuracy: %f" % (lut_accuracy))
print("LUT-Based AVG ROC AUC: %f" % (lut_avg_roc_auc))
modelSave = { 'model_dict': lut_model.state_dict(),
'test_accuracy': lut_accuracy,
'test_avg_roc_auc': lut_avg_roc_auc}
torch.save(modelSave, options_cfg["log_dir"] + "/lut_based_model.pth")
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"])
print("Top level entity stored at: %s/logicnet.v ..." % (options_cfg["log_dir"]))
if args.dump_io:
io_filename = options_cfg["log_dir"] + f"io_{args.dataset_split}.txt"
with open(io_filename, 'w') as f:
pass # Create an empty file.
print(f"Dumping verilog I/O to {io_filename}...")
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"])
verilog_accuracy, verilog_avg_roc_auc = test(lut_model, test_loader, cuda=False)
print("Verilog-Based Model accuracy: %f" % (verilog_accuracy))
print("Verilog-Based AVG ROC AUC: %f" % (verilog_avg_roc_auc))
print("Running out-of-context synthesis")
ret = synthesize_and_get_resource_counts(options_cfg["log_dir"], "logicnet", fpga_part=args.fpga_part, clk_period_ns=args.clock_period, post_synthesis = 1)
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"])
post_synth_accuracy, post_synth_avg_roc_auc = test(lut_model, test_loader, cuda=False)
print("Post-synthesis Verilog-Based Model accuracy: %f" % (post_synth_accuracy))
print("Post-synthesis Verilog-Based AVG ROC AUC: %f" % (post_synth_avg_roc_auc))