[Mlir-commits] [mlir] [mlir][linalg] Add 'compute_element_type' to linalg.elementwise op. (PR #190566)

[Javed Absar]

================
@@ -163,3 +163,28 @@ func.func @ternary(%A : tensor<32x16xi1>, %B: tensor<8x16x32xf32>, %C : tensor<8
       outs(%D: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>
   return %r : tensor<8x16x32xf32>
 }
+// -----
+// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//
+// CHECK: @unary_upcast_exp(%[[A:.+]]: tensor<8x16x32xf16>, %[[B:.+]]: tensor<8x16x32xf32>)
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[IDENTITY]], #[[IDENTITY]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]
+//
+// CHECK-SAME: ins(%[[A]]
+// CHECK-SAME: outs(%[[B]]
+//
+// CHECK: ^{{.*}}(%[[IN:.+]]: f16, %[[OUT:.+]]: f32)
+// CHECK:   %[[EXT:.+]] = arith.extf %[[IN]] : f16 to f32
+// CHECK:   %[[EXP:.+]] = math.exp %[[EXT]] : f32
+// CHECK:   linalg.yield %[[EXP]] : f32
+//
+func.func @unary_upcast_exp(%A : tensor<8x16x32xf16>, %B : tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {
+  %r = linalg.elementwise
+      kind = #linalg.elementwise_kind<exp>
+      compute_element_type = f32
----------------
javedabsar1 wrote:

OK currently (without this PR), if input is f16 and output if f32, it will result in error i.e.
```
$ cat t.mlir
func.func @unary_identity_exp(%A : tensor<8x16x32xf16>, %B: tensor<8x16x32xf32>) ->  tensor<8x16x32xf32> {
  %r = linalg.elementwise
      kind=#linalg.elementwise_kind<exp>
      ins(%A : tensor<8x16x32xf16>)
      outs(%B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>
  return %r : tensor<8x16x32xf32>
}
$ mlir-opt t.mlir
 ...
 error: 'linalg.yield' op type of yield operand 1 ('f16') doesn't match the element type of the enclosing linalg.generic op ('f32')
 ```
 Now, do you want this PR to address just this i.e. the above `linalg.elementwise` example becomes valid, AND therefore the when it is lowered to  `linalg.generic` the body of the linalg.generic will do the extf.
 
 ---
 
 Next is the question why a broader `compute_element_type` is needed. It is for scenarios where input e.g. linalg.exp is f16, but the exponentiation needs to be done on higher precision f32, and the result written out in lower precision f16.  
  
 

https://github.com/llvm/llvm-project/pull/190566