[mlir] Support multi-dimensional vectors in MathToLibm conversion.

[akuegel]

Summary

This resolves a TODO.

[akuegel]

Also add test with multi-dimensional vector

[ftynse]

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

62

You can just call vecType.getNumElements().

63

Nit: LLVM uses unsigned in loops.

68–73

There's this https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Vector/VectorUtils.h#L56

[akuegel]

Use getNumElements() call.

[akuegel]

Thanks for the review :)

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

68–73

This needs computed strides, I cannot just pass it the shape. I am not sure whether it is worth it to compute the strides first, just that I can call this function.

[ftynse]

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

68–73

There's also computeStrides a couple of lines above. :) So you can call two functions and be done with it instead of doing somewhat non-trivial math.

[akuegel]

Use computeStrides and delinearize instead of custom delinearization.

[akuegel]

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

68–73

It is not the perfect utility function for my use case, but ok, I can understand that it is easier to read to just have two function calls than the custom delinearization. Done.

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

 //===-- MathToLibm.cpp - conversion from Math to libm calls ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/MathToLibm/MathToLibm.h"
 
 #include "../PassDetail.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/Dialect/Vector/VectorUtils.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
 
 namespace {
 // Pattern to convert vector operations to scalar operations. This is needed as
 // libm calls require scalars.
VecOpToScalarOp<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {
   auto loc = op.getLoc();
   auto vecType = opType.template dyn_cast<VectorType>();
 
   if (!vecType)
     return failure();
   if (!vecType.hasRank())
     return failure();
   auto shape = vecType.getShape();
-  // TODO: support multidimensional vectors
-  if (shape.size() != 1)
-    return failure();
+  int64_t numElements = vecType.getNumElements();
 
   Value result = rewriter.create<arith::ConstantOp>(
       loc, DenseElementsAttr::get(
                vecType, FloatAttr::get(vecType.getElementType(), 0.0)));
-  for (auto i = 0; i < shape.front(); ++i) {
+  SmallVector<int64_t> ones(shape.size(), 1);
+  SmallVector<int64_t> strides = computeStrides(shape, ones);
+  for (auto linearIndex = 0; linearIndex < numElements; ++linearIndex) {
+    SmallVector<int64_t> positions = delinearize(strides, linearIndex);
     SmallVector<Value> operands;
     for (auto input : op->getOperands())
       operands.push_back(
-          rewriter.create<vector::ExtractElementOp>(loc, input, i));
+          rewriter.create<vector::ExtractOp>(loc, input, positions));
     Value scalarOp =
         rewriter.create<Op>(loc, vecType.getElementType(), operands);
-    result = rewriter.create<vector::InsertElementOp>(loc, scalarOp, result, i);
+    result =
+        rewriter.create<vector::InsertOp>(loc, scalarOp, result, positions);
   }
   rewriter.replaceOp(op, {result});
   return success();
 }
 
 template <typename Op>
 LogicalResult
 ScalarOpToLibmCall<Op>::matchAndRewrite(Op op,

mlir/test/Conversion/MathToLibm/convert-to-libm.mlir

func @expm1_vec_caller(%float: vector<2xf32>, %double: vector<2xf64>) -> (vector<2xf32>, vector<2xf64>) {
   %double_result = math.expm1 %double : vector<2xf64>
   return %float_result, %double_result : vector<2xf32>, vector<2xf64>
 }
 // CHECK-LABEL:   func @expm1_vec_caller(
 // CHECK-SAME:                           %[[VAL_0:.*]]: vector<2xf32>,
 // CHECK-SAME:                           %[[VAL_1:.*]]: vector<2xf64>) -> (vector<2xf32>, vector<2xf64>) {
 // CHECK-DAG:       %[[CVF:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>
 // CHECK-DAG:       %[[CVD:.*]] = arith.constant dense<0.000000e+00> : vector<2xf64>
-// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : i32
-// CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : i32
-// CHECK:           %[[IN0_F32:.*]] = vector.extractelement %[[VAL_0]]{{\[}}%[[C0]] : i32] : vector<2xf32>
+// CHECK:           %[[IN0_F32:.*]] = vector.extract %[[VAL_0]][0] : vector<2xf32>
 // CHECK:           %[[OUT0_F32:.*]] = call @expm1f(%[[IN0_F32]]) : (f32) -> f32
-// CHECK:           %[[VAL_8:.*]] = vector.insertelement %[[OUT0_F32]], %[[CVF]]{{\[}}%[[C0]] : i32] : vector<2xf32>
-// CHECK:           %[[IN1_F32:.*]] = vector.extractelement %[[VAL_0]]{{\[}}%[[C1]] : i32] : vector<2xf32>
+// CHECK:           %[[VAL_8:.*]] = vector.insert %[[OUT0_F32]], %[[CVF]] [0] : f32 into vector<2xf32>
+// CHECK:           %[[IN1_F32:.*]] = vector.extract %[[VAL_0]][1] : vector<2xf32>
 // CHECK:           %[[OUT1_F32:.*]] = call @expm1f(%[[IN1_F32]]) : (f32) -> f32
-// CHECK:           %[[VAL_11:.*]] = vector.insertelement %[[OUT1_F32]], %[[VAL_8]]{{\[}}%[[C1]] : i32] : vector<2xf32>
-// CHECK:           %[[IN0_F64:.*]] = vector.extractelement %[[VAL_1]]{{\[}}%[[C0]] : i32] : vector<2xf64>
+// CHECK:           %[[VAL_11:.*]] = vector.insert %[[OUT1_F32]], %[[VAL_8]] [1] : f32 into vector<2xf32>
+// CHECK:           %[[IN0_F64:.*]] = vector.extract %[[VAL_1]][0] : vector<2xf64>
 // CHECK:           %[[OUT0_F64:.*]] = call @expm1(%[[IN0_F64]]) : (f64) -> f64
-// CHECK:           %[[VAL_14:.*]] = vector.insertelement %[[OUT0_F64]], %[[CVD]]{{\[}}%[[C0]] : i32] : vector<2xf64>
-// CHECK:           %[[IN1_F64:.*]] = vector.extractelement %[[VAL_1]]{{\[}}%[[C1]] : i32] : vector<2xf64>
+// CHECK:           %[[VAL_14:.*]] = vector.insert %[[OUT0_F64]], %[[CVD]] [0] : f64 into vector<2xf64>
+// CHECK:           %[[IN1_F64:.*]] = vector.extract %[[VAL_1]][1] : vector<2xf64>
 // CHECK:           %[[OUT1_F64:.*]] = call @expm1(%[[IN1_F64]]) : (f64) -> f64
-// CHECK:           %[[VAL_17:.*]] = vector.insertelement %[[OUT1_F64]], %[[VAL_14]]{{\[}}%[[C1]] : i32] : vector<2xf64>
+// CHECK:           %[[VAL_17:.*]] = vector.insert %[[OUT1_F64]], %[[VAL_14]] [1] : f64 into vector<2xf64>
 // CHECK:           return %[[VAL_11]], %[[VAL_17]] : vector<2xf32>, vector<2xf64>
 // CHECK:         }
 
+func @expm1_multidim_vec_caller(%float: vector<2x2xf32>) -> (vector<2x2xf32>) {
+  %result = math.expm1 %float : vector<2x2xf32>
+  return %result : vector<2x2xf32>
+}
+// CHECK-LABEL:   func @expm1_multidim_vec_caller(
+// CHECK-SAME:                           %[[VAL:.*]]: vector<2x2xf32>
+// CHECK-DAG:       %[[CVF:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
+// CHECK:           %[[IN0_0_F32:.*]] = vector.extract %[[VAL]][0, 0] : vector<2x2xf32>
+// CHECK:           %[[OUT0_0_F32:.*]] = call @expm1f(%[[IN0_0_F32]]) : (f32) -> f32
+// CHECK:           %[[VAL_1:.*]] = vector.insert %[[OUT0_0_F32]], %[[CVF]] [0, 0] : f32 into vector<2x2xf32>
+// CHECK:           %[[IN0_1_F32:.*]] = vector.extract %[[VAL]][0, 1] : vector<2x2xf32>
+// CHECK:           %[[OUT0_1_F32:.*]] = call @expm1f(%[[IN0_1_F32]]) : (f32) -> f32
+// CHECK:           %[[VAL_2:.*]] = vector.insert %[[OUT0_1_F32]], %[[VAL_1]] [0, 1] : f32 into vector<2x2xf32>
+// CHECK:           %[[IN1_0_F32:.*]] = vector.extract %[[VAL]][1, 0] : vector<2x2xf32>
+// CHECK:           %[[OUT1_0_F32:.*]] = call @expm1f(%[[IN1_0_F32]]) : (f32) -> f32
+// CHECK:           %[[VAL_3:.*]] = vector.insert %[[OUT1_0_F32]], %[[VAL_2]] [1, 0] : f32 into vector<2x2xf32>
+// CHECK:           %[[IN1_1_F32:.*]] = vector.extract %[[VAL]][1, 1] : vector<2x2xf32>
+// CHECK:           %[[OUT1_1_F32:.*]] = call @expm1f(%[[IN1_1_F32]]) : (f32) -> f32
+// CHECK:           %[[VAL_4:.*]] = vector.insert %[[OUT1_1_F32]], %[[VAL_3]] [1, 1] : f32 into vector<2x2xf32>
+// CHECK:           return %[[VAL_4]] : vector<2x2xf32>
+// CHECK:         }