251 lines · python
1# RUN: %PYTHON %s2 3from mlir.dialects import arith, func, linalg4from mlir.dialects.linalg.opdsl.lang import *5from mlir.ir import *6 7 8def run(f):9 print("\nTEST:", f.__name__)10 f()11 return f12 13 14@run15def test_infer_contraction_dimensions_from_ops():16 with Context(), Location.unknown():17 module = Module.create()18 f32 = F32Type.get()19 with InsertionPoint(module.body):20 # === Static shapes ===21 m, n, k = 4, 4, 422 a_type = RankedTensorType.get((m, k), f32)23 b_type = RankedTensorType.get((k, n), f32)24 c_type = RankedTensorType.get((m, n), f32)25 26 @func.FuncOp.from_py_func(a_type, b_type, c_type)27 def contraction_fn(a, b, c):28 zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)29 filled = linalg.fill(zero, outs=[c])30 fill_op = filled.owner31 32 assert not linalg.isa_contraction_op(zero.operation)33 assert not linalg.isa_contraction_op(fill_op)34 assert linalg.infer_contraction_dimensions(fill_op) is None35 36 dim_m = AffineDimExpr.get(0)37 dim_n = AffineDimExpr.get(1)38 dim_k = AffineDimExpr.get(2)39 40 a_map = AffineMap.get(3, 0, [dim_m, dim_k])41 b_map = AffineMap.get(3, 0, [dim_k, dim_n])42 c_map = AffineMap.get(3, 0, [dim_m, dim_n])43 result = linalg.contract(44 a,45 b,46 outs=(filled,),47 indexing_maps=[a_map, b_map, c_map],48 )49 contraction_op = result.owner50 51 assert linalg.isa_contraction_op(contraction_op)52 dims = linalg.infer_contraction_dimensions(contraction_op)53 assert dims is not None54 55 # Expect m=[0], n=[1], k=[2] as per standard matmul.56 assert list(dims.m) == [0], f"Expected m=[0], got {list(dims.m)}"57 assert list(dims.n) == [1], f"Expected n=[1], got {list(dims.n)}"58 assert list(dims.k) == [2], f"Expected k=[2], got {list(dims.k)}"59 assert (60 list(dims.batch) == []61 ), f"Expected batch=[], got {list(dims.batch)}"62 63 # === Dynamic shape case ===64 dyn = ShapedType.get_dynamic_size()65 a_dyn_type = RankedTensorType.get((4, dyn), f32)66 b_dyn_type = RankedTensorType.get((dyn, 4), f32)67 c_type = RankedTensorType.get((4, 4), f32)68 69 @func.FuncOp.from_py_func(a_dyn_type, b_dyn_type, c_type)70 def dynamic_contraction_fn(a, b, c):71 zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)72 filled = linalg.fill(zero, outs=[c])73 dim_m = AffineDimExpr.get(0)74 dim_n = AffineDimExpr.get(1)75 dim_k = AffineDimExpr.get(2)76 77 a_map = AffineMap.get(3, 0, [dim_m, dim_k])78 b_map = AffineMap.get(3, 0, [dim_k, dim_n])79 c_map = AffineMap.get(3, 0, [dim_m, dim_n])80 81 result = linalg.contract(82 a,83 b,84 outs=(filled,),85 indexing_maps=[a_map, b_map, c_map],86 )87 contraction_op = result.owner88 89 assert linalg.isa_contraction_op(contraction_op)90 dims = linalg.infer_contraction_dimensions(contraction_op)91 assert dims is not None92 assert list(dims.m) == [0], f"Expected m=[0], got {list(dims.m)}"93 assert list(dims.n) == [1], f"Expected n=[1], got {list(dims.n)}"94 assert list(dims.k) == [2], f"Expected k=[2], got {list(dims.k)}"95 assert (96 list(dims.batch) == []97 ), f"Expected batch=[], got {list(dims.batch)}"98 99 100@run101def test_infer_convolution_dimensions_from_ops():102 with Context(), Location.unknown():103 module = Module.create()104 f32 = F32Type.get()105 106 with InsertionPoint(module.body):107 # === Static shapes ===108 batch, h, w, c_in, kh, kw, c_out = 1, 8, 8, 4, 3, 3, 16109 input_type = RankedTensorType.get((batch, h, w, c_in), f32)110 filter_type = RankedTensorType.get((kh, kw, c_in, c_out), f32)111 output_type = RankedTensorType.get(112 (batch, h - kh + 1, w - kw + 1, c_out), f32113 )114 115 @func.FuncOp.from_py_func(input_type, filter_type, output_type)116 def conv_fn(input, filter, output):117 zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)118 filled = linalg.fill(zero, outs=[output])119 fill_op = filled.owner120 121 assert not linalg.isa_convolution_op(fill_op)122 assert linalg.infer_convolution_dimensions(fill_op) is None123 124 result = linalg.conv_2d_nhwc_hwcf(input, filter, outs=[filled])125 conv_op = result.owner126 127 assert linalg.isa_convolution_op(conv_op)128 dims = linalg.infer_convolution_dimensions(conv_op)129 assert dims is not None130 assert list(dims.batch) == [0]131 assert list(dims.output_image) == [1, 2]132 assert list(dims.output_channel) == [3]133 assert list(dims.filter_loop) == [4, 5]134 assert list(dims.input_channel) == [6]135 assert list(dims.depth) == []136 assert list(dims.strides) == [1, 1]137 assert list(dims.dilations) == [1, 1]138 139 # === Dynamic shapes ===140 dyn = ShapedType.get_dynamic_size()141 dyn_input_type = RankedTensorType.get((batch, dyn, dyn, c_in), f32)142 dyn_output_type = RankedTensorType.get((batch, dyn, dyn, c_out), f32)143 144 @func.FuncOp.from_py_func(dyn_input_type, filter_type, dyn_output_type)145 def dyn_conv_fn(input, filter, output):146 zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)147 filled = linalg.fill(zero, outs=[output])148 result = linalg.conv_2d_nhwc_hwcf(input, filter, outs=[filled])149 conv_op = result.owner150 151 assert linalg.isa_convolution_op(conv_op)152 dims = linalg.infer_convolution_dimensions(conv_op)153 assert dims is not None154 assert list(dims.batch) == [0]155 assert list(dims.output_image) == [1, 2]156 assert list(dims.output_channel) == [3]157 assert list(dims.filter_loop) == [4, 5]158 assert list(dims.input_channel) == [6]159 assert list(dims.depth) == []160 assert list(dims.strides) == [1, 1]161 assert list(dims.dilations) == [1, 1]162 163 164@run165def test_get_indexing_maps_attr():166 with Context(), Location.unknown():167 module = Module.create()168 f32 = F32Type.get()169 with InsertionPoint(module.body):170 a_type = RankedTensorType.get((4, 8), f32)171 b_type = RankedTensorType.get((8, 16), f32)172 c_type = RankedTensorType.get((4, 16), f32)173 174 dim_m = AffineDimExpr.get(0)175 dim_n = AffineDimExpr.get(1)176 dim_k = AffineDimExpr.get(2)177 178 a_map = AffineMap.get(3, 0, [dim_m, dim_k])179 b_map = AffineMap.get(3, 0, [dim_k, dim_n])180 c_map = AffineMap.get(3, 0, [dim_m, dim_n])181 182 @func.FuncOp.from_py_func(a_type, b_type, c_type)183 def matmul_func(a, b, c):184 zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.0), result=f32)185 assert not linalg.get_indexing_maps(186 zero.operation187 ), "Expected no indexing_maps on non-linalg op"188 189 init = linalg.fill(zero, outs=[c])190 fill_op = init.owner191 fill_maps = linalg.get_indexing_maps(fill_op)192 assert fill_maps is not None193 assert len(fill_maps) == 2194 195 # The fill op should have maps like (d0, d1) -> () and (d0, d1).196 fill_input_map = fill_maps[0].value197 fill_output_map = fill_maps[1].value198 assert fill_input_map == AffineMap.get(2, 0, [])199 assert fill_output_map == AffineMap.get(2, 0, [dim_m, dim_n])200 201 result = linalg.matmul(a, b, outs=(init,))202 matmul_op = result.owner203 matmul_maps = linalg.get_indexing_maps(matmul_op)204 assert matmul_maps is not None205 assert len(matmul_maps) == 3206 207 maps = [map_attr.value for map_attr in matmul_maps]208 assert maps[0] == a_map209 assert maps[1] == b_map210 assert maps[2] == c_map211 212 213@run214def test_infer_contraction_dimensions_from_maps():215 with Context(), Location.unknown():216 module = Module.create()217 with InsertionPoint(module.body):218 # === Test valid contraction (matmul) ===219 dim_m = AffineDimExpr.get(0)220 dim_n = AffineDimExpr.get(1)221 dim_k = AffineDimExpr.get(2)222 a_map = AffineMap.get(3, 0, [dim_m, dim_k])223 b_map = AffineMap.get(3, 0, [dim_k, dim_n])224 c_map = AffineMap.get(3, 0, [dim_m, dim_n])225 226 dims = linalg.infer_contraction_dimensions_from_maps([a_map, b_map, c_map])227 assert dims is not None228 229 # Expect m=[0], n=[1], k=[2] as per standard matmul.230 assert list(dims.m) == [0], f"Expected m=[0], got {list(dims.m)}"231 assert list(dims.n) == [1], f"Expected n=[1], got {list(dims.n)}"232 assert list(dims.k) == [2], f"Expected k=[2], got {list(dims.k)}"233 assert list(dims.batch) == [], f"Expected batch=[], got {list(dims.batch)}"234 235 # === Test invalid input (wrong number of maps) ===236 invalid_dims = linalg.infer_contraction_dimensions_from_maps([a_map, b_map])237 assert invalid_dims is None238 239 # === Test element-wise operation ===240 dim_i = AffineDimExpr.get(0)241 dim_j = AffineDimExpr.get(1)242 elementwise_map = AffineMap.get(2, 0, [dim_i, dim_j])243 elementwise_dims = linalg.infer_contraction_dimensions_from_maps(244 [elementwise_map, elementwise_map, elementwise_map]245 )246 assert elementwise_dims is not None247 assert len(elementwise_dims.m) == 0248 assert len(elementwise_dims.n) == 0249 assert len(elementwise_dims.k) == 0250 assert list(elementwise_dims.batch) == [0, 1]251