Error for compiling bidaf-9 in Krnl-to-Afffine conversion (The situation seems to be similar to #2150)

[negiyas]

Error similar to #2150 occurs occurs while compiling bidaf-9 by onnx-mlir -mcpu=z14 --EmitMLIR bidaf-9.onnx . Patch is needed to avoid #2153, before compilation (see below).

How to reproduce this issue. (1) Patch the following diff to avoid issue#2153.

$ git diff main
diff --git a/src/Conversion/ONNXToKrnl/ML/CategoryMapper.cpp b/src/Conversion/ONNXToKrnl/ML/CategoryMapper.cpp
index 0d63032e..5b3b3a64 100644
--- a/src/Conversion/ONNXToKrnl/ML/CategoryMapper.cpp
+++ b/src/Conversion/ONNXToKrnl/ML/CategoryMapper.cpp
@@ -87,7 +87,6 @@ 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();

diff --git a/src/Dialect/ONNX/ONNXOps/ML/CategoryMapper.cpp b/src/Dialect/ONNX/ONNXOps/ML/CategoryMapper.cpp
index 623b84f4..1148ef8c 100644
--- a/src/Dialect/ONNX/ONNXOps/ML/CategoryMapper.cpp
+++ b/src/Dialect/ONNX/ONNXOps/ML/CategoryMapper.cpp
@@ -47,8 +47,6 @@ 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");

(2) Build onnx-mlir (3) Compile bidaf-9.onnx ( https://github.com/onnx/models/blob/main/text/machine_comprehension/bidirectional_attention_flow/model/bidaf-9.onnx ) by onnx-mlir as follows.

$ ./build/Debug/bin/onnx-mlir -mcpu=z14 --EmitMLIR  bidaf-9.onnx
onnx-mlir: ../llvm-project/mlir/lib/Dialect/Affine/IR/AffineOps.cpp:2999:
    static void mlir::AffineLoadOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::Value, mlir::AffineMap, mlir::ValueRange):
     Assertion `map.getNumInputs() == mapOperands.size() && "inconsistent index info"' failed.
Aborted (core dumped)

The situation The second krnl.load in the following loop %463 = krnl.load %alloc_216[%461#0, %461#1] : memref<1xf32> is invalid, because the krnl.load of has two indices, but the memref rank is one (The rank should be two).

      %454:2 = krnl.define_loops 2
      krnl.iterate(%454#0, %454#1) with (%454#0 -> %arg5 = 0 to 1, %454#1 -> %arg6 = 0 to 1){
        %461:2 = krnl.get_induction_var_value(%454#0, %454#1) : (!krnl.loop, !krnl.loop) -> (index, index)
        %462 = krnl.load %alloc_234[%461#0, %461#1] : memref<1x1xi1>
        %463 = krnl.load %alloc_216[%461#0, %461#1] : memref<1xf32>
        %464 = krnl.load %alloc_215[%461#0, %461#1] : memref<1x1xf32>
        %465 = arith.select %462, %463, %464 : f32
        krnl.store %465, %alloc_236[%461#0, %461#1] : memref<1x1xf32>
      }

[negiyas]

Have confirmed that this issue is fixed by PR#2158 ( https://github.com/onnx/onnx-mlir/pull/2158 ).

[AlexandreEichenberger]

@negiyas thanks for confirming, much appreciated. Note that @2158 is itself subsumed by #2161; I have checked that that PR solve the problem listed in #2158, so I am pretty sure it will work for this PR as well.