MCPcopy Create free account
hub / github.com/KnowingNothing/MatmulTutorial / test_swizzle

Function test_swizzle

cutlass.py/swizzle.py:30–92  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

28
29
30def test_swizzle():
31
32 def get_ind_matrix(rows, cols):
33 return [[(x, y) for y in range(cols)] for x in range(rows)]
34
35 def get_row_major_ind(x, y, rows, cols):
36 return x * cols + y
37
38 def get_row_major_tuple(xy, rows, cols):
39 return (xy // cols, xy % cols)
40
41 def get_col_major_ind(x, y, rows, cols):
42 return x + y * rows
43
44 def get_col_major_tuple(xy, rows, cols):
45 return (xy % rows, xy // rows)
46
47 def print_matrix(mtx, rows, cols, func=lambda x: x, prompt=""):
48 print(prompt)
49 for x in range(rows):
50 for y in range(cols):
51 item = mtx[x][y]
52 item = func(item)
53 print(item, end=" ")
54 print()
55
56 # Swizzle<3, 4, 4>
57 print("Swizzle<3,4,3>")
58 rows = 8
59 cols = 128
60 mtx = get_ind_matrix(rows, cols)
61 print_matrix(mtx, rows, cols, prompt="Original")
62 print()
63 swizzle = Swizzle(3, 4, 3)
64 print_matrix(
65 mtx,
66 rows,
67 cols,
68 lambda tp: get_row_major_tuple(
69 swizzle(get_row_major_ind(tp[0], tp[1], rows, cols)), rows, cols
70 ),
71 prompt="After swizzle",
72 )
73 print()
74
75 # Swizzle<2, 0, -2>
76 print("Swizzle<2,0,-2>")
77 rows = 4
78 cols = 4
79 mtx = get_ind_matrix(rows, cols)
80 print_matrix(mtx, rows, cols, prompt="Original")
81 print()
82 swizzle = Swizzle(2, 0, -2)
83 print_matrix(
84 mtx,
85 rows,
86 cols,
87 lambda tp: get_row_major_tuple(

Callers 1

swizzle.pyFile · 0.85

Calls 5

get_ind_matrixFunction · 0.85
get_row_major_tupleFunction · 0.85
get_row_major_indFunction · 0.85
print_matrixFunction · 0.70
SwizzleClass · 0.70

Tested by

no test coverage detected