search

tenstorrent / tt-mlir

Tenstorrent MLIR compiler
https://tenstorrent.github.io/tt-mlir/
Apache License 2.0
76 stars 13 forks source link

Assertion Failure in AffineMap due to Missing Dimension in RHS of collapsedLinearAffineMap #1343

Open kamalrajkannan78 opened 2 days ago

kamalrajkannan78 commented 2 days ago

Description

Reproduce

git checkout kkannan/avgpool3d_decompose_mlir_issue
git submodule update --recursive
pytest forge/test/mlir/test_ops.py -k "avgpool3d" -svv

Observed Behaviour

python: /proj_sw/user_dev/kkannan/tt-forge-fe/third_party/tt-mlir/lib/Dialect/TT/IR/TTOpsTypes.cpp:445: mlir::AffineMap collapsedLinearAffineMap(::mlir::MLIRContext *, ::llvm::ArrayRef<int64_t>, ::llvm::ArrayRef<int64_t>, ::llvm::ArrayRef<std::pair<std::int64_t, std::int64_t>>): Assertion `found && "Dim does not participate in AffineMap RHS"' failed.
Fatal Python error: Aborted

Logs

TTIR Graph

module @AvgPool3D attributes {tt.system_desc = #tt.system_desc<[{role = host, target_triple = "x86_64-pc-linux-gnu"}], [{arch = <wormhole_b0>, grid = 8x8, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32, l1_unreserved_base = 1024, erisc_l1_unreserved_base = 1024, dram_unreserved_base = 1024, dram_unreserved_end = 1073741824, physical_cores = {worker = [ 0x0,  0x1,  0x2,  0x3,  0x4,  0x5,  0x6,  0x7,  1x0,  1x1,  1x2,  1x3,  1x4,  1x5,  1x6,  1x7,  2x0,  2x1,  2x2,  2x3,  2x4,  2x5,  2x6,  2x7,  3x0,  3x1,  3x2,  3x3,  3x4,  3x5,  3x6,  3x7,  4x0,  4x1,  4x2,  4x3,  4x4,  4x5,  4x6,  4x7,  5x0,  5x1,  5x2,  5x3,  5x4,  5x5,  5x6,  5x7,  6x0,  6x1,  6x2,  6x3,  6x4,  6x5,  6x6,  6x7,  7x0,  7x1,  7x2,  7x3,  7x4,  7x5,  7x6,  7x7] dram = [ 8x0,  9x0,  10x0,  8x1,  9x1,  10x1,  8x2,  9x2,  10x2,  8x3,  9x3,  10x3]}, supported_data_types = [<f32>, <f16>, <bf16>, <bfp_f8>, <bfp_bf8>, <bfp_f4>, <bfp_bf4>, <bfp_f2>, <bfp_bf2>, <u32>, <u16>, <u8>], supported_tile_sizes = [ 4x16,  16x16,  32x16,  4x32,  16x32,  32x32], num_cbs = 32}], [0], [3 : i32], [ 0x0x0x0]>} {
  func.func @forward(%arg0: tensor<1x1x100x54x54xf32> {ttir.name = "x"}, %arg1: tensor<1x1x0x54x54xf32> {ttir.name = "dc.input_tensor.avg_pool3d_0.0"}, %arg2: tensor<1x54x1x100xf32> {ttir.name = "dc.input_tensor.avg_pool3d_0.dc.index.1.2"}, %arg3: tensor<1x1x1x1xf32> {ttir.name = "dc.input_tensor.avg_pool3d_0.dc.avg_pool2d.3.0"}) -> (tensor<1x1x1x54x54xf32> {ttir.name = "AvgPool3D.output_avg_pool3d_0"}) {
    %0 = tensor.empty() : tensor<1x1x54x100x54xf32>
    %1 = "ttir.transpose"(%arg0, %0) <{dim0 = -2 : si32, dim1 = -3 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x1x100x54x54xf32>, tensor<1x1x54x100x54xf32>) -> tensor<1x1x54x100x54xf32>
    %2 = tensor.empty() : tensor<1x54x100x54xf32>
    %3 = "ttir.squeeze"(%1, %2) <{dim = 0 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x1x54x100x54xf32>, tensor<1x54x100x54xf32>) -> tensor<1x54x100x54xf32>
    %4 = tensor.empty() : tensor<1x54x1x54xf32>
    %5 = "ttir.matmul"(%arg2, %3, %4) <{operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x54x1x100xf32>, tensor<1x54x100x54xf32>, tensor<1x54x1x54xf32>) -> tensor<1x54x1x54xf32>
    %6 = tensor.empty() : tensor<1x1x54x1x54xf32>
    %7 = "ttir.unsqueeze"(%5, %6) <{dim = 0 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x54x1x54xf32>, tensor<1x1x54x1x54xf32>) -> tensor<1x1x54x1x54xf32>
    %8 = tensor.empty() : tensor<1x1x54x54x1xf32>
    %9 = "ttir.transpose"(%7, %8) <{dim0 = -2 : si32, dim1 = -1 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x1x54x1x54xf32>, tensor<1x1x54x54x1xf32>) -> tensor<1x1x54x54x1xf32>
    %10 = tensor.empty() : tensor<1x54x54x1xf32>
    %11 = "ttir.conv2d"(%9, %arg3, %10) <{dilation_height = 1 : si32, dilation_width = 1 : si32, groups = 1 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>], padding_bottom = 0 : si32, padding_left = 0 : si32, padding_right = 0 : si32, padding_top = 0 : si32, stride_height = 1 : si32, stride_width = 1 : si32}> {channel_last = 1 : si32} : (tensor<1x1x54x54x1xf32>, tensor<1x1x1x1xf32>, tensor<1x54x54x1xf32>) -> tensor<1x54x54x1xf32>
    %12 = tensor.empty() : tensor<1x54x1x54xf32>
    %13 = "ttir.transpose"(%11, %12) <{dim0 = -2 : si32, dim1 = -1 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x54x54x1xf32>, tensor<1x54x1x54xf32>) -> tensor<1x54x1x54xf32>
    %14 = tensor.empty() : tensor<1x1x54x54xf32>
    %15 = "ttir.transpose"(%13, %14) <{dim0 = -3 : si32, dim1 = -2 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x54x1x54xf32>, tensor<1x1x54x54xf32>) -> tensor<1x1x54x54xf32>
    %16 = tensor.empty() : tensor<1x1x1x54x54xf32>
    %17 = "ttir.unsqueeze"(%15, %16) <{dim = 2 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x1x54x54xf32>, tensor<1x1x1x54x54xf32>) -> tensor<1x1x1x54x54xf32>
    %18 = tensor.empty() : tensor<1x1x1x54x54xf32>
    %19 = "ttir.concat"(%arg1, %17, %18) <{dim = 2 : si32, operand_constraints = [#tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>, #tt.operand_constraint<dram|l1|scalar|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device|any_device_tile|l1_block_sharded>]}> : (tensor<1x1x0x54x54xf32>, tensor<1x1x1x54x54xf32>, tensor<1x1x1x54x54xf32>) -> tensor<1x1x1x54x54xf32>
    return %19 : tensor<1x1x1x54x54xf32>
  }
}