fix test batch accuracy to 100%

Author: WilliamZhang20

test/tpu/torch_backend.py

@@ -1,5 +1,5 @@
 import torch
-from test_tpu import matmul
+from test_tpu import matmul, reset_dut
 from torch._inductor.compile_fx import compile_fx
 from torch._dynamo import register_backend
 from torch.library import custom_op, register_fake
@@ -17,52 +17,75 @@
 def tpu_matmul(a_q: Tensor, b_q: Tensor, bias: Optional[Tensor] = None,
                a_scale: Optional[float] = None, b_scale: Optional[float] = None,
                a_zero: Optional[int] = None, b_zero: Optional[int] = None) -> Tensor:
-    # Prepare centered int8 inputs for the TPU:
-    # 1) convert to int32 so subtraction is safe
-    # 2) subtract zero_point to center around 0
-    # 3) clamp to int8 range and cast back to int8 for the hardware
+    
+    if a_zero is None:
+        a_zero = 0
+    if b_zero is None:
+        b_zero = 0
+    
+    # Convert to int32 for safe arithmetic
     a_i32 = a_q.to(torch.int32)
     b_i32 = b_q.to(torch.int32)
-
-    if a_zero is not None:
-        a_center = a_i32 - int(a_zero)
-    else:
-        a_center = a_i32
-    if b_zero is not None:
-        b_center = b_i32 - int(b_zero)
+    
+    # Subtract zero points
+    a_centered = a_i32 - int(a_zero)
+    b_centered = b_i32 - int(b_zero)
+
+    # For uint8-style quantization (zero_point=-128), shift to signed range
+    shift_amount = 0
+    if a_zero == -128:
+        shift_amount = 128
+        a_hw = (a_centered - shift_amount).clamp(-128, 127).to(torch.int8)
     else:
-        b_center = b_i32
-
-    a_hw = a_center.clamp(-128, 127).to(torch.int8)
-    b_hw = b_center.clamp(-128, 127).to(torch.int8)
-
-    # Call TPU: provide int8 inputs; TPU returns int32 accumulators
+        a_hw = a_centered.clamp(-128, 127).to(torch.int8)
+    
+    b_hw = b_centered.clamp(-128, 127).to(torch.int8)
+    
+    # Call TPU
     future = concurrent.futures.Future()
     async def wrapper():
         try:
+            await reset_dut(dut)
             result_int32 = await matmul(dut, a_hw, b_hw, transpose=True, is_torch=True)
             future.set_result(result_int32)
         except Exception as e:
             future.set_exception(e)
     cocotb.start_soon(wrapper())
     result_int32 = future.result()
-
-    # Dequantize to float using scale product (accumulators are int32)
-    if a_scale is None or b_scale is None:
-        out = result_int32.to(torch.float32)
-    else:
+    
+    # Apply correction for the shift
+    if shift_amount != 0:
+        correction = shift_amount * b_centered.sum(dim=1, dtype=torch.int32)
+        correction = correction.view(1, -1).expand_as(result_int32)
+        result_int32 = result_int32 + correction
+        
+    # Convert to float and scale
+    result_float = result_int32.to(torch.float32)
+    
+    if a_scale is not None and b_scale is not None:
         scale = float(a_scale) * float(b_scale)
-        out = result_int32.to(torch.float32) * scale
-
+        out = result_float * scale
+    else:
+        out = result_float
+    
+    # Add bias
     if bias is not None:
         out = out + bias
-
     return out
 
 @register_fake("tpu::matmul")
 def tpu_matmul_abstract(a: Tensor, b: Tensor, bias: Optional[Tensor] = None, *args, **kwargs) -> Tensor:
-    M, N = a.shape[-1], b.shape[-2]
-    out = a.new_zeros(a.shape[:-1] + (N,), dtype=torch.float32)
+    # a: (..., K)  where K is input features
+    # b: (N, K)    where N is output features (weight matrix, will be transposed)
+    # output: (..., N) after computing a @ b.T
+    
+    # b.shape[0] is N (number of output features)
+    N = b.shape[0]
+    
+    # Output shape: same batch dims as a, but last dim is N
+    out_shape = list(a.shape[:-1]) + [N]
+    out = a.new_zeros(out_shape, dtype=torch.float32)
+    
     if bias is not None:
         out = out + bias
     return out
@@ -76,14 +99,15 @@ def _backend(gm: torch.fx.GraphModule, example_inputs):
         # print("\n=== FX graph received ===")
         # gm.graph.print_tabular()
 
