Giters

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.

Home Page:http://iree.dev/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Wrong `flow.dispatch.tensor.store` op creation; leading to shape mismatch.

pashu123 opened this issue · comments

commented

What happened?

above.mlir:14:12: error: 'iree_linalg_ext.scan' op expected type of operand #1 ('tensor<1x8xi32>') to match type of corresponding result ('tensor<1x?xi32>')
    %8:2 = iree_linalg_ext.scan {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 0]]>} dimension(1) inclusive(true) ins(%3 : tensor<1x8xi32>) outs(%6, %7 : tensor<1x8xi32>, tensor<1xi32>) {
           ^
above.mlir:14:12: note: see current operation:
%12:2 = "iree_linalg_ext.scan"(%7, %9, %11) <{dimension = 1 : i64, inclusive = true, operandSegmentSizes = array<i32: 1, 2>}> ({
^bb0(%arg0: i32, %arg1: i32):
  %13 = "arith.addi"(%arg0, %arg1) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
  "iree_linalg_ext.yield"(%13) : (i32) -> ()
}) {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 0]]>} : (tensor<1x?xi32>, tensor<1x8xi32>, tensor<1xi32>) -> (tensor<1x?xi32>, tensor<1xi32>)

Steps to reproduce your issue

module {
  func.func @jit_eval_dispatch_0_scan_1x8xi32() attributes {translation_info = #iree_codegen.translation_info<CPUDefault>} {
    %c0_i32 = arith.constant 0 : i32
    %c0 = arith.constant 0 : index
    %c64 = arith.constant 64 : index
    %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<1x8xi32>>
    %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<1x8xi32>>
    %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c64) : !flow.dispatch.tensor<writeonly:tensor<1xi32>>
    %3 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [1, 8], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<1x8xi32>> -> tensor<1x8xi32>
    %4 = tensor.empty() : tensor<1xi32>
    %5 = tensor.empty() : tensor<1x8xi32>
    %6 = linalg.fill {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 0], [0, 0], [0, 0], [0, 0]]>} ins(%c0_i32 : i32) outs(%5 : tensor<1x8xi32>) -> tensor<1x8xi32>
    %7 = linalg.fill {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1], [0], [0], [0]]>} ins(%c0_i32 : i32) outs(%4 : tensor<1xi32>) -> tensor<1xi32>
    %8:2 = iree_linalg_ext.scan {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 0]]>} dimension(1) inclusive(true) ins(%3 : tensor<1x8xi32>) outs(%6, %7 : tensor<1x8xi32>, tensor<1xi32>) {
    ^bb0(%arg0: i32, %arg1: i32):
      %9 = arith.addi %arg0, %arg1 : i32
      iree_linalg_ext.yield %9 : i32
    } -> tensor<1x8xi32>, tensor<1xi32>
    flow.dispatch.tensor.store %8#0, %1, offsets = [0, 0], sizes = [1, 8], strides = [1, 1] : tensor<1x8xi32> -> !flow.dispatch.tensor<writeonly:tensor<1x8xi32>>
    flow.dispatch.tensor.store %8#1, %2, offsets = [0], sizes = [1], strides = [1] : tensor<1xi32> -> !flow.dispatch.tensor<writeonly:tensor<1xi32>>
    return
  }
}

RUN: iree-opt --pass-pipeline='builtin.module(func.func(iree-codegen-tile-and-distribute-to-workgroups))' above.mlir

What component(s) does this issue relate to?

No response

Version information

No response

Additional context

No response

commented

@MaheshRavishankar Any suggestion? I will look more into the creation buggy dispatch tensor store op.