Arm backend: Use MXFP op config values directly

Use the weight dtype and block size on MXFP linear and Conv2d directly
instead of accessing them through the config.

Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>
Change-Id: Iffc96d00a887ffbcdf94f3af6d535b4d538451fb

Author: martinlsm

backends/arm/ao_ext/ops/mxfp_conv2d_op.py

@@ -206,11 +206,12 @@ def __init__(
         padding: tuple[int, int],
         dilation: tuple[int, int],
         groups: int,
-        config: MXFPOpConfig,
+        weight_dtype: MXFPDType,
+        block_size: int,
     ) -> None:
         super().__init__()
-        self.config = config
-        self.weight_dtype = mxfp_dtype_to_str(config.weight_dtype)
+        self.weight_dtype = mxfp_dtype_to_str(weight_dtype)
+        self.block_size = block_size
 
         self.register_buffer("weight_qdata", weight_qdata, persistent=True)
         self.register_buffer("weight_scale", weight_scale, persistent=True)
@@ -241,7 +242,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             list(self.padding),
             list(self.dilation),
             self.groups,
-            self.config.block_size,
+            self.block_size,
             self.weight_dtype,
         )
 
@@ -283,5 +284,6 @@ def transform_conv2d_to_mxfp(
         padding,
         dilation,
         module.groups,
-        config,
+        config.weight_dtype,
+        config.block_size,
     )

backends/arm/ao_ext/ops/mxfp_linear_op.py

@@ -137,11 +137,12 @@ def __init__(
         weight_qdata: torch.Tensor,
         weight_scale: torch.Tensor,
         bias: torch.Tensor | None,
-        config: MXFPOpConfig,
+        weight_dtype: MXFPDType,
+        block_size: int,
     ) -> None:
         super().__init__()
-        self.config = config
-        self.weight_dtype = mxfp_dtype_to_str(config.weight_dtype)
+        self.weight_dtype = mxfp_dtype_to_str(weight_dtype)
+        self.block_size = block_size
 
         self.register_buffer("weight_qdata", weight_qdata, persistent=True)
         self.register_buffer("weight_scale", weight_scale, persistent=True)
@@ -163,7 +164,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             self.weight_qdata,
             self.weight_scale,
             self.bias,
-            self.config.block_size,
+            self.block_size,
             self.weight_dtype,
         )
 
@@ -195,4 +196,10 @@ def transform_linear_to_mxfp(
     weight_scale = weight_scale.unsqueeze(0)
 
     bias = module.bias.detach().to(torch.float32) if module.bias is not None else None
-    return MXFPLinearOp(weight_qdata, weight_scale, bias, config)
+    return MXFPLinearOp(
+        weight_qdata,
+        weight_scale,
+        bias,
+        config.weight_dtype,
+        config.block_size,
+    )