[mlir][arith] Extend `floordivsi` expansion

[awarzynski]

Summary

This patch extends createConst so that it can generate constant
vectors (it already generates scalars). This allows the expansion of
arith.floordivsi operating on vectors.

While arith.floordivsi is my main motivation for this change, this
patch should allow other Arith ops to be extended in vector cases too.

[dcaballe]

mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp

29

mlir::getElementTypeOrSelf(type) should help here

40

we can create a vector constant here instead of a broadcast by passing a vector type to the arith::ConstantOp

[awarzynski]

Address Diego's comments

[dcaballe]

LGTM! Just a minor comment that you can fix before landing.

mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp

32

dyn_cast_or_null -> dyn_cast, type can't be null, right?

mlir/include/mlir/Dialect/Arith/Transforms/Passes.td

let options = [
     Option<"alignment", "alignment", "unsigned", /*default=*/"0",
            "Create global memrefs with a specified alignment">,
   ];
 }
 
 def ArithExpandOps : Pass<"arith-expand"> {
   let summary = "Legalize Arith ops to be convertible to LLVM.";
   let constructor = "mlir::arith::createArithExpandOpsPass()";
+  let dependentDialects = ["vector::VectorDialect"];
 }
 
 def ArithUnsignedWhenEquivalent : Pass<"arith-unsigned-when-equivalent"> {
   let summary = "Replace signed ops with unsigned ones where they are proven equivalent";
   let description = [{
     Replace signed ops with their unsigned equivalents when integer range analysis
     determines that their arguments and results are all guaranteed to be
     non-negative when interpreted as signed integers. When this occurs,

mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp

 //===- ExpandOps.cpp - Pass to legalize Arith ops for LLVM lowering --===//
 //
 // 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/Dialect/Arith/Transforms/Passes.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace arith {
 #define GEN_PASS_DEF_ARITHEXPANDOPS
 #include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
 } // namespace arith
 } // namespace mlir
 
 using namespace mlir;
 
 /// Create an integer or index constant.
 static Value createConst(Location loc, Type type, int value,
                          PatternRewriter &rewriter) {
-  return rewriter.create<arith::ConstantOp>(
-      loc, rewriter.getIntegerAttr(type, value));
+  Type elTy = type;
+  ShapedType vecTy = llvm::dyn_cast_or_null<ShapedType>(type);
+  if (vecTy)
+    elTy = vecTy.getElementType();
+
+  auto constantOp = rewriter.create<arith::ConstantOp>(
+      loc, rewriter.getIntegerAttr(elTy, value));
+
+  if (!vecTy)
+    return constantOp;
+
+  auto broadCastOp = rewriter.create<vector::BroadcastOp>(
+      loc, VectorType::get(vecTy.getShape(), elTy), constantOp);
+
+  return broadCastOp;
 }
 
 namespace {
 
 /// Expands CeilDivUIOp (n, m) into
 ///  n == 0 ? 0 : ((n-1) / m) + 1
 struct CeilDivUIOpConverter : public OpRewritePattern<arith::CeilDivUIOp> {
   using OpRewritePattern::OpRewritePattern;

mlir/test/Dialect/Arith/expand-ops.mlir

 // CHECK:           [[TERM1:%.+]] = arith.andi [[NNEG]], [[MPOS]] : i1
 // CHECK:           [[TERM2:%.+]] = arith.andi [[NPOS]], [[MNEG]] : i1
 // CHECK:           [[CMP2:%.+]] = arith.ori [[TERM1]], [[TERM2]] : i1
 // CHECK:           [[RES:%.+]] = arith.select [[CMP2]], [[TRUE3]], [[FALSE]] : index
 }
 
 // -----
 
+// Test floor divide with vector
+// CHECK-LABEL:   func.func @floordivi_vec(
+// CHECK-SAME:                             %[[VAL_0:.*]]: vector<4xi32>,
+// CHECK-SAME:                             %[[VAL_1:.*]]: vector<4xi32>) -> vector<4xi32> {
+func.func @floordivi_vec(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> (vector<4xi32>) {
+  %res = arith.floordivsi %arg0, %arg1 : vector<4xi32>
+  return %res : vector<4xi32>
+// CHECK:           %[[VAL_2:.*]] = arith.constant 1 : i32
+// CHECK:           %[[VAL_3:.*]] = arith.constant dense<1> : vector<4xi32>
+// CHECK:           %[[VAL_4:.*]] = arith.constant 0 : i32
+// CHECK:           %[[VAL_5:.*]] = arith.constant dense<0> : vector<4xi32>
+// CHECK:           %[[VAL_6:.*]] = arith.constant -1 : i32
+// CHECK:           %[[VAL_7:.*]] = arith.constant dense<-1> : vector<4xi32>
+// CHECK:           %[[VAL_8:.*]] = arith.cmpi slt, %[[VAL_1]], %[[VAL_5]] : vector<4xi32>
+// CHECK:           %[[VAL_9:.*]] = arith.select %[[VAL_8]], %[[VAL_3]], %[[VAL_7]] : vector<4xi1>, vector<4xi32>
+// CHECK:           %[[VAL_10:.*]] = arith.subi %[[VAL_9]], %[[VAL_0]] : vector<4xi32>
+// CHECK:           %[[VAL_11:.*]] = arith.divsi %[[VAL_10]], %[[VAL_1]] : vector<4xi32>
+// CHECK:           %[[VAL_12:.*]] = arith.subi %[[VAL_7]], %[[VAL_11]] : vector<4xi32>
+// CHECK:           %[[VAL_13:.*]] = arith.divsi %[[VAL_0]], %[[VAL_1]] : vector<4xi32>
+// CHECK:           %[[VAL_14:.*]] = arith.cmpi slt, %[[VAL_0]], %[[VAL_5]] : vector<4xi32>
+// CHECK:           %[[VAL_15:.*]] = arith.cmpi sgt, %[[VAL_0]], %[[VAL_5]] : vector<4xi32>
+// CHECK:           %[[VAL_16:.*]] = arith.cmpi slt, %[[VAL_1]], %[[VAL_5]] : vector<4xi32>
+// CHECK:           %[[VAL_17:.*]] = arith.cmpi sgt, %[[VAL_1]], %[[VAL_5]] : vector<4xi32>
+// CHECK:           %[[VAL_18:.*]] = arith.andi %[[VAL_14]], %[[VAL_17]] : vector<4xi1>
+// CHECK:           %[[VAL_19:.*]] = arith.andi %[[VAL_15]], %[[VAL_16]] : vector<4xi1>
+// CHECK:           %[[VAL_20:.*]] = arith.ori %[[VAL_18]], %[[VAL_19]] : vector<4xi1>
+// CHECK:           %[[VAL_21:.*]] = arith.select %[[VAL_20]], %[[VAL_12]], %[[VAL_13]] : vector<4xi1>, vector<4xi32>
+// CHECK:           return %[[VAL_21]] : vector<4xi32>
+// CHECK:         }
+}
+
+// -----
+
 // Test ceil divide with unsigned integer
 // CHECK-LABEL:       func @ceildivui
 // CHECK-SAME:     ([[ARG0:%.+]]: i32, [[ARG1:%.+]]: i32) -> i32 {
 func.func @ceildivui(%arg0: i32, %arg1: i32) -> (i32) {
   %res = arith.ceildivui %arg0, %arg1 : i32
   return %res : i32
 // CHECK:           [[ZERO:%.+]] = arith.constant 0 : i32
 // CHECK:           [[ISZERO:%.+]] = arith.cmpi eq, %arg0, [[ZERO]] : i32