[mlir][linalg] Add a few unary operations.

[bixia]

Summary

Add operations abs, ceil, floor, and neg to the C++ API and Python API.

Add test cases.

[aartbik]

LGTM, but please wait for Nicolas, since this really is his playground ;-)

[gysit]

Thanks for adding!

There are also a couple of places on the C++ side that need to be updated.

1. there is an enum that defines all the Unary function types:

https://github.com/llvm/llvm-project/blob/9fac110bf70896a5411cf57b3ff34fab5798ab78/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td#L62
the names have to match the function names used on the Python side!

2. there is a helper to build the functions:

https://github.com/llvm/llvm-project/blob/9fac110bf70896a5411cf57b3ff34fab5798ab78/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L138
it is the counterpart to emitter.py if the operation is generated from the yaml file.

3. is an example that tests the C++ code generation path

https://github.com/llvm/llvm-project/blob/eb6a3c0c0c71ab44141e71112ecd0a2ae2848037/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir#L290

I think it would be great to update the C++ side in the same revision. Doing it in a follow up commit is also fine for me.

[bixia]

Added the operations to C++.

[bixia]

Thanks for the review and pointers for the C++ changes! I added the C++ change and modified the PR description. PTAL.

[gysit]

perfect, thanks!

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

let extraClassDeclaration = [{
     private:
       llvm::StringMap<RegionBuilderFunType> namedStructuredOpRegionBuilders;
   }];
 }
 
 // Define the function attribute enums matching the OpDSL functions.
 def UnaryFn : I32EnumAttr<"UnaryFn", "", [
   I32EnumAttrCase<"exp", 0>,
-  I32EnumAttrCase<"log", 1>
+  I32EnumAttrCase<"log", 1>,
+  I32EnumAttrCase<"abs", 2>,
+  I32EnumAttrCase<"ceil", 3>,
+  I32EnumAttrCase<"floor", 4>,
+  I32EnumAttrCase<"negf", 5>
 ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::linalg";
 }
 def BinaryFn : I32EnumAttr<"BinaryFn", "", [
   I32EnumAttrCase<"add", 0>,
   I32EnumAttrCase<"sub", 1>,
   I32EnumAttrCase<"mul", 2>,

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Value buildUnaryFn(UnaryFn unaryFn, Value arg) {
     if (!isFloatingPoint(arg))
       llvm_unreachable("unsupported non numeric type");
     OpBuilder builder = getBuilder();
     switch (unaryFn) {
     case UnaryFn::exp:
       return builder.create<math::ExpOp>(arg.getLoc(), arg);
     case UnaryFn::log:
       return builder.create<math::LogOp>(arg.getLoc(), arg);
+    case UnaryFn::abs:
+      return builder.create<math::AbsOp>(arg.getLoc(), arg);
+    case UnaryFn::ceil:
+      return builder.create<math::CeilOp>(arg.getLoc(), arg);
+    case UnaryFn::floor:
+      return builder.create<math::FloorOp>(arg.getLoc(), arg);
+    case UnaryFn::negf:
+      return builder.create<arith::NegFOp>(arg.getLoc(), arg);
     }
     llvm_unreachable("unsupported unary function");
   }
 
   // Build the binary functions defined by OpDSL.
   Value buildBinaryFn(BinaryFn binaryFn, Value arg0, Value arg1) {
     bool allFloatingPoint = isFloatingPoint(arg0) && isFloatingPoint(arg1);
     bool allInteger = isInteger(arg0) && isInteger(arg1);

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

class UnaryFnType:
   def __repr__(self):
     return f"{self.fn_name}"
 
 
 class UnaryFn:
   """Unary function namespace."""
   exp = UnaryFnType("exp")
   log = UnaryFnType("log")
+  abs = UnaryFnType("abs")
+  ceil = UnaryFnType("ceil")
+  floor = UnaryFnType("floor")
+  negf = UnaryFnType("negf")
 
 
 class BinaryFnType:
   """Binary function.
 
   A binary function takes two tensor expressions and returns the
   function evaluation result.
   """

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

if _is_floating_point_type(x.type):
       return math.ExpOp(x).result
     raise NotImplementedError("Unsupported 'exp' operand: {x}")
 
   def _unary_log(self, x: Value) -> Value:
     if _is_floating_point_type(x.type):
       return math.LogOp(x).result
     raise NotImplementedError("Unsupported 'log' operand: {x}")
 
+  def _unary_abs(self, x: Value) -> Value:
+    if _is_floating_point_type(x.type):
+      return math.AbsOp(x).result
+    raise NotImplementedError("Unsupported 'abs' operand: {x}")
+
+  def _unary_ceil(self, x: Value) -> Value:
+    if _is_floating_point_type(x.type):
+      return math.CeilOp(x).result
+    raise NotImplementedError("Unsupported 'ceil' operand: {x}")
+
+  def _unary_floor(self, x: Value) -> Value:
+    if _is_floating_point_type(x.type):
+      return math.FloorOp(x).result
+    raise NotImplementedError("Unsupported 'floor' operand: {x}")
+
+  def _unary_negf(self, x: Value) -> Value:
+    if _is_floating_point_type(x.type):
+      return arith.NegFOp(x).result
+    raise NotImplementedError("Unsupported 'negf' operand: {x}")
+
   def _binary_add(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.AddFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.AddIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'add' operands: {lhs}, {rhs}")
 
   def _binary_sub(self, lhs: Value, rhs: Value) -> Value:

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

func @generalize_elemwise_log(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
   return %0: tensor<4x8xf32>
 }
 
 // CHECK-LABEL: @generalize_elemwise_log
 // CHECK:        = math.log
 
 // -----
 
+// Verifies the fun attribute controls the unary function used.
+func @generalize_elemwise_abs(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
+  %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<abs>}
+                              ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
+  return %0: tensor<4x8xf32>
+}
+
+// CHECK-LABEL: @generalize_elemwise_abs
+// CHECK:        = math.abs
+
+// -----
+
+// Verifies the fun attribute controls the unary function used.
+func @generalize_elemwise_ceil(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
+  %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<ceil>}
+                              ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
+  return %0: tensor<4x8xf32>
+}
+
+// CHECK-LABEL: @generalize_elemwise_ceil
+// CHECK:        = math.ceil
+
+// -----
+
+// Verifies the fun attribute controls the unary function used.
+func @generalize_elemwise_floor(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
+  %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<floor>}
+                              ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
+  return %0: tensor<4x8xf32>
+}
+
+// CHECK-LABEL: @generalize_elemwise_floor
+// CHECK:        = math.floor
+
+// -----
+
+// Verifies the fun attribute controls the unary function used.
+func @generalize_elemwise_negf(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
+  %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<negf>}
+                              ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
+  return %0: tensor<4x8xf32>
+}
+
+// CHECK-LABEL: @generalize_elemwise_negf
+// CHECK:        = arith.negf
+
+// -----
+
 // Verifies the default value of the fun attribute is an add op.
 func @generalize_elemwise_add(%lhs : tensor<4x8xf32>, %rhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
   %0 = linalg.elemwise_binary ins(%lhs, %rhs: tensor<4x8xf32>, tensor<4x8xf32>)
                               outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
   return %0: tensor<4x8xf32>
 }
 
 // CHECK-LABEL: @generalize_elemwise_add

