Skip to content

Commit 260dbb4

Browse files
authored
[mlir][sparse] add COO to python tests (llvm#70090)
also typo fix
1 parent 117041d commit 260dbb4

File tree

3 files changed

+11
-8
lines changed

3 files changed

+11
-8
lines changed

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -132,13 +132,15 @@ def main():
132132

133133
# CHECK-LABEL: TEST: testSDDMMM
134134
print("\nTEST: testSDDMMM")
135+
count = 0
135136
with ir.Context() as ctx, ir.Location.unknown():
136-
count = 0
137137
# Loop over various ways to compile and annotate the SDDMM kernel with
138138
# a *single* sparse tensor. Note that we deliberate do not exhaustively
139139
# search the full state space to reduce runtime of the test. It is
140140
# straightforward to adapt the code below to explore more combinations.
141+
# For these simple orderings, dim2lvl and lvl2dim are the same.
141142
levels = [
143+
[st.DimLevelType.compressed_nu, st.DimLevelType.singleton],
142144
[st.DimLevelType.dense, st.DimLevelType.dense],
143145
[st.DimLevelType.dense, st.DimLevelType.compressed],
144146
[st.DimLevelType.compressed, st.DimLevelType.dense],
@@ -154,15 +156,15 @@ def main():
154156
for iwidth in [32]:
155157
for e in [True]:
156158
attr = st.EncodingAttr.get(
157-
level, ordering, None, pwidth, iwidth
159+
level, ordering, ordering, pwidth, iwidth
158160
)
159161
opt = f"parallelization-strategy=none"
160162
compiler = sparse_compiler.SparseCompiler(
161163
options=opt, opt_level=0, shared_libs=[support_lib]
162164
)
163165
build_compile_and_run_SDDMMM(attr, compiler)
164166
count = count + 1
165-
# CHECK: Passed 8 tests
167+
# CHECK: Passed 10 tests
166168
print("Passed ", count, "tests")
167169

168170

mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -115,17 +115,18 @@ def main():
115115

116116
# CHECK-LABEL: TEST: testSpMM
117117
print("\nTEST: testSpMM")
118+
count = 0
118119
with ir.Context() as ctx, ir.Location.unknown():
119-
count = 0
120120
# Loop over various ways to compile and annotate the SpMM kernel with
121121
# a *single* sparse tensor. Note that we deliberate do not exhaustively
122122
# search the full state space to reduce runtime of the test. It is
123123
# straightforward to adapt the code below to explore more combinations.
124-
124+
# For these simple orderings, dim2lvl and lvl2dim are the same.
125125
vl = 1
126126
e = False
127127
opt = f"parallelization-strategy=none"
128128
levels = [
129+
[st.DimLevelType.compressed_nu, st.DimLevelType.singleton],
129130
[st.DimLevelType.dense, st.DimLevelType.dense],
130131
[st.DimLevelType.dense, st.DimLevelType.compressed],
131132
[st.DimLevelType.compressed, st.DimLevelType.dense],
@@ -144,11 +145,11 @@ def main():
144145
for pwidth in bitwidths:
145146
for iwidth in bitwidths:
146147
attr = st.EncodingAttr.get(
147-
level, ordering, None, pwidth, iwidth
148+
level, ordering, ordering, pwidth, iwidth
148149
)
149150
build_compile_and_run_SpMM(attr, compiler)
150151
count = count + 1
151-
# CHECK: Passed 8 tests
152+
# CHECK: Passed 10 tests
152153
print("Passed ", count, "tests")
153154

154155

mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ def __call__(self, module: ir.Module):
2424
self.compile(module)
2525

2626
def compile(self, module: ir.Module):
27-
"""Compiles the module by invoking the sparse copmiler pipeline."""
27+
"""Compiles the module by invoking the sparse compiler pipeline."""
2828
passmanager.PassManager.parse(self.pipeline).run(module.operation)
2929

3030
def jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:

0 commit comments

Comments
 (0)