Merge branch 'main' into gpu_wait

requirements.txt

@@ -5,8 +5,8 @@ toml<0.11
 pytest-cov
 coverage<8.0.0
 ipykernel
-pytest<8.1
-nbval<0.11
+pytest<8.2
+nbval<0.12
 filecheck<0.0.25
 lit<18.0.0
 pre-commit==3.6.2

tests/dialects/test_bufferization.py

@@ -0,0 +1,31 @@
+from xdsl.dialects.bufferization import AllocTensorOp, ToTensorOp
+from xdsl.dialects.builtin import MemRefType, TensorType, UnitAttr, f64
+from xdsl.dialects.test import TestOp
+
+
+def test_to_tensor():
+    memref_t = MemRefType(f64, [10, 20, 30])
+    tensor_t = TensorType(f64, [10, 20, 30])
+    memref_v = TestOp(result_types=[memref_t]).res[0]
+
+    to_tensor = ToTensorOp(memref_v)
+    assert to_tensor.memref == memref_v
+    assert to_tensor.restrict is None
+    assert to_tensor.writable is None
+    assert to_tensor.tensor.type == tensor_t
+
+    to_tensor = ToTensorOp(memref_v, writable=True, restrict=True)
+    assert to_tensor.memref == memref_v
+    assert to_tensor.restrict == UnitAttr()
+    assert to_tensor.writable == UnitAttr()
+    assert to_tensor.tensor.type == tensor_t
+
+
+def test_alloc_tensor_static():
+    t = TensorType(f64, [10, 20, 30])
+    alloc_tensor = AllocTensorOp(t)
+
+    assert alloc_tensor.tensor.type == t
+    assert alloc_tensor.dynamic_sizes == ()
+    assert alloc_tensor.copy is None
+    assert alloc_tensor.size_hint is None

tests/dialects/test_tensor.py

@@ -0,0 +1,67 @@
+from xdsl.dialects.builtin import DenseArrayBase, TensorType, f64, i64
+from xdsl.dialects.tensor import ExtractSliceOp, InsertSliceOp
+from xdsl.dialects.test import TestOp
+
+
+def test_extract_slice_static():
+    input_t = TensorType(f64, [10, 20, 30])
+    input_v = TestOp(result_types=[input_t]).res[0]
+
+    extract_slice = ExtractSliceOp.from_static_parameters(input_v, [1, 2, 3], [4, 5, 6])
+
+    assert extract_slice.source is input_v
+    assert extract_slice.static_offsets == DenseArrayBase.from_list(i64, [1, 2, 3])
+    assert extract_slice.static_sizes == DenseArrayBase.from_list(i64, [4, 5, 6])
+    assert extract_slice.static_strides == DenseArrayBase.from_list(i64, [1, 1, 1])
+    assert extract_slice.offsets == ()
+    assert extract_slice.sizes == ()
+    assert extract_slice.strides == ()
+    assert extract_slice.result.type == TensorType(f64, [4, 5, 6])
+
+    extract_slice = ExtractSliceOp.from_static_parameters(
+        input_v, [1, 2, 3], [4, 5, 6], [8, 9, 10]
+    )
+
+    assert extract_slice.source is input_v
+    assert extract_slice.static_offsets == DenseArrayBase.from_list(i64, [1, 2, 3])
+    assert extract_slice.static_sizes == DenseArrayBase.from_list(i64, [4, 5, 6])
+    assert extract_slice.static_strides == DenseArrayBase.from_list(i64, [8, 9, 10])
+    assert extract_slice.offsets == ()
+    assert extract_slice.sizes == ()
+    assert extract_slice.strides == ()
+    assert extract_slice.result.type == TensorType(f64, [4, 5, 6])
+
+
+def test_insert_slice_static():
+    source_t = TensorType(f64, [10, 20])
+    source_v = TestOp(result_types=[source_t]).res[0]
+    dest_t = TensorType(f64, [10, 20, 30])
+    dest_v = TestOp(result_types=[dest_t]).res[0]
+
+    insert_slice = InsertSliceOp.from_static_parameters(
+        source_v, dest_v, [1, 2], [4, 5]
+    )
+
+    assert insert_slice.source is source_v
+    assert insert_slice.dest is dest_v
+    assert insert_slice.static_offsets == DenseArrayBase.from_list(i64, [1, 2])
+    assert insert_slice.static_sizes == DenseArrayBase.from_list(i64, [4, 5])
+    assert insert_slice.static_strides == DenseArrayBase.from_list(i64, [1, 1])
+    assert insert_slice.offsets == ()
+    assert insert_slice.sizes == ()
+    assert insert_slice.strides == ()
+    assert insert_slice.result.type == dest_t
+
+    insert_slice = InsertSliceOp.from_static_parameters(
+        source_v, dest_v, [1, 2], [4, 5], [8, 9]
+    )
+
+    assert insert_slice.source is source_v
+    assert insert_slice.dest is dest_v
+    assert insert_slice.static_offsets == DenseArrayBase.from_list(i64, [1, 2])
+    assert insert_slice.static_sizes == DenseArrayBase.from_list(i64, [4, 5])
+    assert insert_slice.static_strides == DenseArrayBase.from_list(i64, [8, 9])
+    assert insert_slice.offsets == ()
+    assert insert_slice.sizes == ()
+    assert insert_slice.strides == ()
+    assert insert_slice.result.type == dest_t

