[MLIR] Support lowering n-D arith.index_cast to LLVM

[pzread]

Summary

Previously we can only lower arith.index_cast with 1-D vectors to LLVM. This change added the support for n-D vectors.

[ftynse]

mlir/lib/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.cpp

137–143

Can't this just be typeConverter->convertType(getElementTypeOrSelf(op.getIn().getType()) ?

[pzread]

mlir/lib/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.cpp

137–143

getElementTypeOrSelf can only extracts one level of a ShapedType.

[ftynse]

mlir/lib/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.cpp

137–143

I know, and the original type op.getIn().getType() is a single-level vector type because we can't nest vectors. The code does literally the same for targetElementType in line 135.

[hanchung]

mlir/lib/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.cpp

137–143

The op.getIn().getType() is vector types in this case. It's not the converted type. (The method you use is the converted type -- adaptor.getIn().getType(). As ftynse's suggestion, I think the typeConverter can handle the case?

145–146

I think we can use getIntOrFloatBitWidth(). It looks cleaner to me.

... = targetElementType.getIntOrFloatBitWidth();
...

[pzread]

Use getIntOrFloatBitWidth to get bit width.

[pzread]

mlir/lib/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.cpp

137–143

Oh, I'm sorry, I misread the op.getIn() as adaptor.getIn().

I'm not familiar with rewrite, but will it be a case where the operand of arith.index_cast has been rewritten into another type with a different element type (e.g., extension)? Therefore the element type from op.getIn() and adaptor.getIn() are different.

[pzread]

Simplify the method to derive source type.

mlir/lib/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.cpp

 
 //===----------------------------------------------------------------------===//
 // IndexCastOpLowering
 //===----------------------------------------------------------------------===//
 
 LogicalResult IndexCastOpLowering::matchAndRewrite(
     arith::IndexCastOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
-  auto targetType = typeConverter->convertType(op.getResult().getType());
+  auto resultType = op.getResult().getType();
   auto targetElementType =
-      typeConverter->convertType(getElementTypeOrSelf(op.getResult()))
-          .cast<IntegerType>();
-  auto sourceElementType =
-      getElementTypeOrSelf(adaptor.getIn()).cast<IntegerType>();
-  unsigned targetBits = targetElementType.getWidth();
-  unsigned sourceBits = sourceElementType.getWidth();
+      typeConverter->convertType(getElementTypeOrSelf(resultType));
+  Type sourceElementType = adaptor.getIn().getType();
+  // Get the element type of an n-D array.
+  while (sourceElementType.isa<LLVM::LLVMArrayType>()) {
+    sourceElementType =
+        sourceElementType.cast<LLVM::LLVMArrayType>().getElementType();
+  }
+  sourceElementType = getElementTypeOrSelf(sourceElementType);
+
+  unsigned targetBits = targetElementType.getIntOrFloatBitWidth();
+  unsigned sourceBits = sourceElementType.getIntOrFloatBitWidth();
 
-  if (targetBits == sourceBits)
+  if (targetBits == sourceBits) {
     rewriter.replaceOp(op, adaptor.getIn());
-  else if (targetBits < sourceBits)
-    rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, targetType, adaptor.getIn());
+    return success();
+  }
+
+  // Handle the scalar and 1D vector cases.
+  auto operandType = adaptor.getIn().getType();
+  if (!operandType.isa<LLVM::LLVMArrayType>()) {
+    auto targetType = typeConverter->convertType(resultType);
+    if (targetBits < sourceBits)
+      rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, targetType,
+                                                 adaptor.getIn());
     else
-    rewriter.replaceOpWithNewOp<LLVM::SExtOp>(op, targetType, adaptor.getIn());
+      rewriter.replaceOpWithNewOp<LLVM::SExtOp>(op, targetType,
+                                                adaptor.getIn());
     return success();
   }
 
