1# RUN: %PYTHON %s | FileCheck %s 2 3from mlir.ir import * 4from mlir.dialects import builtin 5from mlir.dialects import func 6from mlir.dialects import linalg 7 8from mlir.dialects.linalg.opdsl.lang import * 9 10# This tests miscellaneous features of the emitter that are not tested by the 11# fill, matmul, convolution, or pooling tests. The features include: 12# - constant defined in the body 13# - fix/predefined types 14# - some math/arith functions, including abs, ceil, exp, floor, log, and negf 15# - custom op names. 16 17 18@linalg_structured_op 19def test_const(O=TensorDef(F32, S.M, S.N, output=True)): 20 O[D.m, D.n] = TypeFn.cast_unsigned(F32, const(42)) + TypeFn.cast_unsigned( 21 F32, const(2.3283064e-10)) 22 23 24@linalg_structured_op 25def test_index(O=TensorDef(I32, S.M, S.N, output=True)): 26 O[D.m, D.n] = TypeFn.cast_signed(I32, index(D.m)) + TypeFn.cast_signed( 27 I32, index(D.n)) 28 29 30@linalg_structured_op 31def elemwise_unary_poly( 32 I=TensorDef(T), 33 O=TensorDef(U, output=True), 34 fun=UnaryFnAttrDef(default=UnaryFn.exp), 35 cast=TypeFnAttrDef(default=TypeFn.cast_signed)): 36 O[None] = fun(cast(U, I[None])) 37 38 39@linalg_structured_op(op_name="custom_op_name") 40def non_default_op_name(I=TensorDef(T, S.N), O=TensorDef(T, S.N, output=True)): 41 O[D.n] = I[D.n] 42 43 44with Context() as ctx, Location.unknown(): 45 module = Module.create() 46 f32 = F32Type.get() 47 c32 = ComplexType.get(f32) 48 i32 = IntegerType.get_signless(32) 49 with InsertionPoint(module.body): 50 51 # CHECK-LABEL: @test_f32_const 52 # CHECK-DAG: %[[CST0:.+]] = arith.constant 42 : i64 53 # CHECK-DAG: %[[CST0_CAST:.+]] = arith.uitofp %[[CST0]] : i64 to f32 54 # CHECK-DAG: %[[CST1:.+]] = arith.constant 2.3283063999999999E-10 : f64 55 # CHECK-DAG: %[[CST1_CAST:.+]] = arith.truncf %[[CST1]] : f64 to f32 56 # CHECK-DAG: %[[SUM:.+]] = arith.addf %[[CST0_CAST]], %[[CST1_CAST]] : f32 57 # CHECK-NEXT: linalg.yield %[[SUM]] : f32 58 @func.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32)) 59 def test_f32_const(init_result): 60 return test_const(outs=[init_result]) 61 62 # CHECK-LABEL: @test_i32_index 63 # CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index 64 # CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index 65 # CHECK-DAG: %[[IDX0_CAST:.+]] = arith.index_cast %[[IDX0]] : index to i32 66 # CHECK-DAG: %[[IDX1_CAST:.+]] = arith.index_cast %[[IDX1]] : index to i32 67 # CHECK-DAG: %[[SUM:.+]] = arith.addi %[[IDX0_CAST]], %[[IDX1_CAST]] : i32 68 # CHECK-NEXT: linalg.yield %[[SUM]] : i32 69 @func.FuncOp.from_py_func(RankedTensorType.get((4, 16), i32)) 70 def test_i32_index(init_result): 71 return test_index(outs=[init_result]) 72 73 # CHECK-LABEL: @test_f32_elemwise_exp 74 # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) 75 # CHECK-NEXT: %[[EXP:.+]] = math.exp %[[IN]] : f32 76 # CHECK-NEXT: linalg.yield %[[EXP]] : f32 77 # CHECK-NEXT: -> tensor<4x16xf32> 78 @func.FuncOp.from_py_func( 79 RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) 80 def test_f32_elemwise_exp(input, init_result): 81 return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.exp) 82 83 # CHECK-LABEL: @test_f32_elemwise_log 84 # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) 85 # CHECK-NEXT: %[[LOG:.+]] = math.log %[[IN]] : f32 86 # CHECK-NEXT: linalg.yield %[[LOG]] : f32 87 # CHECK-NEXT: -> tensor<4x16xf32> 88 @func.FuncOp.from_py_func( 89 RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) 90 def test_f32_elemwise_log(input, init_result): 91 return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.log) 92 93 # CHECK-LABEL: @test_f32_elemwise_abs 94 # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) 95 # CHECK-NEXT: %[[EXP:.+]] = math.abs %[[IN]] : f32 96 # CHECK-NEXT: linalg.yield %[[EXP]] : f32 97 # CHECK-NEXT: -> tensor<4x16xf32> 98 @func.FuncOp.from_py_func( 99 RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) 100 def test_f32_elemwise_abs(input, init_result): 101 return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.abs) 102 103 # CHECK-LABEL: @test_f32_elemwise_ceil 104 # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) 105 # CHECK-NEXT: %[[EXP:.+]] = math.ceil %[[IN]] : f32 106 # CHECK-NEXT: linalg.yield %[[EXP]] : f32 107 # CHECK-NEXT: -> tensor<4x16xf32> 108 @func.FuncOp.from_py_func( 109 RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) 110 def test_f32_elemwise_ceil(input, init_result): 111 return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.ceil) 112 113 # CHECK-LABEL: @test_f32_elemwise_floor 114 # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) 115 # CHECK-NEXT: %[[EXP:.+]] = math.floor %[[IN]] : f32 116 # CHECK-NEXT: linalg.yield %[[EXP]] : f32 117 # CHECK-NEXT: -> tensor<4x16xf32> 118 @func.FuncOp.from_py_func( 119 RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) 120 def test_f32_elemwise_floor(input, init_result): 121 return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.floor) 122 123 # CHECK-LABEL: @test_f32_elemwise_neg 124 # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) 125 # CHECK-NEXT: %[[EXP:.+]] = arith.negf %[[IN]] : f32 126 # CHECK-NEXT: linalg.yield %[[EXP]] : f32 127 # CHECK-NEXT: -> tensor<4x16xf32> 128 @func.FuncOp.from_py_func( 129 RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) 130 def test_f32_elemwise_neg(input, init_result): 131 return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.negf) 132 133 # CHECK-LABEL: @test_c32_elemwise_neg 134 # CHECK: ^{{.*}}(%[[IN:.+]]: complex<f32>, %[[OUT:.+]]: complex<f32>) 135 # CHECK-NEXT: %[[EXP:.+]] = complex.neg %[[IN]] : complex<f32> 136 # CHECK-NEXT: linalg.yield %[[EXP]] : complex<f32> 137 # CHECK-NEXT: -> tensor<4x16xcomplex<f32>> 138 @func.FuncOp.from_py_func( 139 RankedTensorType.get((4, 16), c32), RankedTensorType.get((4, 16), c32)) 140 def test_c32_elemwise_neg(input, init_result): 141 return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.negf) 142 143 # Just check that we don't assert out on name mismatch. 144 # CHECK-LABEL: @test_non_default_op_name 145 @func.FuncOp.from_py_func( 146 RankedTensorType.get((42,), f32), RankedTensorType.get((42,), f32)) 147 def test_non_default_op_name(input, init_result): 148 return non_default_op_name(input, outs=[init_result]) 149 150 151print(module) 152