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

[rengolin]

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/end-to-end/KernelBench/test_kernel_bench.py --test 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.