From faf78965dd552f8f9571d1e038328a7213adb224 Mon Sep 17 00:00:00 2001
From: Florent Kermarrec <florent@enjoy-digital.fr>
Date: Wed, 25 Nov 2020 15:15:08 +0100
Subject: [PATCH] bench/test_etherbone: allow direct CommUDP use with --udp.

Useful to compare performance with RemoteClient and CommUDP.
---
 bench/test_etherbone.py | 42 +++++++++++++++++++++++++----------------
 1 file changed, 26 insertions(+), 16 deletions(-)

diff --git a/bench/test_etherbone.py b/bench/test_etherbone.py
index 54185e0..33019f6 100755
--- a/bench/test_etherbone.py
+++ b/bench/test_etherbone.py
@@ -13,14 +13,22 @@ import time
 import argparse
 
 from litex import RemoteClient
+from litex.tools.remote.comm_udp import CommUDP
+
+# Constants ----------------------------------------------------------------------------------------
 
 KiB = 1024
 MiB = 1024*KiB
 
+comms = {
+    "cli": RemoteClient,
+    "udp": CommUDP
+}
+
 # Identifier Test ----------------------------------------------------------------------------------
 
-def ident_test(port):
-    wb = RemoteClient(port=port)
+def ident_test(comm, port):
+    wb = comms[comm](port=port, csr_csv="csr.csv")
     wb.open()
 
     fpga_identifier = ""
@@ -37,8 +45,8 @@ def ident_test(port):
 
 # Access Test --------------------------------------------------------------------------------------
 
-def access_test(port):
-    wb = RemoteClient(port=port)
+def access_test(comm, port):
+    wb = comms[comm](port=port, csr_csv="csr.csv")
     wb.open()
 
     data = 0x12345678
@@ -52,8 +60,8 @@ def access_test(port):
 
 # SRAM Test ----------------------------------------------------------------------------------------
 
-def sram_test(port):
-    wb = RemoteClient(port=port)
+def sram_test(comm, port):
+    wb = comms[comm](port=port, csr_csv="csr.csv")
     wb.open()
 
     def mem_dump(base, length):
@@ -100,12 +108,12 @@ def sram_test(port):
 
 # Speed Test ---------------------------------------------------------------------------------------
 
-def speed_test(port):
-    wb = RemoteClient(port=port)
+def speed_test(comm, port):
+    wb = comms[comm](port=port, csr_csv="csr.csv")
     wb.open()
 
     test_size  = 16*KiB
-    burst_size = 128
+    burst_size = 200
 
     print("Testing write speed... ", end="")
     start = time.time()
@@ -113,15 +121,15 @@ def speed_test(port):
         wb.write(wb.mems.sram.base, [j for j in range(burst_size)])
     end = time.time()
     duration = (end - start)
-    print("{:3.2f} KiB/s".format(test_size/(duration*KiB)))
+    print("{:8.2f} KiB/s".format(test_size/(duration*KiB)))
 
-    print("Testing read speed... ", end="")
+    print("Testing read speed...  ", end="")
     start = time.time()
     for i in range(test_size//(4*burst_size)):
         wb.read(wb.mems.sram.base, length=burst_size)
     end = time.time()
     duration = (end - start)
-    print("{:3.2f} KiB/s".format(test_size/(duration*KiB)))
+    print("{:8.2f} KiB/s".format(test_size/(duration*KiB)))
 
     wb.close()
 
@@ -130,6 +138,7 @@ def speed_test(port):
 def main():
     parser = argparse.ArgumentParser(description="LiteEth Etherbone test utility")
     parser.add_argument("--port",   default="1234",      help="Host bind port")
+    parser.add_argument("--udp",    action="store_true", help="Use CommUDP directly instead of RemoteClient")
     parser.add_argument("--ident",  action="store_true", help="Read FPGA identifier")
     parser.add_argument("--access", action="store_true", help="Test single Write/Read access over Etherbone")
     parser.add_argument("--sram",   action="store_true", help="Test SRAM access over Etherbone")
@@ -137,18 +146,19 @@ def main():
     args = parser.parse_args()
 
     port = int(args.port, 0)
+    comm = "cli" if not args.udp else "udp"
 
     if args.ident:
-        ident_test(port=port)
+        ident_test(comm=comm, port=port)
 
     if args.access:
-        access_test(port=port)
+        access_test(comm=comm, port=port)
 
     if args.sram:
-        sram_test(port=port)
+        sram_test(comm=comm, port=port)
 
     if args.speed:
-        speed_test(port=port)
+        speed_test(comm=comm, port=port)
 
 if __name__ == "__main__":
     main()