[mlir] support nD vector forms of shifts in std-to-llvm conversion

These ops were not ported to the nD vector conversion when it was introduced
and nobody needed them so far.

Reviewed By: gysit

Differential Revision: https://reviews.llvm.org/D107750

2 files changed, 45 insertions, 3 deletions

diff --git a/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp b/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
index 4079e09..b4725d2 100644
--- a/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
+++ b/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
@@ -397,13 +397,13 @@ using SelectOpLowering = VectorConvertToLLVMPattern<SelectOp, LLVM::SelectOp>;
 using SignExtendIOpLowering =
     VectorConvertToLLVMPattern<SignExtendIOp, LLVM::SExtOp>;
 using ShiftLeftOpLowering =
-    OneToOneConvertToLLVMPattern<ShiftLeftOp, LLVM::ShlOp>;
+    VectorConvertToLLVMPattern<ShiftLeftOp, LLVM::ShlOp>;
 using SignedDivIOpLowering =
     VectorConvertToLLVMPattern<SignedDivIOp, LLVM::SDivOp>;
 using SignedRemIOpLowering =
     VectorConvertToLLVMPattern<SignedRemIOp, LLVM::SRemOp>;
 using SignedShiftRightOpLowering =
-    OneToOneConvertToLLVMPattern<SignedShiftRightOp, LLVM::AShrOp>;
+    VectorConvertToLLVMPattern<SignedShiftRightOp, LLVM::AShrOp>;
 using SubFOpLowering = VectorConvertToLLVMPattern<SubFOp, LLVM::FSubOp>;
 using SubIOpLowering = VectorConvertToLLVMPattern<SubIOp, LLVM::SubOp>;
 using TruncateIOpLowering =
@@ -414,7 +414,7 @@ using UnsignedDivIOpLowering =
 using UnsignedRemIOpLowering =
     VectorConvertToLLVMPattern<UnsignedRemIOp, LLVM::URemOp>;
 using UnsignedShiftRightOpLowering =
-    OneToOneConvertToLLVMPattern<UnsignedShiftRightOp, LLVM::LShrOp>;
+    VectorConvertToLLVMPattern<UnsignedShiftRightOp, LLVM::LShrOp>;
 using XOrOpLowering = VectorConvertToLLVMPattern<XOrOp, LLVM::XOrOp>;
 using ZeroExtendIOpLowering =
     VectorConvertToLLVMPattern<ZeroExtendIOp, LLVM::ZExtOp>;
diff --git a/mlir/test/Conversion/StandardToLLVM/convert-nd-vector-to-llvmir.mlir b/mlir/test/Conversion/StandardToLLVM/convert-nd-vector-to-llvmir.mlir
index fe98d29..5ec2692 100644
--- a/mlir/test/Conversion/StandardToLLVM/convert-nd-vector-to-llvmir.mlir
+++ b/mlir/test/Conversion/StandardToLLVM/convert-nd-vector-to-llvmir.mlir
@@ -138,3 +138,45 @@ func @trunci_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi16> {
   %0 = trunci %arg0: vector<1x2x3xi64> to vector<1x2x3xi16>
   return %0 : vector<1x2x3xi16>
 }
+
+// CHECK-LABEL: @shl
+func @shl_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi64> {
+  // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.shl %{{.*}}, %{{.*}} : vector<3xi64>
+  // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.shl %{{.*}}, %{{.*}} : vector<3xi64>
+  // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  %c1 = constant dense<1> : vector<1x2x3xi64>
+  %0 = shift_left %arg0, %c1 : vector<1x2x3xi64>
+  return %0 : vector<1x2x3xi64>
+}
+
+// CHECK-LABEL: @shrs
+func @shrs_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi64> {
+  // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.ashr %{{.*}}, %{{.*}} : vector<3xi64>
+  // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.ashr %{{.*}}, %{{.*}} : vector<3xi64>
+  // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  %c1 = constant dense<1> : vector<1x2x3xi64>
+  %0 = shift_right_signed %arg0, %c1 : vector<1x2x3xi64>
+  return %0 : vector<1x2x3xi64>
+}
+
+// CHECK-LABEL: @shru
+func @shru_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi64> {
+  // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.lshr %{{.*}}, %{{.*}} : vector<3xi64>
+  // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  // CHECK: llvm.lshr %{{.*}}, %{{.*}} : vector<3xi64>
+  // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
+  %c1 = constant dense<1> : vector<1x2x3xi64>
+  %0 = shift_right_unsigned %arg0, %c1 : vector<1x2x3xi64>
+  return %0 : vector<1x2x3xi64>
+}