[mlir][StandardOps] Add folders to AddFOp and SubFOp.

[hanchung]

Summary

• addf(x,0) -> x
• subf(x,0) -> x
• subf(x,x) -> 0

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

299

Please avoid hardcoding the constant op(hardcoding limits composability), as well as using getZeroAttr(getZeroAttr creates an attribute, which involves uniquing and thus acquiring a lock). operands contains attributes for operands of this operation that are constant, use that instead:

if (FloatAttr rhs = operands[1].dyn_cast_or_null<FloatAttr>())
  if (rhs.getValue().isZero())
    return op.lhs();

[hanchung]

address comments

[hanchung]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

299

I was trying to address vector cases. If this limits composability, let's go with your suggestion.

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2766

Is this valid even when x is NaN or Infinity?

[hanchung]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2766

I'm not so familiar with floating point here. In LLVM, this is true when fast-math flag is on. It is always valid from the perspective of Math, but not LLVM IR floating-point operations.

After revisiting some sections in LLVM doc, I also found that other folders in this patch are might not valid.

For example, ‘x-0’->’x’ is not a valid transformation if the rounding mode is “round.downward” or “round.dynamic” because if the value of ‘x’ is +0 then ‘x-0’ should evaluate to ‘-0’ when rounding downward. However, this transformation is legal for all other rounding modes.

https://llvm.org/docs/LangRef.html#constrained-floating-point-intrinsics

I think the question would be, do we map subf/addf to LLVM intrinsics or math operations?

I also noticed that there is a TODO in the doc of addf

TODO: In the distant future, this will accept optional attributes for fast math, contraction, rounding mode, and other controls.

In this context, it seems that we follow LLVM intrinsics, and this is only valid if we have fast math attribute (which is set to true)?

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2766

By "LLVM intrinsics" you mean "LLVM instructions" here I think? (intrinsics and instructions are different classes in LLVM)?

Unfortunately floating point isn't just arbitrary math and we're limited here.

[hanchung]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2766

Yes, I meant LLVM instructions. Got you, thanks!

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static Type getTensorTypeFromMemRefType(Type type) {
   return NoneType::get(type.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // AddFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult AddFOp::fold(ArrayRef<Attribute> operands) {
+  // addf(x,0) -> x
+  if (auto op = rhs().getDefiningOp<ConstantOp>())
+    if (op.value() == Builder(getContext()).getZeroAttr(getType()))
+      return lhs();
+
   return constFoldBinaryOp<FloatAttr>(
       operands, [](APFloat a, APFloat b) { return a + b; });
 }
 
 //===----------------------------------------------------------------------===//
 // AddIOp
 //===----------------------------------------------------------------------===//
 
LogicalResult StoreOp::fold(ArrayRef<Attribute> cstOperands,
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // SubFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SubFOp::fold(ArrayRef<Attribute> operands) {
+  // subf(x,x) -> 0
+  if (getOperand(0) == getOperand(1))
+    return Builder(getContext()).getZeroAttr(getType());
+  // subf(x,0) -> x
+  if (auto op = rhs().getDefiningOp<ConstantOp>())
+    if (op.value() == Builder(getContext()).getZeroAttr(getType()))
+      return lhs();
+
   return constFoldBinaryOp<FloatAttr>(
       operands, [](APFloat a, APFloat b) { return a - b; });
 }
 
 //===----------------------------------------------------------------------===//
 // SubIOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Standard/canonicalize.mlir

 // Tensor cast is folded away.
 //   CHECK-NOT:   tensor.cast
 //       CHECK:   return %[[S]] : tensor<4x6x16x32xi8>
 func @rank_reducing_subtensor_insert_of_cast(%a : tensor<16x32xi8>, %b : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {
   %cast = tensor.cast %a : tensor<16x32xi8> to tensor<?x32xi8>
   %res = subtensor_insert %cast into %b[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor<?x32xi8> into tensor<4x6x16x32xi8>
   return %res : tensor<4x6x16x32xi8>
 }
+
+// -----
+
+// CHECK-LABEL: func @addf_zero
+//  CHECK-SAME:   %[[ARG:.[a-z0-9A-Z_]+]]: f32
+//       CHECK:   return %[[ARG]] : f32
+func @addf_zero(%arg0: f32) -> f32 {
+  %cst = constant 0.0 : f32
+  %0 = addf %arg0, %cst : f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: func @addf_zero_vector
+//  CHECK-SAME:   %[[ARG:.[a-z0-9A-Z_]+]]: vector<4xf32>
+//       CHECK:   return %[[ARG]] : vector<4xf32>
+func @addf_zero_vector(%arg0: vector<4xf32>) -> vector<4xf32> {
+  %cst = constant dense<0.0> : vector<4xf32>
+  %0 = addf %arg0, %cst : vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @subf_zero
+//  CHECK-SAME:   %[[ARG:.[a-z0-9A-Z_]+]]: f32
+//       CHECK:   return %[[ARG]] : f32
+func @subf_zero(%arg0: f32) -> f32 {
+  %cst = constant 0.0 : f32
+  %0 = subf %arg0, %cst : f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: func @subf_zero_vector
+//  CHECK-SAME:   %[[ARG:.[a-z0-9A-Z_]+]]: vector<4xf32>
+//       CHECK:   return %[[ARG]] : vector<4xf32>
+func @subf_zero_vector(%arg0: vector<4xf32>) -> vector<4xf32> {
+  %cst = constant dense<0.0> : vector<4xf32>
+  %0 = subf %arg0, %cst : vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: func @subf_same
+//       CHECK:   %[[ZERO:.+]] = constant 0.000000e+00 : f32
+//       CHECK:   return %[[ZERO]] : f32
+func @subf_same(%arg0: f32) -> f32 {
+  %0 = subf %arg0, %arg0 : f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: func @subf_same_vector
+//       CHECK:   %[[ZERO:.+]] = constant dense<0.000000e+00> : vector<4xf32>
+//       CHECK:   return %[[ZERO]] : vector<4xf32>
+func @subf_same_vector(%arg0: vector<4xf32>) -> vector<4xf32> {
+  %0 = subf %arg0, %arg0 : vector<4xf32>
+  return %0 : vector<4xf32>
+}