[python] Simplify iree_gpu dialect bindings tests. NFC. (iree-org#19104)

Use a decoration to manage context / module / insertion point creation.

compiler/bindings/python/test/ir/dialects_test.py

@@ -10,136 +10,133 @@
 from iree.compiler.dialects import flow, hal, stream, vm, util, iree_gpu
 
 
-@lambda _: _()
-def gpu_pipeline_options_attr():
-    with ir.Context() as ctx, ir.Location.unknown():
+def run(fn):
+    with ir.Context(), ir.Location.unknown():
         module = ir.Module.create()
         with ir.InsertionPoint(module.body):
-            reorder_attr = iree_gpu.ReorderWorkgroupsStrategyAttr.get(
-                iree_gpu.ReorderWorkgroupsStrategy.Transpose, ctx
-            )
-            assert reorder_attr.value == iree_gpu.ReorderWorkgroupsStrategy.Transpose
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get(
-                True,
-                False,
-                False,
-                reorder_attr,
-            )
-            assert type(gpu_attr) is iree_gpu.PipelineOptionsAttr
-            assert gpu_attr.prefetch_shared_memory
-            assert not gpu_attr.no_reduce_shared_memory_bank_conflicts
-            assert not gpu_attr.use_igemm_convolution
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get(
-                False,
-                True,
-                True,
-                iree_gpu.ReorderWorkgroupsStrategyAttr.get(
-                    iree_gpu.ReorderWorkgroupsStrategy.Transpose, ctx
-                ),
-            )
-            assert not gpu_attr.prefetch_shared_memory
-            assert gpu_attr.no_reduce_shared_memory_bank_conflicts
-            assert gpu_attr.use_igemm_convolution
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get()
-            assert (
-                gpu_attr.prefetch_shared_memory is None
-                and gpu_attr.no_reduce_shared_memory_bank_conflicts is None
-                and gpu_attr.use_igemm_convolution is None
-                and gpu_attr.reorder_workgroups_strategy is None
-            )
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get(True)
-            assert gpu_attr.prefetch_shared_memory
-            assert (
-                gpu_attr.no_reduce_shared_memory_bank_conflicts is None
-                and gpu_attr.use_igemm_convolution is None
-                and gpu_attr.reorder_workgroups_strategy is None
-            )
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get(True, False)
-            assert (
-                gpu_attr.use_igemm_convolution is None
-                and gpu_attr.reorder_workgroups_strategy is None
-            )
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get(True, False, False)
-            assert gpu_attr.reorder_workgroups_strategy is None
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get(
-                no_reduce_shared_memory_bank_conflicts=False
-            )
-            assert (
-                gpu_attr.no_reduce_shared_memory_bank_conflicts is not None
-                and not gpu_attr.no_reduce_shared_memory_bank_conflicts
-            )
-            assert gpu_attr.prefetch_shared_memory is None
-            assert gpu_attr.use_igemm_convolution is None
-            assert gpu_attr.reorder_workgroups_strategy is None
-
-            gpu_attr = iree_gpu.PipelineOptionsAttr.get(
-                reorder_workgroups_strategy=reorder_attr
-            )
-            assert gpu_attr.reorder_workgroups_strategy is not None
-            assert (
-                gpu_attr.reorder_workgroups_strategy.value
-                # unfortunately not `is`
-                == iree_gpu.ReorderWorkgroupsStrategy.Transpose
-            )
-
-
-@lambda _: _()
+            print("\nTEST:", fn.__name__)
+            fn()
+    return fn
+
+
+@run
+def gpu_pipeline_options_attr():
+    reorder_attr = iree_gpu.ReorderWorkgroupsStrategyAttr.get(
+        iree_gpu.ReorderWorkgroupsStrategy.Transpose
+    )
+    assert reorder_attr.value == iree_gpu.ReorderWorkgroupsStrategy.Transpose
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get(
+        True,
+        False,
+        False,
+        reorder_attr,
+    )
+    assert type(gpu_attr) is iree_gpu.PipelineOptionsAttr
+    assert gpu_attr.prefetch_shared_memory
+    assert not gpu_attr.no_reduce_shared_memory_bank_conflicts
+    assert not gpu_attr.use_igemm_convolution
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get(
+        False,
+        True,
+        True,
+        iree_gpu.ReorderWorkgroupsStrategyAttr.get(
+            iree_gpu.ReorderWorkgroupsStrategy.Transpose
+        ),
+    )
+    assert not gpu_attr.prefetch_shared_memory
+    assert gpu_attr.no_reduce_shared_memory_bank_conflicts
+    assert gpu_attr.use_igemm_convolution
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get()
+    assert (
+        gpu_attr.prefetch_shared_memory is None
+        and gpu_attr.no_reduce_shared_memory_bank_conflicts is None
+        and gpu_attr.use_igemm_convolution is None
+        and gpu_attr.reorder_workgroups_strategy is None
+    )
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get(True)
+    assert gpu_attr.prefetch_shared_memory
+    assert (
+        gpu_attr.no_reduce_shared_memory_bank_conflicts is None
+        and gpu_attr.use_igemm_convolution is None
+        and gpu_attr.reorder_workgroups_strategy is None
+    )
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get(True, False)
+    assert (
+        gpu_attr.use_igemm_convolution is None
+        and gpu_attr.reorder_workgroups_strategy is None
+    )
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get(True, False, False)
+    assert gpu_attr.reorder_workgroups_strategy is None
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get(
+        no_reduce_shared_memory_bank_conflicts=False
+    )
+    assert (
+        gpu_attr.no_reduce_shared_memory_bank_conflicts is not None
+        and not gpu_attr.no_reduce_shared_memory_bank_conflicts
+    )
+    assert gpu_attr.prefetch_shared_memory is None
+    assert gpu_attr.use_igemm_convolution is None
+    assert gpu_attr.reorder_workgroups_strategy is None
+
+    gpu_attr = iree_gpu.PipelineOptionsAttr.get(
+        reorder_workgroups_strategy=reorder_attr
+    )
+    assert gpu_attr.reorder_workgroups_strategy is not None
+    assert (
+        gpu_attr.reorder_workgroups_strategy.value
+        # unfortunately not `is`
+        == iree_gpu.ReorderWorkgroupsStrategy.Transpose
+    )
+
+
+@run
 def mma_intrinsic_attr():
-    with ir.Context() as ctx, ir.Location.unknown():
-        module = ir.Module.create()
-        with ir.InsertionPoint(module.body):
-            mma_intrinsic_attr = iree_gpu.MMAIntrinsicAttr.get(
-                iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16, ctx
-            )
-            assert mma_intrinsic_attr is not None
-            assert (
-                str(mma_intrinsic_attr)
-                == "#iree_gpu<mma_intrinsic MFMA_F32_32x32x8_F16>"
-            )
-
-            raw_value = mma_intrinsic_attr.raw_value
-            assert raw_value == iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16
-            value = mma_intrinsic_attr.value
-            assert str(value) == "MFMA_F32_32x32x8_F16"
-            assert int(value) == raw_value
-
-            mma_attr = iree_gpu.MMAAttr.get(raw_value, ctx)
-            assert mma_attr is not None
-
-            f16 = ir.F16Type.get()
-            f32 = ir.F32Type.get()
-            a_type, b_type, c_type = mma_attr.abc_element_types
-            assert a_type == f16
-            assert b_type == f16
-            assert c_type == f32
-
-            vec_4xf16 = ir.VectorType.get((4,), f16)
-            a_vec_type, b_vec_type, _c_vec_type = mma_attr.abc_vector_types
-            assert a_vec_type == vec_4xf16
-            assert b_vec_type == vec_4xf16
-
-            M, N, K = mma_attr.mnk_shape
-            assert M == 32
-            assert N == 32
-            assert K == 8
-
-            assert mma_intrinsic_attr.mma == mma_attr
-
-
-@lambda _: _()
+    mma_intrinsic_attr = iree_gpu.MMAIntrinsicAttr.get(
+        iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16
+    )
+    assert mma_intrinsic_attr is not None
+    assert str(mma_intrinsic_attr) == "#iree_gpu<mma_intrinsic MFMA_F32_32x32x8_F16>"
+
+    raw_value = mma_intrinsic_attr.raw_value
+    assert raw_value == iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16
+    value = mma_intrinsic_attr.value
+    assert str(value) == "MFMA_F32_32x32x8_F16"
+    assert int(value) == raw_value
+
+    mma_attr = iree_gpu.MMAAttr.get(raw_value)
+    assert mma_attr is not None
+
+    f16 = ir.F16Type.get()
+    f32 = ir.F32Type.get()
+    a_type, b_type, c_type = mma_attr.abc_element_types
+    assert a_type == f16
+    assert b_type == f16
+    assert c_type == f32
+
+    vec_4xf16 = ir.VectorType.get((4,), f16)
+    a_vec_type, b_vec_type, _c_vec_type = mma_attr.abc_vector_types
+    assert a_vec_type == vec_4xf16
+    assert b_vec_type == vec_4xf16
+
+    M, N, K = mma_attr.mnk_shape
+    assert M == 32
+    assert N == 32
+    assert K == 8
+
+    assert mma_intrinsic_attr.mma == mma_attr
+
+
+@run
 def lowering_config_attr():
-    with ir.Context() as ctx, ir.Location.unknown():
-        module = ir.Module.create()
-        with ir.InsertionPoint(module.body):
-            attributes = ir.DictAttr.get({"reduction": ir.ArrayAttr.get([])}, ctx)
-            lowering_config = iree_gpu.LoweringConfigAttr.get(attributes, ctx)
-            assert lowering_config is not None
+    attributes = ir.DictAttr.get({"reduction": ir.ArrayAttr.get([])})
+    lowering_config = iree_gpu.LoweringConfigAttr.get(attributes)
+    assert lowering_config is not None
 
-            assert lowering_config.attributes == attributes
+    assert lowering_config.attributes == attributes