Merge pull request #53 from InfiniTensor/pass-block-sizes-to-the-second-mm-arrangement-in-addmm-arrangement

voltjia · web-flow · commit 0d6c7f79ba39 · 2025-10-13T16:29:34.000+08:00
Pass block sizes to the second `mm.arrangement` in `addmm.arrangement`
diff --git a/src/ntops/kernels/addmm.py b/src/ntops/kernels/addmm.py
@@ -38,7 +38,13 @@ def arrangement(
     )
 
     x_arranged, y_arranged, output_arranged, _ = mm.arrangement(
-        x, y, output, input_precision
+        x,
+        y,
+        output,
+        input_precision,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
     )
 
     input_precision_arranged = input_precision

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,13 @@ def arrangement(`
`38`	`38`	`)`
`39`	`39`
`40`	`40`	`x_arranged, y_arranged, output_arranged, _ = mm.arrangement(`
`41`		`- x, y, output, input_precision`
	`41`	`+ x,`
	`42`	`+ y,`
	`43`	`+ output,`
	`44`	`+ input_precision,`
	`45`	`+ block_size_m=block_size_m,`
	`46`	`+ block_size_n=block_size_n,`
	`47`	`+ block_size_k=block_size_k,`
`42`	`48`	`)`
`43`	`49`
`44`	`50`	`input_precision_arranged = input_precision`