tmc.py: add track of stallguard

Signed-off-by: Timofey Titovets <nefelim4ag@gmail.com>

4 files changed, 102 insertions, 0 deletions

diff --git a/klippy/extras/tmc.py b/klippy/extras/tmc.py
index 1d8599e2..aad19cf4 100644
--- a/klippy/extras/tmc.py
+++ b/klippy/extras/tmc.py
@@ -5,6 +5,7 @@
 # This file may be distributed under the terms of the GNU GPLv3 license.
 import logging, collections
 import stepper
+from . import bulk_sensor
 
 
 ######################################################################
@@ -220,6 +221,96 @@ class TMCErrorCheck:
             self.last_drv_fields = {n: v for n, v in fields.items() if v}
         return {'drv_status': self.last_drv_fields, 'temperature': temp}
 
+######################################################################
+# Record driver status
+######################################################################
+
+class TMCStallguardDump:
+    def __init__(self, config, mcu_tmc):
+        self.printer = config.get_printer()
+        self.stepper_name = ' '.join(config.get_name().split()[1:])
+        self.mcu_tmc = mcu_tmc
+        self.mcu = self.mcu_tmc.get_mcu()
+        self.fields = self.mcu_tmc.get_fields()
+        self.sg2_supp = False
+        self.sg4_reg_name = None
+        # It is possible to support TMC2660, just disable it for now
+        if not self.fields.all_fields.get("DRV_STATUS", None):
+            return
+        # Collect driver capabilities
+        if self.fields.all_fields["DRV_STATUS"].get("sg_result", None):
+            self.sg2_supp = True
+        # New drivers have separate register for SG4 result
+        if self.mcu_tmc.name_to_reg.get("SG_RESULT", 0):
+            self.sg4_reg_name = "SG_RESULT"
+        # 2240 supports both SG2 & SG4
+        if self.sg4_reg_name is None:
+            if self.mcu_tmc.name_to_reg.get("SG4_RESULT", 0):
+                self.sg4_reg_name = "SG4_RESULT"
+        # TMC2208
+        if self.sg2_supp is None and self.sg4_reg_name is None:
+            return
+        self.optimized_spi = False
+        # Bulk API
+        self.samples = []
+        self.query_timer = None
+        self.error = None
+        self.batch_bulk = bulk_sensor.BatchBulkHelper(
+            self.printer, self._dump, self._start, self._stop)
+        api_resp = {'header': ('time', 'sg_result', 'cs_actual')}
+        self.batch_bulk.add_mux_endpoint("tmc/stallguard_dump", "name",
+                                         self.stepper_name, api_resp)
+    def _start(self):
+        self.error = None
+        status = self.mcu_tmc.get_register_raw("DRV_STATUS")
+        if status.get("spi_status"):
+            self.optimized_spi = True
+        reactor = self.printer.get_reactor()
+        self.query_timer = reactor.register_timer(self._query_tmc,
+                                                  reactor.NOW)
+    def _stop(self):
+        self.printer.get_reactor().unregister_timer(self.query_timer)
+        self.query_timer = None
+        self.samples = []
+    def _query_tmc(self, eventtime):
+        sg_result = -1
+        cs_actual = -1
+        recv_time = eventtime
+        try:
+            if self.optimized_spi or self.sg4_reg_name == "SG4_RESULT":
+                #TMC2130/TMC5160/TMC2240
+                status = self.mcu_tmc.get_register_raw("DRV_STATUS")
+                reg_val = status["data"]
+                cs_actual = self.fields.get_field("cs_actual", reg_val)
+                sg_result = self.fields.get_field("sg_result", reg_val)
+                is_stealth = self.fields.get_field("stealth", reg_val)
+                recv_time = status["#receive_time"]
+                if is_stealth and self.sg4_reg_name == "SG4_RESULT":
+                    sg4_ret = self.mcu_tmc.get_register_raw("SG4_RESULT")
+                    sg_result = sg4_ret["data"]
+                    recv_time = sg4_ret["#receive_time"]
+            else:
+                # TMC2209
+                if self.sg4_reg_name == "SG_RESULT":
+                    sg4_ret = self.mcu_tmc.get_register_raw("SG_RESULT")
+                    sg_result = sg4_ret["data"]
+                    recv_time = sg4_ret["#receive_time"]
+        except self.printer.command_error as e:
+            self.error = e
+            return self.printer.get_reactor().NEVER
+        print_time = self.mcu.estimated_print_time(recv_time)
+        self.samples.append((print_time, sg_result, cs_actual))
+        if self.optimized_spi:
+            return eventtime + 0.001
+        # UART queried as fast as possible
+        return eventtime + 0.005
+    def _dump(self, eventtime):
+            if self.error:
+                raise self.error
+            samples = self.samples
+            self.samples = []
+            return {"data": samples}
+
 
 ######################################################################
 # G-Code command helpers