-        # ---- replace every linear but try to preserve quantized buffers + scale info ----
+        # Replace linear operations with TPU matmul
         dequant_op = torch.ops.quantized_decomposed.dequantize_per_tensor.default
+        
         for node in list(gm.graph.nodes):
             if node.target == torch.ops.aten.linear.default:
                 x_node, w_node, bias = node.args
 
                 def unwrap_dequant(n):
-                    # If node is a dequantize_per_tensor node, return (q_tensor, scale, zero_point, node)
+                    # If node is a dequantize_per_tensor node, return quantization params
                     if isinstance(n, torch.fx.Node) and n.target == dequant_op:
                         # args: (quantized_tensor, scale, zero_point, min, max, dtype)
                         q_tensor = n.args[0]
@@ -99,23 +123,25 @@ def unwrap_dequant(n):
                     if x_un and w_un:
                         q_x, x_scale, x_zp, x_deq_node = x_un
                         q_w, w_scale, w_zp, w_deq_node = w_un
+                        
+                        # Verify zero points are 0 for symmetric quantization
+                        print(f"Replacing linear: x_zp={x_zp}, w_zp={w_zp}")
+                        
                         new_node = gm.graph.call_function(
                             torch.ops.tpu.matmul,
                             args=(q_x, q_w, bias, x_scale, w_scale, x_zp, w_zp),
                         )
                         node.replace_all_uses_with(new_node)
                         gm.graph.erase_node(node)
-                        # try to erase the now-unused dequantize nodes
-                        try:
+                        
+                        # Clean up unused dequantize nodes
+                        if len(x_deq_node.users) == 0:
                             gm.graph.erase_node(x_deq_node)
-                        except Exception:
-                            pass
-                        try:
+                        if len(w_deq_node.users) == 0:
                             gm.graph.erase_node(w_deq_node)
-                        except Exception:
-                            pass
                     else:
-                        # Fallback: if inputs are already raw tensors/float, call tpu.matmul with them
+                        # Fallback for non-quantized tensors
+                        print(f"Warning: Linear layer not quantized, using fallback")
                         new_node = gm.graph.call_function(
                             torch.ops.tpu.matmul,
                             args=(x_node, w_node, bias),

test/tpu/train_qat_model.py

@@ -14,7 +14,7 @@
   get_symmetric_quantization_config,
   XNNPACKQuantizer,
 )
-
+from test_tpu import reset_dut
 import cocotb
 from cocotb.clock import Clock
 import numpy as np
@@ -35,7 +35,18 @@ def forward(self, x):
         x = self.relu(x)
         x = self.fc2(x)
         return x
-    
+
+def compute_accuracy(model, dataloader):
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for images, labels in dataloader:
+            outputs = model(images)
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+    return correct / total
+
 def get_quantized_model():
     transform = transforms.Compose([transforms.ToTensor()])
     train_ds = torchvision.datasets.MNIST(root='./data', train=True,
@@ -70,18 +81,22 @@ def get_quantized_model():
     criterion = nn.CrossEntropyLoss()
     optimizer = torch.optim.SGD(prepared_model.parameters(), lr=0.01, momentum=0.9)
 
-    prepared_model = move_exported_model_to_train(prepared_model)
-    for epoch in range(10):
+    for epoch in range(3):
+        prepared_model = move_exported_model_to_train(prepared_model)
+
         for images, labels in train_loader:
             optimizer.zero_grad()
-            out = prepared_model(images) # Input is 1x1x28x28, the graph handles the flatten.
+            out = prepared_model(images)
             loss = criterion(out, labels)
             loss.backward()
             optimizer.step()
 
+        # ---- Accuracy after epoch ----
+        prepared_model = move_exported_model_to_eval(prepared_model)
+        acc = compute_accuracy(prepared_model, train_loader)
+        print(f"Epoch {epoch+1} training accuracy: {acc:.4f}")
+
     print("Training over")
-    
-    prepared_model = move_exported_model_to_eval(prepared_model)
 
     # --- FULL TORCHAO PT2E Convert Step ---
     print("Converting prepared model to quantized model...")
@@ -98,6 +113,7 @@ async def tpu_torch_test(dut):
     model = get_quantized_model()
     clock = Clock(dut.clk, 20, units="ns")
     cocotb.start_soon(clock.start())
+    await reset_dut(dut)
 
     # compile it with backend
     from torch_backend import make_backend