mlir/test/python/dialects/linalg/opdsl/emit_misc.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import func
 from mlir.dialects import linalg
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 # This tests miscellaneous features of the emitter that are not tested by the
 # fill, matmul, convolution, or pooling tests. The features include:
 # - constant defined in the body
 # - fix/predefined types
-# - exponential functions
+# - some math/arith functions, including abs, ceil, exp, floor, log, and negf
 # - custom op names.
 
 
 @linalg_structured_op
 def test_const(O=TensorDef(F32, S.M, S.N, output=True)):
   O[D.m, D.n] = TypeFn.cast_unsigned(F32, const(42)) + TypeFn.cast_unsigned(
       F32, const(2.3283064e-10))
 
with InsertionPoint(module.body):
     # CHECK-NEXT:   %[[LOG:.+]] = math.log %[[IN]] : f32
     # CHECK-NEXT:   linalg.yield %[[LOG]] : f32
     # CHECK-NEXT: -> tensor<4x16xf32>
     @builtin.FuncOp.from_py_func(
         RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
     def test_f32_elemwise_log(input, init_result):
       return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.log)
 
+    # CHECK-LABEL: @test_f32_elemwise_abs
+    # CHECK:      ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32)
+    # CHECK-NEXT:   %[[EXP:.+]] = math.abs %[[IN]] : f32
+    # CHECK-NEXT:   linalg.yield %[[EXP]] : f32
+    # CHECK-NEXT: -> tensor<4x16xf32>
+    @builtin.FuncOp.from_py_func(
+        RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
+    def test_f32_elemwise_abs(input, init_result):
+      return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.abs)
+
+    # CHECK-LABEL: @test_f32_elemwise_ceil
+    # CHECK:      ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32)
+    # CHECK-NEXT:   %[[EXP:.+]] = math.ceil %[[IN]] : f32
+    # CHECK-NEXT:   linalg.yield %[[EXP]] : f32
+    # CHECK-NEXT: -> tensor<4x16xf32>
+    @builtin.FuncOp.from_py_func(
+        RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
+    def test_f32_elemwise_ceil(input, init_result):
+      return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.ceil)
+
+    # CHECK-LABEL: @test_f32_elemwise_floor
+    # CHECK:      ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32)
+    # CHECK-NEXT:   %[[EXP:.+]] = math.floor %[[IN]] : f32
+    # CHECK-NEXT:   linalg.yield %[[EXP]] : f32
+    # CHECK-NEXT: -> tensor<4x16xf32>
+    @builtin.FuncOp.from_py_func(
+        RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
+    def test_f32_elemwise_floor(input, init_result):
+      return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.floor)
+
+    # CHECK-LABEL: @test_f32_elemwise_neg
+    # CHECK:      ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32)
+    # CHECK-NEXT:   %[[EXP:.+]] = arith.negf %[[IN]] : f32
+    # CHECK-NEXT:   linalg.yield %[[EXP]] : f32
+    # CHECK-NEXT: -> tensor<4x16xf32>
+    @builtin.FuncOp.from_py_func(
+        RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
+    def test_f32_elemwise_neg(input, init_result):
+      return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.negf)
+
     # Just check that we don't assert out on name mismatch.
     # CHECK-LABEL: @test_non_default_op_name
     @builtin.FuncOp.from_py_func(
         RankedTensorType.get((42,), f32), RankedTensorType.get((42,), f32))
     def test_non_default_op_name(input, init_result):
       return non_default_op_name(input, outs=[init_result])
 
 
 print(module)