Support lowering of index-cast on vector types.

[arpith-jacob]

Summary

The index cast operation accepts vector types. Implement its lowering in this patch.

[ftynse]

Could we update the doc of the index cast op while we are here? It current says "casts between <...> scalars".

[arpith-jacob]

Updated docs.

[arpith-jacob]

Done.

Could you commit this patch? Thanks!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 
 //===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
 def IndexCastOp : ArithmeticCastOp<"index_cast">, Arguments<(ins AnyType:$in)> {
   let summary = "cast between index and integer types";
   let description = [{
-    Casts between integer scalars and 'index' scalars. Index is an integer of
-    platform-specific bit width. If casting to a wider integer, the value is
-    sign-extended. If casting to a narrower integer, the value is truncated.
+    Casts between scalar or vector integers and corresponding 'index' scalar or
+    vectors. Index is an integer of platform-specific bit width. If casting to
+    a wider integer, the value is sign-extended. If casting to a narrower
+    integer, the value is truncated.
   }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MulFOp

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

struct IndexCastOpLowering : public ConvertOpToLLVMPattern<IndexCastOp> {
   using ConvertOpToLLVMPattern<IndexCastOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(IndexCastOp indexCastOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     IndexCastOpAdaptor transformed(operands);
 
     auto targetType =
-        typeConverter->convertType(indexCastOp.getResult().getType())
+        typeConverter->convertType(indexCastOp.getResult().getType());
+    auto targetElementType =
+        typeConverter
+            ->convertType(getElementTypeOrSelf(indexCastOp.getResult()))
             .cast<IntegerType>();
-    auto sourceType = transformed.in().getType().cast<IntegerType>();
-    unsigned targetBits = targetType.getWidth();
-    unsigned sourceBits = sourceType.getWidth();
+    auto sourceElementType =
+        getElementTypeOrSelf(transformed.in()).cast<IntegerType>();
+    unsigned targetBits = targetElementType.getWidth();
+    unsigned sourceBits = sourceElementType.getWidth();
 
     if (targetBits == sourceBits)
       rewriter.replaceOp(indexCastOp, transformed.in());
     else if (targetBits < sourceBits)
       rewriter.replaceOpWithNewOp<LLVM::TruncOp>(indexCastOp, targetType,
                                                  transformed.in());
     else
       rewriter.replaceOpWithNewOp<LLVM::SExtOp>(indexCastOp, targetType,

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

 func @index_cast(%arg0: index, %arg1: i1) {
 // CHECK-NEXT: = llvm.trunc %arg0 : i{{.*}} to i1
   %0 = index_cast %arg0: index to i1
 // CHECK-NEXT: = llvm.sext %arg1 : i1 to i{{.*}}
   %1 = index_cast %arg1: i1 to index
   return
 }
 
+// CHECK-LABEL: @vector_index_cast
+func @vector_index_cast(%arg0: vector<2xindex>, %arg1: vector<2xi1>) {
+// CHECK-NEXT: = llvm.trunc %{{.*}} : vector<2xi{{.*}}> to vector<2xi1>
+  %0 = index_cast %arg0: vector<2xindex> to vector<2xi1>
+// CHECK-NEXT: = llvm.sext %{{.*}} : vector<2xi1> to vector<2xi{{.*}}>
+  %1 = index_cast %arg1: vector<2xi1> to vector<2xindex>
+  return
+}
+
 // Checking conversion of signed integer types to floating point.
 // CHECK-LABEL: @sitofp
 func @sitofp(%arg0 : i32, %arg1 : i64) {
 // CHECK-NEXT: = llvm.sitofp {{.*}} : i32 to f32
   %0 = sitofp %arg0: i32 to f32
 // CHECK-NEXT: = llvm.sitofp {{.*}} : i32 to f64
   %1 = sitofp %arg0: i32 to f64
 // CHECK-NEXT: = llvm.sitofp {{.*}} : i64 to f32