llvm / llvm/torch-mlir

Dialect `tm_tensor` not found for custom op 'tm_tensor.scan'

Open
#4,574 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
1.9k
Forks
736
Avg merge
5d 22h
Merged PRs (30d)
15

Description

When lowering some KernelBench programs (level1, *cumsum*), I get this error:

Dialect 'tm_tensor' not found for custom op 'tm_tensor.scan'

Looking at the torch-mlir repo, I realized this is specific to this code. While it's ok to have repo-specific dialects in torch-mlir, it should never emit MLIR code with it in the output, only upstream dialects. This is requesting linalg-on-tensors output.

Reproducer:

  1. Clone Lighthouse
  2. Init uv as in README
  3. Run: $ uv run examples/KernelBench/test-kernel-bench.py --kernel level1/89_

Full error:

mlir._mlir_libs._mlir.ir.MLIRError: Unable to parse module assembly:
error: "-":8:12: Dialect `tm_tensor' not found for custom op 'tm_tensor.scan' 
 note: "-":8:12: Available dialects: acc, affine, amdgpu, arith, arm_neon, arm_sme, arm_sve, async, bufferization, builtin, cf, complex, dlti, emitc, func, gpu, index, irdl, linalg, llvm, math, memref, ml_program, mpi, nvgpu, nvvm, omp, pdl, pdl_interp, ptr, quant, rocdl, scf, shape, shard, smt, sparse_tensor, spirv, tensor, tosa, transform, transform_ext (*), transform_smt_ext (*), ub, vcix, vector, wasmssa, x86, xegpu, xevm (* corresponding to loaded but unregistered dialects); for more info on dialect registration see https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management

Full generated file:

module {
  func.func @main(%arg0: tensor<32768x32768xf32>) -> tensor<32768x32768xf32> {
    %cst = arith.constant 0.000000e+00 : f32
    %0 = tensor.empty() : tensor<32768x32768xf32>
    %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<32768x32768xf32>) -> tensor<32768x32768xf32>
    %2 = tensor.empty() : tensor<32768xf32>
    %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<32768xf32>) -> tensor<32768xf32>
    %4:2 = tm_tensor.scan dimension(1) inclusive(true) ins(%arg0 : tensor<32768x32768xf32>) outs(%1, %3 : tensor<32768x32768xf32>, tensor<32768xf32>) {
    ^bb0(%arg1: f32, %arg2: f32):
      %5 = arith.addf %arg1, %arg2 : f32
      tm_tensor.yield %5 : f32
    } -> tensor<32768x32768xf32>, tensor<32768xf32>
    return %4#0 : tensor<32768x32768xf32>
  }
}

Expected output:

  • Only linalg operations.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the failure with examples/KernelBench/test-kernel-bench.py and level1/89_ as described. Trace the lowering that produces tm_tensor.scan and compare the generated module with the requested linalg-on-tensors output; done means the output contains only linalg operations and parses with the available dialects.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.