dialects: (memref) build a dim with an init rather than a static method (#6384)

Author: dipo101

tests/dialects/test_memref.py

@@ -209,13 +209,24 @@ def test_memref_dealloc():
 def test_memref_dim():
     idx = arith.ConstantOp.from_int_and_width(1, IndexType())
     alloc0 = AllocOp.get(i32, 64, [3, 1, 2])
-    dim_1 = memref.DimOp.from_source_and_index(alloc0, idx)
+    dim_1 = memref.DimOp(alloc0, idx)
 
     assert dim_1.source is alloc0.memref
     assert dim_1.index is idx.result
     assert isinstance(dim_1.result.type, IndexType)
 
 
+def test_memref_dim_from_source_and_index_is_deprecated():
+    idx = arith.ConstantOp.from_int_and_width(1, IndexType())
+    alloc0 = AllocOp.get(i32, 64, [3, 1, 2])
+
+    with pytest.deprecated_call():
+        dim_1 = memref.DimOp.from_source_and_index(alloc0, idx)  # pyright: ignore[reportDeprecated]
+
+    assert dim_1.source is alloc0.memref
+    assert dim_1.index is idx.result
+
+
 def test_memref_rank():
     alloc0 = AllocOp.get(i32, 64, [3, 1, 2])
     dim_1 = memref.RankOp.from_memref(alloc0)
@@ -244,10 +255,10 @@ def matmul(args: tuple[BlockArgument, ...]) -> None:
 
             lit0 = arith.ConstantOp.from_int_and_width(0, builtin.IndexType())
             lit1 = arith.ConstantOp.from_int_and_width(1, builtin.IndexType())
-            dim_a0 = memref.DimOp.from_source_and_index(a, lit0)
-            dim_a1 = memref.DimOp.from_source_and_index(a, lit1)
-            dim_b0 = memref.DimOp.from_source_and_index(b, lit0)
-            dim_b1 = memref.DimOp.from_source_and_index(b, lit1)
+            dim_a0 = memref.DimOp(a, lit0)
+            dim_a1 = memref.DimOp(a, lit1)
+            dim_b0 = memref.DimOp(b, lit0)
+            dim_b1 = memref.DimOp(b, lit1)
             out = memref.AllocaOp.get(
                 builtin.f64, None, [DYNAMIC_INDEX] * 2, [dim_a0, dim_b1]
             )

xdsl/dialects/linalg/transforms/tiling.py

@@ -260,11 +260,9 @@ def _build_loop_range(
     assert isa(source_type, MemRefType | TensorType)
     match source_type:
         case MemRefType():
-            dim_op = memref.DimOp.from_source_and_index(source, position_op)
+            dim_op = memref.DimOp(source, position_op)
         case TensorType():
-            dim_op = tensor.DimOp.build(
-                operands=[source, position_op], result_types=[IndexType()]
-            )
+            dim_op = tensor.DimOp(source, position_op)
         case _:
             assert_never(source_type)
 

xdsl/dialects/memref.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
 import abc
-from collections.abc import Iterable, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from typing import ClassVar, cast
 
-from typing_extensions import Self
+from typing_extensions import Self, deprecated
 
 from xdsl.dialects.builtin import (
     DYNAMIC_INDEX,
@@ -489,11 +489,22 @@ class DimOp(IRDLOperation):
 
     assembly_format = "$source `,` $index attr-dict `:` type($source)"
 
+    def __init__(
+        self,
+        source: SSAValue | Operation,
+        index: SSAValue | Operation,
+        attributes: Mapping[str, Attribute] | None = None,
+    ):
+        super().__init__(
+            operands=(source, index), result_types=(IndexType(),), attributes=attributes
+        )
+
     @staticmethod
+    @deprecated("Use DimOp(source, index) instead")
     def from_source_and_index(
         source: SSAValue | Operation, index: SSAValue | Operation
-    ):
-        return DimOp.build(operands=[source, index], result_types=[IndexType()])
+    ) -> DimOp:
+        return DimOp(source, index)
 
 
 @irdl_op_definition

xdsl/transforms/convert_linalg_to_loops.py

@@ -43,7 +43,7 @@ def materialize_loop_bound(
         dim_index_op = arith.ConstantOp.from_int_and_width(dim_index, IndexType())
         rewriter.insert(dim_index_op, insertion_point)
 
-        dim_op = memref.DimOp.from_source_and_index(operand, dim_index_op.result)
+        dim_op = memref.DimOp(operand, dim_index_op.result)
         rewriter.insert(dim_op, insertion_point)
         return dim_op.result
 

xdsl/transforms/convert_memref_to_ptr.py

@@ -105,9 +105,7 @@ def get_strides(
                 arith.ConstantOp.from_int_and_width(i + 1, _index_type)
             )
             dim_idx.result.name_hint = "dim_idx"
-            dim_size = builder.insert(
-                memref.DimOp.from_source_and_index(memref_val, dim_idx.result)
-            ).result
+            dim_size = builder.insert(memref.DimOp(memref_val, dim_idx.result)).result
         prev = strides[i + 1]
         match (prev, dim_size):
             case (int(p), int(d)):

xdsl/transforms/experimental/Apply1DMPIToStencil.py

@@ -49,7 +49,7 @@ def match_and_rewrite(
             0, builtin.IndexType()
         )
         dim_zero_const = arith.ConstantOp(int_attr, builtin.IndexType())
-        dim_zero_size_op = memref.DimOp.from_source_and_index(op.field, dim_zero_const)
+        dim_zero_size_op = memref.DimOp(op.field, dim_zero_const)
         dim_zero_i32_op = arith.IndexCastOp(dim_zero_size_op, builtin.i32)
         dim_zero_i64_op = arith.IndexCastOp(dim_zero_size_op, builtin.i64)