Fix an error on krnl-to-affine conversion for onnx.CategoryMapper with rank-2+ inputs (#2086)

This PR includes fixes an error on krnl-to-affine conversion for onnx.CategoryMapper with rank 2 inputs and their lit tests.

Signed-off-by: Yasushi Negishi <[email protected]>

Author: negiyas

src/Conversion/ONNXToKrnl/ML/CategoryMapper.cpp

@@ -87,6 +87,7 @@ struct ONNXCategoryMapperOpLowering
 
     // Basic information.
     int64_t rank = memRefType.getShape().size();
+    assert(((rank == 1) || (rank == 2)) && "Invalid rank of input");
     ShapedType inputType = X.getType().cast<ShapedType>();
     Type elementType = inputType.getElementType();
 
@@ -259,8 +260,14 @@ struct ONNXCategoryMapperOpLowering
           MathBuilder createMath(createKrnl);
           Value zero = createMath.constant(
               createMath.getBuilder().getIntegerType(64), 0);
+          ArrayRef<int64_t> shape =
+              memref.getType().cast<ShapedType>().getShape();
+          SmallVector<int64_t, 4> newShape;
+          for (uint64_t i = 0; i < shape.size(); i++)
+            newShape.emplace_back(
+                (shape[i] == ShapedType::kDynamic) ? 1 : shape[i]);
           auto memRefType = MemRefType::get(
-              {rank}, krnl::StringType::get(elementType.getContext()));
+              newShape, krnl::StringType::get(elementType.getContext()));
           Value stringMemRef = createKrnl.getRef(memRefType, memref, zero);
           inputElem = createKrnl.load(stringMemRef, loopInd);
         })

src/Dialect/ONNX/ONNXOps/ML/CategoryMapper.cpp

@@ -47,6 +47,8 @@ LogicalResult ONNXCategoryMapperOp::verify() {
   }
 
   ShapedType inputType = X.getType().cast<ShapedType>();
+  if ((inputType.getRank() != 1) && (inputType.getRank() != 2))
+    return emitOpError("input rank must be one or two");
   Type elementType = inputType.getElementType();
   if (!elementType.isInteger(64) && !elementType.isa<ONNXStringType>())
     return emitOpError("input must be a tensor of int64 or string");

src/Dialect/ONNX/ONNXOps/Math/ElementwiseUnary.cpp

@@ -32,6 +32,8 @@ namespace onnx_mlir {
 /// Handle shape inference for unary element-wise operators.
 LogicalResult inferShapeForUnaryOps(Operation *op) {
   Value input = op->getOperand(0);
+  if (!hasShapeAndRank(input))
+    return success();
   RankedTensorType inputType = input.getType().dyn_cast<RankedTensorType>();
   return inferShapeForUnaryOps(
       op, inputType.getElementType(), inputType.getEncoding());
@@ -41,6 +43,8 @@ LogicalResult inferShapeForUnaryOps(Operation *op) {
 /// type.
 LogicalResult inferShapeForUnaryOps(Operation *op, Type elementType) {
   Value input = op->getOperand(0);
+  if (!hasShapeAndRank(input))
+    return success();
   RankedTensorType inputType = input.getType().dyn_cast<RankedTensorType>();
   return inferShapeForUnaryOps(op, elementType, inputType.getEncoding());
 }

test/mlir/onnx/onnx_lowering_category_mapper.mlir

@@ -20,8 +20,8 @@ func.func private @test_category_mapper_string_to_int64(%arg0 : tensor<2x2x!onnx
   // CHECK-DAG: [[LOOP_0:%.+]]:2 = krnl.define_loops 2
   // CHECK:     krnl.iterate([[LOOP_0]]#0, [[LOOP_0]]#1) with ([[LOOP_0]]#0 -> [[I_0:%.+]] = 0 to 2, [[LOOP_0]]#1 -> [[I_1:%.+]] = 0 to 2){  
   // CHECK:     [[IVS:%.+]]:2 = krnl.get_induction_var_value([[LOOP_0]]#0, [[LOOP_0]]#1) : (!krnl.loop, !krnl.loop) -> (index, index)
-  // CHECK:     [[REF:%.+]] = "krnl.getref"(%arg0, [[ZERO_i64]]) : (memref<2x2x!krnl.string>, i64) -> memref<2x!krnl.string>
-  // CHECK:     [[LOAD1:%.+]] = krnl.load [[REF]]{{.}}[[IVS]]#0, [[IVS]]#1{{.}} : memref<2x!krnl.string>
+  // CHECK:     [[REF:%.+]] = "krnl.getref"(%arg0, [[ZERO_i64]]) : (memref<2x2x!krnl.string>, i64) -> memref<2x2x!krnl.string>
+  // CHECK:     [[LOAD1:%.+]] = krnl.load [[REF]]{{.}}[[IVS]]#0, [[IVS]]#1{{.}} : memref<2x2x!krnl.string>
   // CHECK:     [[INDEX:%.+]] = "krnl.find_index"([[LOAD1]], [[G]], [[V]], [[LEN]]) : (!krnl.string, memref<3xi32>, memref<3xi32>, i32) -> index
   // CHECK:     [[LOAD2:%.+]] = krnl.load [[CAT_STRINGS]]{{.}}[[INDEX]]{{.}} : memref<3x!krnl.string>
   // CHECK:     [[STRLEN:%.+]] = "krnl.strlen"([[LOAD2]]) : (!krnl.string) -> i64