tests/filecheck/dialects/linalg/linalg_ops.mlir

@@ -38,10 +38,10 @@ linalg.generic {indexing_maps = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1)
 // CHECK-GENERIC-NEXT:  ^0(%{{.*}}: f32, %{{.*}}: f32):
 // CHECK-GENERIC-NEXT:      "linalg.yield"(%{{.*}}) : (f32) -> ()
 // CHECK-GENERIC-NEXT:  }) : (f32, memref<1x256xf32>) -> ()
-// CHECK-GENERIC-NEXT:  "linalg.generic"(%0, %1) <{"indexing_maps" = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> (d0, d1)>], "iterator_types" = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>], "operandSegmentSizes" = array<i32: 1, 1>}> ({
+// CHECK-GENERIC-NEXT:  "linalg.generic"(%0, %1) <{"indexing_maps" = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> (d0, d1)>], "iterator_types" = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>], "doc" = "a_docstring", "library_call" = "a_library_call", "operandSegmentSizes" = array<i32: 1, 1>}> ({
 // CHECK-GENERIC-NEXT:  ^1(%arg3_1 : f32, %arg4_1 : f32):
 // CHECK-GENERIC-NEXT:    "linalg.yield"(%arg3_1) : (f32) -> ()
-// CHECK-GENERIC-NEXT:  }) {"doc" = "a_docstring", "library_call" = "a_library_call"} : (f32, memref<1x256xf32>) -> ()
+// CHECK-GENERIC-NEXT:  }) : (f32, memref<1x256xf32>) -> ()
 
 // CHECK-GENERIC:       "linalg.generic"(%0, %1) <{"indexing_maps" = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> (d0, d1)>], "iterator_types" = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>], "operandSegmentSizes" = array<i32: 1, 1>}> ({
 // CHECK-GENERIC-NEXT:  ^{{.*}}(%{{.*}}: f32, %{{.*}}: f32):

tests/filecheck/mlir-conversion/with-mlir/dialects/bufferization/ops.mlir

@@ -0,0 +1,13 @@
+// RUN: xdsl-opt %s | mlir-opt --mlir-print-op-generic | xdsl-opt | filecheck %s
+
+module{
+    %t = "bufferization.alloc_tensor"() <{"operandSegmentSizes" = array<i32: 0, 0, 0>}> : () -> tensor<10x20x30xf64>
+    %m = "test.op"() : () -> memref<30x20x10xf32>
+    %m_t = "bufferization.to_tensor"(%m) : (memref<30x20x10xf32>) -> tensor<30x20x10xf32>
+}
+
+// CHECK:       builtin.module {
+// CHECK-NEXT:    %0 = "bufferization.alloc_tensor"() <{"operandSegmentSizes" = array<i32: 0, 0, 0>}> : () -> tensor<10x20x30xf64>
+// CHECK-NEXT:    %1 = "test.op"() : () -> memref<30x20x10xf32>
+// CHECK-NEXT:    %2 = "bufferization.to_tensor"(%1) : (memref<30x20x10xf32>) -> tensor<30x20x10xf32>
+// CHECK-NEXT:  }

tests/filecheck/mlir-conversion/with-mlir/dialects/tensor/ops.mlir

