[mlir][memref] Missing type conversion in memref.reshape llvm lowering

[Hardcode84]

Summary

Shape can be memref of index type, so memref::LoadOp result need to be converted into llvm type.

[Hardcode84]

assert msg

[bondhugula]

Shape can be memref of index type

Trying to understand if it can be anything else?

[Hardcode84]

Trying to understand if it can be anything else?

It can be also a memref of plain integer type (there are tests for i64, but I doubt existing code will work for i32 or other integer types).

[ftynse]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1163–1169

It somehow feels that this (ab)uses the rewriter API. There is no remapped value for dimSize because it is created in the line above, so this just forces the rewriter to call the target materialization in case of type mismatch. Given the amount of TODO in that part of the rewriter, it is not unlikely that the behavior may change and break this in unexpected ways.

I would rather call the materialization directly from type converter if the loaded type is index. ReshapeOp only supports memref of integer or index anyway.

[Hardcode84]

use materializeTargetConversion directly

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

if (shapeMemRefType.hasStaticShape()) {
         int64_t size = targetMemRefType.getDimSize(i);
         // If the size of this dimension is dynamic, then load it at runtime
         // from the shape operand.
         if (!ShapedType::isDynamic(size)) {
           dimSize = createIndexConstant(rewriter, loc, size);
         } else {
           Value shapeOp = reshapeOp.getShape();
           Value index = createIndexConstant(rewriter, loc, i);
           dimSize = rewriter.create<memref::LoadOp>(loc, shapeOp, index);
+          Type indexType = getIndexType();
+          if (dimSize.getType() != indexType)
+            dimSize = typeConverter->materializeTargetConversion(
+                rewriter, loc, indexType, dimSize);
+          assert(dimSize && "Invalid memref element type");
         }
 
         desc.setSize(rewriter, loc, i, dimSize);
         desc.setStride(rewriter, loc, i, stride);
 
         // Prepare the stride value for the next dimension.
         stride = rewriter.create<LLVM::MulOp>(loc, stride, dimSize);
       }
 

mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir

func.func @memref.reshape.dynamic.dim(%arg: memref<?x?x?xf32>, %shape: memref<4xi64>) -> memref<?x?x12x32xf32> {
   // CHECK: %[[insert10:.*]] = llvm.insertvalue %[[mul]], %[[insert9]][4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
 
   // CHECK: %[[result_cast:.*]] = builtin.unrealized_conversion_cast %[[insert10]] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)> to memref<?x?x12x32xf32>
   %0 = memref.reshape %arg(%shape) : (memref<?x?x?xf32>, memref<4xi64>) -> memref<?x?x12x32xf32>
 
   return %0 : memref<?x?x12x32xf32>
   // CHECK: return %[[result_cast]] : memref<?x?x12x32xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @memref.reshape_index
+// CHECK-SAME:    %[[arg:.*]]: memref<?x?xi32>, %[[shape:.*]]: memref<1xindex>
+func.func @memref.reshape_index(%arg0: memref<?x?xi32>, %shape: memref<1xindex>) ->  memref<?xi32> {
+  // CHECK: %[[arg_cast:.*]] = builtin.unrealized_conversion_cast %[[arg]] : memref<?x?xi32> to !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: %[[shape_cast:.*]] = builtin.unrealized_conversion_cast %[[shape]] : memref<1xindex> to !llvm.struct<(ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[undef:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[alloc_ptr:.*]] = llvm.extractvalue %[[arg_cast]][0] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: %[[align_ptr:.*]] = llvm.extractvalue %[[arg_cast]][1] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: %[[insert0:.*]] = llvm.insertvalue %[[alloc_ptr]], %[[undef:.*]][0] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[insert1:.*]] = llvm.insertvalue %[[align_ptr]], %[[insert0:.*]][1] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)>
+
+  // CHECK: %[[zero0:.*]] = llvm.mlir.constant(0 : index) : i64
+  // CHECK: %[[insert2:.*]] = llvm.insertvalue %[[zero0]], %[[insert1:.*]][2] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)>
+
+  // CHECK: %[[one0:.*]] = llvm.mlir.constant(1 : index) : i64
+  // CHECK: %[[zero1:.*]] = llvm.mlir.constant(0 : index) : i64
+
+  // CHECK: %[[shape_ptr0:.*]] = llvm.extractvalue %[[shape_cast:.*]][1] : !llvm.struct<(ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[shape_gep0:.*]] = llvm.getelementptr %[[shape_ptr0:.*]][%[[zero1:.*]]] : (!llvm.ptr<i64>, i64) -> !llvm.ptr<i64>
+  // CHECK: %[[shape_load0:.*]] = llvm.load %[[shape_gep0:.*]] : !llvm.ptr<i64>
+  // CHECK: %[[insert3:.*]] = llvm.insertvalue %[[shape_load0:.*]], %[[insert2:.*]][3, 0] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[insert4:.*]] = llvm.insertvalue %[[one0:.*]], %[[insert3:.*]][4, 0] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)>
+
+  // CHECK: %[[result_cast:.*]] = builtin.unrealized_conversion_cast %[[insert4:.*]] : !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi32>
+  // CHECK: return %[[result_cast:.*]] : memref<?xi32>
+
+  %1 = memref.reshape %arg0(%shape) : (memref<?x?xi32>, memref<1xindex>) -> memref<?xi32>
+  return %1 : memref<?xi32>
+}