+  auto vectorType = resultType.dyn_cast<VectorType>();
+  if (!vectorType)
+    return rewriter.notifyMatchFailure(op, "expected vector result type");
+
+  return LLVM::detail::handleMultidimensionalVectors(
+      op.getOperation(), adaptor.getOperands(), *getTypeConverter(),
+      [&](Type llvm1DVectorTy, ValueRange operands) -> Value {
+        OpAdaptor adaptor(operands);
+        if (targetBits < sourceBits) {
+          return rewriter.create<LLVM::TruncOp>(op.getLoc(), llvm1DVectorTy,
+                                                adaptor.getIn());
+        }
+        return rewriter.create<LLVM::SExtOp>(op.getLoc(), llvm1DVectorTy,
+                                             adaptor.getIn());
+      },
+      rewriter);
+}
+
 //===----------------------------------------------------------------------===//
 // CmpIOpLowering
 //===----------------------------------------------------------------------===//
 
 // Convert arith.cmp predicate into the LLVM dialect CmpPredicate. The two enums
 // share numerical values so just cast.
 template <typename LLVMPredType, typename PredType>
 static LLVMPredType convertCmpPredicate(PredType pred) {

mlir/test/Conversion/ArithmeticToLLVM/convert-nd-vector-to-llvmir.mlir

func.func @select_2d(%arg0 : vector<4x3xi1>, %arg1 : vector<4x3xi32>, %arg2 : vector<4x3xi32>) {
   // CHECK: %[[EXTRACT1:.*]] = llvm.extractvalue %[[ARG0]][0] : !llvm.array<4 x vector<3xi1>>
   // CHECK: %[[EXTRACT2:.*]] = llvm.extractvalue %[[ARG1]][0] : !llvm.array<4 x vector<3xi32>>
   // CHECK: %[[EXTRACT3:.*]] = llvm.extractvalue %[[ARG2]][0] : !llvm.array<4 x vector<3xi32>>
   // CHECK: %[[SELECT:.*]] = llvm.select %[[EXTRACT1]], %[[EXTRACT2]], %[[EXTRACT3]] : vector<3xi1>, vector<3xi32>
   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[SELECT]], %{{.*}}[0] : !llvm.array<4 x vector<3xi32>>
   %0 = arith.select %arg0, %arg1, %arg2 : vector<4x3xi1>, vector<4x3xi32>
   func.return
 }
+
+// CHECK-LABEL: func @index_cast_2d(
+// CHECK-SAME:    %[[ARG0:.*]]: vector<1x2x3xi1>)
+func.func @index_cast_2d(%arg0: vector<1x2x3xi1>) {
+  // CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]
+  // CHECK: %[[EXTRACT1:.*]] = llvm.extractvalue %[[SRC]][0, 0] : !llvm.array<1 x array<2 x vector<3xi1>>>
+  // CHECK: %[[SEXT1:.*]] = llvm.sext %[[EXTRACT1]] : vector<3xi1> to vector<3xi{{.*}}>
+  // CHECK: %[[INSERT1:.*]] = llvm.insertvalue %[[SEXT1]], %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
+  // CHECK: %[[EXTRACT2:.*]] = llvm.extractvalue %[[SRC]][0, 1] : !llvm.array<1 x array<2 x vector<3xi1>>>
+  // CHECK: %[[SEXT2:.*]] = llvm.sext %[[EXTRACT2]] : vector<3xi1> to vector<3xi{{.*}}>
+  // CHECK: %[[INSERT2:.*]] = llvm.insertvalue %[[SEXT2]], %[[INSERT1]][0, 1] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
+  %0 = arith.index_cast %arg0: vector<1x2x3xi1> to vector<1x2x3xindex>
+  // CHECK: %[[EXTRACT3:.*]] = llvm.extractvalue %[[INSERT2]][0, 0] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
+  // CHECK: %[[TRUNC1:.*]] = llvm.trunc %[[EXTRACT3]] : vector<3xi{{.*}}> to vector<3xi1>
+  // CHECK: %[[INSERT3:.*]] = llvm.insertvalue %[[TRUNC1]], %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi1>>>
+  // CHECK: %[[EXTRACT4:.*]] = llvm.extractvalue %[[INSERT2]][0, 1] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
+  // CHECK: %[[TRUNC2:.*]] = llvm.trunc %[[EXTRACT4]] : vector<3xi{{.*}}> to vector<3xi1>
+  // CHECK: %[[INSERT4:.*]] = llvm.insertvalue %[[TRUNC2]], %[[INSERT3]][0, 1] : !llvm.array<1 x array<2 x vector<3xi1>>>
+  %1 = arith.index_cast %0: vector<1x2x3xindex> to vector<1x2x3xi1>
+  return
+}