@@ -7,6 +7,8 @@
 %src = "test.op"() {"value" = dense<1.000000e-01> : tensor<4x1xf32>} : () -> tensor<4x1xf32>
 %shape = "test.op"() : () -> (tensor<1xi32>)
 %t4 = tensor.reshape %src(%shape) : (tensor<4x1xf32>, tensor<1xi32>) -> tensor<4xf32>
+%inserted_slice = "tensor.insert_slice"(%t2, %t1) <{"static_offsets" = array<i64: 0, 1>, "static_sizes" = array<i64: 1, 2>, "static_strides" = array<i64: 1, 1>, "operandSegmentSizes" = array<i32: 1, 1, 0, 0, 0>}> : (tensor<2xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
+%extracted_slice = "tensor.extract_slice"(%t1) <{"static_offsets" = array<i64: 0, 1>, "static_sizes" = array<i64: 1, 2>, "static_strides" = array<i64: 1, 1>, "operandSegmentSizes" = array<i32: 1, 0, 0, 0>}> : (tensor<2x3xf32>) -> tensor<2xf32>
 
 
 // CHECK:       module {
@@ -17,4 +19,6 @@
 // CHECK-NEXT:  %4 = "test.op"() {value = dense<1.000000e-01> : tensor<4x1xf32>} : () -> tensor<4x1xf32>
 // CHECK-NEXT:  %5 = "test.op"() : () -> tensor<1xi32>
 // CHECK-NEXT:  %reshape = tensor.reshape %4(%5) : (tensor<4x1xf32>, tensor<1xi32>) -> tensor<4xf32>
+// CHECK-NEXT:  %inserted_slice = tensor.insert_slice %1 into %0[0, 1] [1, 2] [1, 1] : tensor<2xf32> into tensor<2x3xf32>
+// CHECK-NEXT:  %extracted_slice = tensor.extract_slice %0[0, 1] [1, 2] [1, 1] : tensor<2x3xf32> to tensor<2xf32>
 // CHECK-NEXT: }

tests/filecheck/mlir-conversion/with-mlir/mlir_opt.mlir

@@ -11,6 +11,9 @@
 // RUN: xdsl-opt %s -p mlir-opt{generic=false\ arguments='--cse','--mlir-print-op-generic'} --print-op-generic | filecheck %s
 // RUN: xdsl-opt %s -p mlir-opt{generic=true\ arguments='--cse','--mlir-print-op-generic'} --print-op-generic | filecheck %s
 
+// Check that manually passing an executable works
+// RUN: xdsl-opt %s -p mlir-opt{executable=mlir-opt\ generic=true\ arguments='--cse','--mlir-print-op-generic'} --print-op-generic | filecheck %s
+
 "builtin.module"() ({
   "func.func"() ({
     %0 = "arith.constant"() {"value" = 1 : i32} : () -> i32

tests/filecheck/mlir-conversion/with-mlir/mlir_opt_fail.mlir

@@ -1,5 +1,6 @@
 // RUN: xdsl-opt %s -p mlir-opt{arguments='--hello','--mlir-print-op-generic'} --print-op-generic --verify-diagnostics | filecheck %s
 // RUN: xdsl-opt %s -p mlir-opt[this-probably-will-never-be-an-MLIR-pass-name] --print-op-generic --verify-diagnostics | filecheck %s
+// RUN: xdsl-opt %s -p mlir-opt{executable='"false"'} --print-op-generic --verify-diagnostics | filecheck %s
 
 "builtin.module"() ({
   "func.func"() ({

xdsl/dialects/bufferization.py

@@ -0,0 +1,94 @@
+from typing import Any
+
+from xdsl.dialects.builtin import (
+    ContainerType,
+    IndexType,
+    MemRefType,
+    ShapedType,
+    TensorType,
+    UnitAttr,
+    UnrankedMemrefType,
+    UnrankedTensorType,
+)
+from xdsl.ir import Attribute, Dialect, Operation, SSAValue
+from xdsl.irdl import (
+    AnyOf,
+    AttrSizedOperandSegments,
+    IRDLOperation,
+    irdl_op_definition,
+    operand_def,
+    opt_operand_def,
+    opt_prop_def,
+    result_def,
+    var_operand_def,
+)
+
+
+@irdl_op_definition
+class AllocTensorOp(IRDLOperation):
+    name = "bufferization.alloc_tensor"
+
+    dynamic_sizes = var_operand_def(IndexType())
+    copy = opt_operand_def(AnyOf((TensorType, UnrankedTensorType)))
+    size_hint = opt_operand_def(IndexType())
+
+    tensor = result_def(AnyOf((TensorType, UnrankedTensorType)))
+
+    irdl_options = [AttrSizedOperandSegments(as_property=True)]
+
+    def __init__(
+        self,
+        result_type: Attribute,
+        dynamic_sizes: list[Operation | SSAValue] | None = None,
+        copy: SSAValue | Operation | None = None,
+        size_hint: SSAValue | Operation | None = None,
+    ):
+        super().__init__(
+            operands=(dynamic_sizes, copy, size_hint),
+            result_types=(result_type,),
+        )
+
+
+@irdl_op_definition
+class ToTensorOp(IRDLOperation):
+    name = "bufferization.to_tensor"
+
+    memref = operand_def(AnyOf((MemRefType, UnrankedMemrefType)))
+    tensor = result_def(AnyOf((TensorType, UnrankedTensorType)))
+    writable = opt_prop_def(UnitAttr)
+    restrict = opt_prop_def(UnitAttr)
+
+    def __init__(
+        self,
+        memref: SSAValue | Operation,
+        restrict: bool = False,
+        writable: bool = False,
+    ):
+        memref_v = SSAValue.get(memref)
+        memref_t = memref_v.type
+        if not isinstance(memref_t, ContainerType):
+            raise ValueError(f"Expected ContainerType, got {memref_t}")
+        if not isinstance(memref_t, ShapedType):
+            raise ValueError(f"Expected ShapedType, got {memref_t}")
+        properties = dict[str, Attribute]()
+        if restrict:
+            properties["restrict"] = UnitAttr()
+        if writable:
+            properties["writable"] = UnitAttr()
+        super().__init__(
+            operands=(memref,),
+            result_types=(
+                TensorType[Any](memref_t.get_element_type(), memref_t.get_shape()),
+            ),
+            properties=properties,
+        )
+
+
+Bufferization = Dialect(
+    "bufferization",
+    [
+        AllocTensorOp,
+        ToTensorOp,
+    ],
+    [],
+)

xdsl/dialects/linalg.py

@@ -35,7 +35,7 @@
     irdl_attr_definition,
     irdl_op_definition,
     operand_def,
-    opt_attr_def,
+    opt_prop_def,
     prop_def,
     region_def,
     var_operand_def,
@@ -96,8 +96,8 @@ class Generic(IRDLOperation):
     # Trait attributes
     indexing_maps: ArrayAttr[AffineMapAttr] = prop_def(ArrayAttr[AffineMapAttr])
     iterator_types: ArrayAttr[IteratorTypeAttr] = prop_def(ArrayAttr[IteratorTypeAttr])
-    doc: StringAttr | None = opt_attr_def(StringAttr)
-    library_call: StringAttr | None = opt_attr_def(StringAttr)
+    doc: StringAttr | None = opt_prop_def(StringAttr)
+    library_call: StringAttr | None = opt_prop_def(StringAttr)
 
     irdl_options = [AttrSizedOperandSegments(as_property=True)]
 
@@ -118,8 +118,6 @@ def __init__(
             properties={
                 "indexing_maps": ArrayAttr(indexing_maps),
                 "iterator_types": ArrayAttr(iterator_types),
-            },
-            attributes={
                 "doc": doc,
                 "library_call": library_call,
             },
@@ -238,7 +236,14 @@ def print(self, printer: Printer):
 
         if self.res:
             printer.print_string(" -> ")
-            printer.print_list(self.res, lambda res: printer.print_attribute(res.type))
+            if len(self.res) == 1:
+                printer.print_attribute(self.res[0].type)
+            else:
+                printer.print("(")
+                printer.print_list(
+                    self.res, lambda res: printer.print_attribute(res.type)
+                )
+                printer.print(")")
 
     @classmethod
     def parse(cls, parser: Parser) -> Self:
@@ -360,7 +365,6 @@ def parse(cls, parser: Parser) -> Self:
             doc,
             library_call,
         )
-        generic.attributes |= attrs
         generic.attributes |= extra_attrs
 
         return generic