1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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
64
65
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
|
import ctypes
import math
import mlir_finch.execution_engine
import mlir_finch.passmanager
from mlir_finch import ir
from mlir_finch.dialects import arith, complex, func, linalg, sparse_tensor, tensor
import numpy as np
from ._array import Array
from ._common import as_shape, fn_cache
from ._core import CWD, DEBUG, OPT_LEVEL, SHARED_LIBS, ctx, pm
from ._dtypes import DType, IeeeComplexFloatingDType, IeeeRealFloatingDType, IntegerDType
from .formats import ConcreteFormat, _determine_format
@fn_cache
def get_add_module(
a_tensor_type: ir.RankedTensorType,
b_tensor_type: ir.RankedTensorType,
out_tensor_type: ir.RankedTensorType,
dtype: DType,
) -> ir.Module:
with ir.Location.unknown(ctx):
module = ir.Module.create()
if isinstance(dtype, IeeeRealFloatingDType):
arith_op = arith.AddFOp
elif isinstance(dtype, IeeeComplexFloatingDType):
arith_op = complex.AddOp
elif isinstance(dtype, IntegerDType):
arith_op = arith.AddIOp
else:
raise RuntimeError(f"Can not add {dtype=}.")
dtype = dtype._get_mlir_type()
max_rank = out_tensor_type.rank
with ir.InsertionPoint(module.body):
@func.FuncOp.from_py_func(a_tensor_type, b_tensor_type)
def add(a, b):
out = tensor.empty(out_tensor_type.shape, dtype, encoding=out_tensor_type.encoding)
generic_op = linalg.GenericOp(
[out_tensor_type],
[a, b],
[out],
ir.ArrayAttr.get(
[
ir.AffineMapAttr.get(ir.AffineMap.get_minor_identity(max_rank, t.rank))
for t in (a_tensor_type, b_tensor_type, out_tensor_type)
]
),
ir.ArrayAttr.get([ir.Attribute.parse("#linalg.iterator_type<parallel>")] * max_rank),
)
block = generic_op.regions[0].blocks.append(dtype, dtype, dtype)
with ir.InsertionPoint(block):
a, b, o = block.arguments
res = sparse_tensor.BinaryOp(dtype, a, b)
overlap = res.regions[0].blocks.append(dtype, dtype)
with ir.InsertionPoint(overlap):
arg0, arg1 = overlap.arguments
overlap_res = arith_op(arg0, arg1)
sparse_tensor.YieldOp([overlap_res])
left_region = res.regions[1].blocks.append(dtype)
with ir.InsertionPoint(left_region):
(arg0,) = left_region.arguments
sparse_tensor.YieldOp([arg0])
right_region = res.regions[2].blocks.append(dtype)
with ir.InsertionPoint(right_region):
(arg0,) = right_region.arguments
sparse_tensor.YieldOp([arg0])
linalg.YieldOp([res])
return generic_op.result
add.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
if DEBUG:
(CWD / "add_module.mlir").write_text(str(module))
pm.run(module.operation)
if DEBUG:
(CWD / "add_module_opt.mlir").write_text(str(module))
return mlir_finch.execution_engine.ExecutionEngine(module, opt_level=OPT_LEVEL, shared_libs=SHARED_LIBS)
@fn_cache
def get_reshape_module(
a_tensor_type: ir.RankedTensorType,
shape_tensor_type: ir.RankedTensorType,
out_tensor_type: ir.RankedTensorType,
) -> ir.Module:
with ir.Location.unknown(ctx):
module = ir.Module.create()
with ir.InsertionPoint(module.body):
@func.FuncOp.from_py_func(a_tensor_type, shape_tensor_type)
def reshape(a, shape):
return tensor.reshape(out_tensor_type, a, shape)
reshape.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
if DEBUG:
(CWD / "reshape_module.mlir").write_text(str(module))
pm.run(module.operation)
if DEBUG:
(CWD / "reshape_module_opt.mlir").write_text(str(module))
return mlir_finch.execution_engine.ExecutionEngine(module, opt_level=OPT_LEVEL, shared_libs=SHARED_LIBS)
@fn_cache
def get_broadcast_to_module(
in_tensor_type: ir.RankedTensorType,
out_tensor_type: ir.RankedTensorType,
dimensions: tuple[int, ...],
) -> ir.Module:
with ir.Location.unknown(ctx):
module = ir.Module.create()
with ir.InsertionPoint(module.body):
@func.FuncOp.from_py_func(in_tensor_type)
def broadcast_to(in_tensor):
out = tensor.empty(
out_tensor_type.shape, out_tensor_type.element_type, encoding=out_tensor_type.encoding
)
return linalg.broadcast(in_tensor, outs=[out], dimensions=dimensions)
broadcast_to.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
if DEBUG:
(CWD / "broadcast_to_module.mlir").write_text(str(module))
pm.run(module.operation)
if DEBUG:
(CWD / "broadcast_to_module_opt.mlir").write_text(str(module))
return mlir_finch.execution_engine.ExecutionEngine(module, opt_level=OPT_LEVEL, shared_libs=SHARED_LIBS)
@fn_cache
def get_convert_module(
in_tensor_type: ir.RankedTensorType,
out_tensor_type: ir.RankedTensorType,
):
with ir.Location.unknown(ctx):
module = ir.Module.create()
with ir.InsertionPoint(module.body):
@func.FuncOp.from_py_func(in_tensor_type)
def convert(in_tensor):
return sparse_tensor.convert(out_tensor_type, in_tensor)
convert.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
if DEBUG:
(CWD / "convert_module.mlir").write_text(str(module))
pm.run(module.operation)
if DEBUG:
(CWD / "convert_module.mlir").write_text(str(module))
return mlir_finch.execution_engine.ExecutionEngine(module, opt_level=OPT_LEVEL, shared_libs=SHARED_LIBS)
def add(x1: Array, x2: Array, /) -> Array:
# TODO: Determine output format via autoscheduler
ret_storage_format = _determine_format(x1.format, x2.format, dtype=x1.dtype, union=True)
ret_storage = ret_storage_format._get_ctypes_type(owns_memory=True)()
out_tensor_type = ret_storage_format._get_mlir_type(shape=np.broadcast_shapes(x1.shape, x2.shape))
add_module = get_add_module(
x1._get_mlir_type(),
x2._get_mlir_type(),
out_tensor_type=out_tensor_type,
dtype=x1.dtype,
)
add_module.invoke(
"add",
ctypes.pointer(ctypes.pointer(ret_storage)),
*x1._to_module_arg(),
*x2._to_module_arg(),
)
return Array(storage=ret_storage, shape=tuple(out_tensor_type.shape))
def asformat(x: Array, /, format: ConcreteFormat) -> Array:
if format.rank != x.ndim:
raise ValueError(f"`format.rank != `self.ndim`, {format.rank=}, {x.ndim=}")
if format == x.format:
return x
out_tensor_type = format._get_mlir_type(shape=x.shape)
ret_storage = format._get_ctypes_type(owns_memory=True)()
convert_module = get_convert_module(
x._get_mlir_type(),
out_tensor_type,
)
convert_module.invoke(
"convert",
ctypes.pointer(ctypes.pointer(ret_storage)),
*x._to_module_arg(),
)
return Array(storage=ret_storage, shape=x.shape)
def reshape(x: Array, /, shape: tuple[int, ...]) -> Array:
from ._conversions import _from_numpy
shape = as_shape(shape)
if math.prod(x.shape) != math.prod(shape):
raise ValueError(f"`math.prod(x.shape) != math.prod(shape)`, {x.shape=}, {shape=}")
ret_storage_format = _determine_format(x.format, dtype=x.dtype, union=len(shape) > x.ndim, out_ndim=len(shape))
shape_array = _from_numpy(np.asarray(shape, dtype=np.uint64))
out_tensor_type = ret_storage_format._get_mlir_type(shape=shape)
ret_storage = ret_storage_format._get_ctypes_type(owns_memory=True)()
reshape_module = get_reshape_module(x._get_mlir_type(), shape_array._get_mlir_type(), out_tensor_type)
reshape_module.invoke(
"reshape",
ctypes.pointer(ctypes.pointer(ret_storage)),
*x._to_module_arg(),
*shape_array._to_module_arg(),
)
return Array(storage=ret_storage, shape=shape)
|