@@ -233,6 +324,7 @@ class TMCCommandHelper:
         self.mcu_tmc = mcu_tmc
         self.current_helper = current_helper
         self.echeck_helper = TMCErrorCheck(config, mcu_tmc)
+        self.record_helper = TMCStallguardDump(config, mcu_tmc)
         self.fields = mcu_tmc.get_fields()
         self.read_registers = self.read_translate = None
         self.toff = None
diff --git a/klippy/extras/tmc2130.py b/klippy/extras/tmc2130.py
index b41eba02..181fe07a 100644
--- a/klippy/extras/tmc2130.py
+++ b/klippy/extras/tmc2130.py
@@ -232,6 +232,8 @@ class MCU_TMC_SPI_chain:
         pr = pr[(self.chain_len - chain_pos) * 5 :
                 (self.chain_len - chain_pos + 1) * 5]
         return (pr[1] << 24) | (pr[2] << 16) | (pr[3] << 8) | pr[4]
+    def get_mcu(self):
+        return self.spi.get_mcu()
 
 # Helper to setup an spi daisy chain bus from settings in a config section
 def lookup_tmc_spi_chain(config):
@@ -292,6 +294,8 @@ class MCU_TMC_SPI:
             "Unable to write tmc spi '%s' register %s" % (self.name, reg_name))
     def get_tmc_frequency(self):
         return self.tmc_frequency
+    def get_mcu(self):
+        return self.tmc_spi.get_mcu()
 
 
 ######################################################################
diff --git a/klippy/extras/tmc2660.py b/klippy/extras/tmc2660.py
index 7fd01546..c97f1eff 100644
--- a/klippy/extras/tmc2660.py
+++ b/klippy/extras/tmc2660.py
@@ -231,6 +231,8 @@ class MCU_TMC2660_SPI:
             self.spi.spi_send(msg, minclock)
     def get_tmc_frequency(self):
         return None
+    def get_mcu(self):
+        return self.spi.get_mcu()
 
 
 ######################################################################
diff --git a/klippy/extras/tmc_uart.py b/klippy/extras/tmc_uart.py
index 93935042..fa0d6262 100644
--- a/klippy/extras/tmc_uart.py
+++ b/klippy/extras/tmc_uart.py
@@ -187,6 +187,8 @@ class MCU_TMC_uart_bitbang:
             self.analog_mux.activate(instance_id)
         msg = self._encode_write(0xf5, addr, reg | 0x80, val)
         self.tmcuart_send_cmd.send([self.oid, msg, 0], minclock=minclock)
+    def get_mcu(self):
+        return self.mcu
 
 # Lookup a (possibly shared) tmc uart
 def lookup_tmc_uart_bitbang(config, max_addr):
@@ -261,3 +263,5 @@ class MCU_TMC_uart:
             "Unable to write tmc uart '%s' register %s" % (self.name, reg_name))
     def get_tmc_frequency(self):
         return self.tmc_frequency
+    def get_mcu(self):
+        return self.mcu_uart.get_mcu()