Skip to content

Linalg

linalg

LinalgFunctions dataclass

Bases: InterpreterFunctions

Source code in xdsl/interpreters/linalg.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
@register_impls
class LinalgFunctions(InterpreterFunctions):
    @impl(linalg.GenericOp)
    def run_generic(
        self, interpreter: Interpreter, op: linalg.GenericOp, args: tuple[Any, ...]
    ) -> PythonValues:
        if op.library_call is not None:
            raise NotImplementedError(
                "library_call not yet supported in linalg.generic interpreter"
            )
        return run_linalg_structured_op(interpreter, op, args)

    @impl_terminator(linalg.YieldOp)
    def run_yield(
        self, interpreter: Interpreter, op: linalg.YieldOp, args: tuple[Any, ...]
    ):
        return ReturnedValues(args), ()

    @impl(linalg.AddOp)
    def run_add(
        self, interpreter: Interpreter, op: linalg.AddOp, args: tuple[Any, ...]
    ) -> tuple[Any, ...]:
        return run_linalg_structured_op(interpreter, op, args)

    @impl(linalg.FillOp)
    def run_fill(
        self, interpreter: Interpreter, op: linalg.FillOp, args: tuple[Any, ...]
    ) -> tuple[Any, ...]:
        operand, res = args[0], args[1]
        assert isinstance(operand, ShapedArray)
        assert isinstance(res, ShapedArray)
        operand = cast(ShapedArray[float], operand)
        res = cast(ShapedArray[float], res)
        if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
            raise NotImplementedError()
        for i in range(len(res.data)):
            res.data_ptr[i] = operand.data_ptr[0]
        if len(op.results) > 0:
            return (res,)
        return ()

    @impl(linalg.MulOp)
    def run_mul(
        self, interpreter: Interpreter, op: linalg.MulOp, args: tuple[Any, ...]
    ) -> tuple[Any, ...]:
        return run_linalg_structured_op(interpreter, op, args)

    @impl(linalg.TransposeOp)
    def run_transpose(
        self, interpreter: Interpreter, op: linalg.TransposeOp, args: tuple[Any, ...]
    ) -> tuple[Any, ...]:
        operand, res = args[0], args[1]
        assert isinstance(operand, ShapedArray)
        assert isinstance(res, ShapedArray)
        operand = cast(ShapedArray[float], operand)
        res = cast(ShapedArray[float], res)
        if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
            raise NotImplementedError()
        assert len(operand.shape) == 2
        assert len(res.shape) == 2
        rows, cols = operand.shape
        for i in range(rows):
            for j in range(cols):
                res.data_ptr[j * rows + i] = operand.data_ptr[i * cols + j]
        if len(op.results) > 0:
            return (res,)
        return ()

    @impl(linalg.MatmulOp)
    def run_mat_mul(
        self, interpreter: Interpreter, op: linalg.MatmulOp, args: tuple[Any, ...]
    ) -> tuple[Any, ...]:
        return run_linalg_structured_op(interpreter, op, args)

    @impl(linalg.PoolingNchwMaxOp)
    def run_pooling_nchw_max(
        self,
        interpreter: Interpreter,
        op: linalg.PoolingNchwMaxOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        input, kernel_filter, res = args[0], args[1], args[2]
        assert isinstance(input, ShapedArray)
        assert isinstance(kernel_filter, ShapedArray)
        assert isinstance(res, ShapedArray)
        input = cast(ShapedArray[float], input)
        kernel_filter = cast(ShapedArray[float], kernel_filter)
        res = cast(ShapedArray[float], res)
        if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
            raise NotImplementedError()
        strides_type = op.strides.type
        assert isinstance(strides_type, TensorType)
        (strides_shape,) = strides_type.get_shape()
        strides = op.strides.get_values()
        if strides_shape != 2:
            raise NotImplementedError("Only 2d max pooling supported")

        m_height, m_width = input.shape[2:]
        ky, kx = kernel_filter.shape[0], kernel_filter.shape[1]

        # convert input into a numpy like array
        input_data = [
            [input.data[r * m_width + c] for c in range(m_width)]
            for r in range(m_height)
        ]

        output: list[float] = []
        for k in range(0, m_height - ky + 1, strides[0]):
            for l in range(0, m_width - kx + 1, strides[0]):
                block_max_value = float("-inf")
                for i in range(k, k + ky):
                    for j in range(l, l + kx):
                        block_max_value = max(block_max_value, input_data[i][j])
                output.append(block_max_value)
        for i in range(len(output)):
            res.data_ptr[i] = output[i]
        if len(op.results) > 0:
            return (res,)
        return ()

    @impl(linalg.Conv2DNchwFchwOp)
    def run_conv_2d_nchw_fchw(
        self,
        interpreter: Interpreter,
        op: linalg.Conv2DNchwFchwOp,
        args: tuple[Any, ...],
    ) -> tuple[Any, ...]:
        input, kernel_filter, res = args[0], args[1], args[2]
        assert isinstance(input, ShapedArray)
        assert isinstance(kernel_filter, ShapedArray)
        assert isinstance(res, ShapedArray)
        input = cast(ShapedArray[float], input)
        kernel_filter = cast(ShapedArray[float], kernel_filter)
        res = cast(ShapedArray[float], res)
        if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
            raise NotImplementedError()
        m_height, m_width = input.shape[2:]
        ky, kx = kernel_filter.shape[2], kernel_filter.shape[3]
        strides = op.strides.get_values()
        # convert input into a numpy like array
        input_data = [
            [input.data[r * m_width + c] for c in range(m_width)]
            for r in range(m_height)
        ]
        # convert kernel into a numpy like array
        kernel_data = [
            [
                kernel_filter.data[r * kernel_filter.shape[3] + c]
                for c in range(kernel_filter.shape[3])
            ]
            for r in range(kernel_filter.shape[2])
        ]
        output: list[float] = []
        for k in range(0, m_height - ky + 1, strides[0]):
            for l in range(0, m_width - kx + 1, strides[0]):
                conv_value: float = 0.0
                for i in range(k, k + ky):
                    for j in range(l, l + kx):
                        conv_value += input_data[i][j] * kernel_data[i - k][j - l]
                output.append(conv_value)
        for i in range(len(output)):
            res.data_ptr[i] = output[i]
        if len(op.results) > 0:
            return (res,)
        return ()

run_generic(interpreter: Interpreter, op: linalg.GenericOp, args: tuple[Any, ...]) -> PythonValues

Source code in xdsl/interpreters/linalg.py 
68
69
70
71
72
73
74
75
76
@impl(linalg.GenericOp)
def run_generic(
    self, interpreter: Interpreter, op: linalg.GenericOp, args: tuple[Any, ...]
) -> PythonValues:
    if op.library_call is not None:
        raise NotImplementedError(
            "library_call not yet supported in linalg.generic interpreter"
        )
    return run_linalg_structured_op(interpreter, op, args)

run_yield(interpreter: Interpreter, op: linalg.YieldOp, args: tuple[Any, ...])

Source code in xdsl/interpreters/linalg.py 
78
79
80
81
82
@impl_terminator(linalg.YieldOp)
def run_yield(
    self, interpreter: Interpreter, op: linalg.YieldOp, args: tuple[Any, ...]
):
    return ReturnedValues(args), ()

run_add(interpreter: Interpreter, op: linalg.AddOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Source code in xdsl/interpreters/linalg.py 
84
85
86
87
88
@impl(linalg.AddOp)
def run_add(
    self, interpreter: Interpreter, op: linalg.AddOp, args: tuple[Any, ...]
) -> tuple[Any, ...]:
    return run_linalg_structured_op(interpreter, op, args)

run_fill(interpreter: Interpreter, op: linalg.FillOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Source code in xdsl/interpreters/linalg.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@impl(linalg.FillOp)
def run_fill(
    self, interpreter: Interpreter, op: linalg.FillOp, args: tuple[Any, ...]
) -> tuple[Any, ...]:
    operand, res = args[0], args[1]
    assert isinstance(operand, ShapedArray)
    assert isinstance(res, ShapedArray)
    operand = cast(ShapedArray[float], operand)
    res = cast(ShapedArray[float], res)
    if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
        raise NotImplementedError()
    for i in range(len(res.data)):
        res.data_ptr[i] = operand.data_ptr[0]
    if len(op.results) > 0:
        return (res,)
    return ()

run_mul(interpreter: Interpreter, op: linalg.MulOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Source code in xdsl/interpreters/linalg.py 
107
108
109
110
111
@impl(linalg.MulOp)
def run_mul(
    self, interpreter: Interpreter, op: linalg.MulOp, args: tuple[Any, ...]
) -> tuple[Any, ...]:
    return run_linalg_structured_op(interpreter, op, args)

run_transpose(interpreter: Interpreter, op: linalg.TransposeOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Source code in xdsl/interpreters/linalg.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
@impl(linalg.TransposeOp)
def run_transpose(
    self, interpreter: Interpreter, op: linalg.TransposeOp, args: tuple[Any, ...]
) -> tuple[Any, ...]:
    operand, res = args[0], args[1]
    assert isinstance(operand, ShapedArray)
    assert isinstance(res, ShapedArray)
    operand = cast(ShapedArray[float], operand)
    res = cast(ShapedArray[float], res)
    if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
        raise NotImplementedError()
    assert len(operand.shape) == 2
    assert len(res.shape) == 2
    rows, cols = operand.shape
    for i in range(rows):
        for j in range(cols):
            res.data_ptr[j * rows + i] = operand.data_ptr[i * cols + j]
    if len(op.results) > 0:
        return (res,)
    return ()

run_mat_mul(interpreter: Interpreter, op: linalg.MatmulOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Source code in xdsl/interpreters/linalg.py 
134
135
136
137
138
@impl(linalg.MatmulOp)
def run_mat_mul(
    self, interpreter: Interpreter, op: linalg.MatmulOp, args: tuple[Any, ...]
) -> tuple[Any, ...]:
    return run_linalg_structured_op(interpreter, op, args)

run_pooling_nchw_max(interpreter: Interpreter, op: linalg.PoolingNchwMaxOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Source code in xdsl/interpreters/linalg.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
@impl(linalg.PoolingNchwMaxOp)
def run_pooling_nchw_max(
    self,
    interpreter: Interpreter,
    op: linalg.PoolingNchwMaxOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    input, kernel_filter, res = args[0], args[1], args[2]
    assert isinstance(input, ShapedArray)
    assert isinstance(kernel_filter, ShapedArray)
    assert isinstance(res, ShapedArray)
    input = cast(ShapedArray[float], input)
    kernel_filter = cast(ShapedArray[float], kernel_filter)
    res = cast(ShapedArray[float], res)
    if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
        raise NotImplementedError()
    strides_type = op.strides.type
    assert isinstance(strides_type, TensorType)
    (strides_shape,) = strides_type.get_shape()
    strides = op.strides.get_values()
    if strides_shape != 2:
        raise NotImplementedError("Only 2d max pooling supported")

    m_height, m_width = input.shape[2:]
    ky, kx = kernel_filter.shape[0], kernel_filter.shape[1]

    # convert input into a numpy like array
    input_data = [
        [input.data[r * m_width + c] for c in range(m_width)]
        for r in range(m_height)
    ]

    output: list[float] = []
    for k in range(0, m_height - ky + 1, strides[0]):
        for l in range(0, m_width - kx + 1, strides[0]):
            block_max_value = float("-inf")
            for i in range(k, k + ky):
                for j in range(l, l + kx):
                    block_max_value = max(block_max_value, input_data[i][j])
            output.append(block_max_value)
    for i in range(len(output)):
        res.data_ptr[i] = output[i]
    if len(op.results) > 0:
        return (res,)
    return ()

run_conv_2d_nchw_fchw(interpreter: Interpreter, op: linalg.Conv2DNchwFchwOp, args: tuple[Any, ...]) -> tuple[Any, ...]

Source code in xdsl/interpreters/linalg.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
@impl(linalg.Conv2DNchwFchwOp)
def run_conv_2d_nchw_fchw(
    self,
    interpreter: Interpreter,
    op: linalg.Conv2DNchwFchwOp,
    args: tuple[Any, ...],
) -> tuple[Any, ...]:
    input, kernel_filter, res = args[0], args[1], args[2]
    assert isinstance(input, ShapedArray)
    assert isinstance(kernel_filter, ShapedArray)
    assert isinstance(res, ShapedArray)
    input = cast(ShapedArray[float], input)
    kernel_filter = cast(ShapedArray[float], kernel_filter)
    res = cast(ShapedArray[float], res)
    if not all(res.data_ptr[i] == 0.0 for i in range(len(res.data))):
        raise NotImplementedError()
    m_height, m_width = input.shape[2:]
    ky, kx = kernel_filter.shape[2], kernel_filter.shape[3]
    strides = op.strides.get_values()
    # convert input into a numpy like array
    input_data = [
        [input.data[r * m_width + c] for c in range(m_width)]
        for r in range(m_height)
    ]
    # convert kernel into a numpy like array
    kernel_data = [
        [
            kernel_filter.data[r * kernel_filter.shape[3] + c]
            for c in range(kernel_filter.shape[3])
        ]
        for r in range(kernel_filter.shape[2])
    ]
    output: list[float] = []
    for k in range(0, m_height - ky + 1, strides[0]):
        for l in range(0, m_width - kx + 1, strides[0]):
            conv_value: float = 0.0
            for i in range(k, k + ky):
                for j in range(l, l + kx):
                    conv_value += input_data[i][j] * kernel_data[i - k][j - l]
            output.append(conv_value)
    for i in range(len(output)):
        res.data_ptr[i] = output[i]
    if len(op.results) > 0:
        return (res,)
    return ()

run_linalg_structured_op(interpreter: Interpreter, op: linalg.LinalgStructuredOperation, args: tuple[ShapedArray[float] | float, ...])

Helper function for interpreting ops inheriting from LinalgStructuredOperation.

Source code in xdsl/interpreters/linalg.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def run_linalg_structured_op(
    interpreter: Interpreter,
    op: linalg.LinalgStructuredOperation,
    args: tuple[ShapedArray[float] | float, ...],
):
    """
    Helper function for interpreting ops inheriting from
    [`LinalgStructuredOperation`][xdsl.dialects.linalg.LinalgStructuredOperation].
    """
    body = op.body
    inputs_count = len(op.inputs)
    input_args = args[:inputs_count]
    output_args = args[inputs_count:]
    results = op.results
    indexing_maps = tuple(attr.data for attr in op.get_indexing_maps())
    loop_ranges = op.get_static_loop_ranges()

    if any(not isinstance(arg, ShapedArray) for arg in output_args):
        raise NotImplementedError("Only shaped out results are implemented")
    output_args = cast(tuple[ShapedArray[float], ...], output_args)
    if results:
        # If there are results, they must be tensors, initialised with the
        # `output_args`. If not, the results are stored in output_args directly.
        outputs = tuple(arg.copy() for arg in output_args)
    else:
        outputs = output_args

    loop_shaped_args = input_args + outputs

    output_indexing_maps = indexing_maps[inputs_count:]

    for indices in product(*(range(loop_range) for loop_range in loop_ranges)):
        loop_scalar_args = tuple(
            (
                i.load(indexing_map.eval(indices, ()))
                if isinstance(i, ShapedArray)
                else i
            )
            for i, indexing_map in zip(loop_shaped_args, indexing_maps, strict=True)
        )
        loop_results = interpreter.run_ssacfg_region(body, loop_scalar_args, "for_loop")
        for res, indexing_map in zip(loop_results, output_indexing_maps, strict=True):
            result_indices = indexing_map.eval(indices, ())
            outputs[0].store(result_indices, res)

    return outputs if results else ()