[mlir][sparse] enable Python BSR test (#72325)

Author: aartbik

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

@@ -32,11 +32,28 @@ def boilerplate(attr: st.EncodingAttr):
 def expected(id_map):
     """Returns expected contents of output.
 
+    +-----+-----+-----+-----+-----+
+    | 1 0 | . . | . . | . . | 0 3 |
+    | 0 2 | . . | . . | . . | 0 0 |
+    +-----+-----+-----+-----+-----+
+    | . . | . . | . . | . . | . . |
+    | . . | . . | . . | . . | . . |
+    +-----+-----+-----+-----+-----+
+    | . . | . . | 5 0 | . . | . . |
+    | . . | . . | 0 0 | . . | . . |
+    +-----+-----+-----+-----+-----+
+    | . . | . . | . . | . . | . . |
+    | . . | . . | . . | . . | . . |
+    +-----+-----+-----+-----+-----+
+    | 0 0 | . . | . . | . . | . . |
+    | 4 0 | . . | . . | . . | . . |
+    +-----+-----+-----+-----+-----+
+
     Output appears as dimension coordinates but lexicographically
-    sorted by level coordinates.
+    sorted by level coordinates. For BSR, the blocks are filled.
     """
-    return (
-        f"""# extended FROSTT format
+    if id_map is 0:
+        return f"""# extended FROSTT format
 2 5
 10 10
 1 1 1
@@ -45,8 +62,8 @@ def expected(id_map):
 5 5 5
 10 1 4
 """
-        if id_map
-        else f"""# extended FROSTT format
+    if id_map is 1:
+        return f"""# extended FROSTT format
 2 5
 10 10
 1 1 1
@@ -55,7 +72,28 @@ def expected(id_map):
 5 5 5
 1 10 3
 """
-    )
+    if id_map is 2:
+        return f"""# extended FROSTT format
+2 16
+10 10
+1 1 1
+1 2 0
+2 1 0
+2 2 2
+1 9 0
+1 10 3
+2 9 0
+2 10 0
+5 5 5
+5 6 0
+6 5 0
+6 6 0
+9 1 0
+9 2 0
+10 1 4
+10 2 0
+"""
+    raise AssertionError("unexpected id_map")
 
 
 def build_compile_and_run_output(attr: st.EncodingAttr, compiler, expected):
@@ -93,10 +131,10 @@ def main():
             [st.DimLevelType.compressed, st.DimLevelType.compressed],
         ]
         orderings = [
-            (ir.AffineMap.get_permutation([0, 1]), True),
-            (ir.AffineMap.get_permutation([1, 0]), False),
+            (ir.AffineMap.get_permutation([0, 1]), 0),
+            (ir.AffineMap.get_permutation([1, 0]), 1),
         ]
-        bitwidths = [8, 16, 32, 64]
+        bitwidths = [8, 64]
         compiler = sparse_compiler.SparseCompiler(
             options="", opt_level=2, shared_libs=[support_lib]
         )
@@ -135,11 +173,10 @@ def main():
         l3 = ir.AffineDimExpr.get(3)
         lvl2dim = ir.AffineMap.get(4, 0, [2 * l0 + l2, 2 * l1 + l3])
         attr = st.EncodingAttr.get(level, dim2lvl, lvl2dim, 0, 0)
-        # TODO: enable this one CONVERSION on BSR is working
-        # build_compile_and_run_output(attr, compiler, block_expected())
+        build_compile_and_run_output(attr, compiler, expected(2))
         count = count + 1
 
-    # CHECK: Passed 33 tests
+    # CHECK: Passed 17 tests
     print("Passed", count, "tests")