dynamic_reshape and dynamic_broadcast_in_dim can't be lowered to a dialect upstream of mlir.

[linuxlonelyeagle]

We seem to be in a bit of a bind. There is stablehlo as a backend in torch-mlir, and chlo-legalize-to-stablehlo did solve some problems before, but there are still some remaining issues.
Give an example of e2e.HardswishModule_basic
Here are all the pipelines it runs.

builtin.module(torchscript-module-to-torch-backend-pipeline{backend-legal-ops= extra-library=})
builtin.module(torch-backend-to-stablehlo-backend-pipeline)
builtin.module(func.func(chlo-legalize-to-stablehlo),canonicalize,stablehlo-legalize-to-linalg)
builtin.module(func.func(refback-generalize-tensor-pad),func.func(refback-generalize-tensor-concat),func.func(linalg-generalize-named-ops),func.func(linalg-fuse-elementwise-ops),convert-shape-to-std,sparse-assembler,sparsification-and-bufferization,sparse-storage-specifier-to-llvm,inline,func.func(scf-bufferize),func.func(tm-tensor-bufferize),func.func(empty-tensor-to-alloc-tensor),func.func(linalg-bufferize),func-bufferize,arith-bufferize,refback-mlprogram-bufferize,func.func(tensor-bufferize),func.func(finalizing-bufferize),func.func(buffer-deallocation),refback-munge-calling-conventions,func.func(tm-tensor-to-loops),func.func(refback-munge-memref-copy),func.func(convert-linalg-to-loops),func.func(lower-affine),convert-scf-to-cf,func.func(refback-expand-ops-for-llvm),func.func(arith-expand),func.func(convert-math-to-llvm),convert-math-to-libm,expand-strided-metadata,finalize-memref-to-llvm,lower-affine,convert-bufferization-to-memref,finalize-memref-to-llvm,func.func(convert-arith-to-llvm),convert-func-to-llvm,convert-cf-to-llvm,convert-complex-to-llvm,reconcile-unrealized-casts)

      %22 = bufferization.to_tensor %alloc_28 : memref<2xindex>
      %23 = stablehlo.dynamic_broadcast_in_dim %18, %22, dims = [0, 1] : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
      %24 = bufferization.to_memref %23 : memref<?x?xf32>
      %25 = bufferization.to_memref %23 : memref<?x?xf32>
      %26 = bufferization.to_memref %23 : memref<?x?xf32>
      %27 = stablehlo.dynamic_broadcast_in_dim %0, %22, dims = [0, 1] : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
      %28 = bufferization.to_memref %27 : memref<?x?xf32>
      %dim_29 = memref.dim %25, %c0 : memref<?x?xf32>
      %dim_30 = memref.dim %24, %c1 : memref<?x?xf32>
      %alloc_31 = memref.alloc(%dim_29, %dim_30) {alignment = 64 : i64} : memref<?x?xf32>
      linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%26, %28 : memref<?x?xf32>, memref<?x?xf32>) outs(%alloc_31 : memref<?x?xf32>) {
      ^bb0(%in: f32, %in_32: f32, %out: f32):
        %29 = arith.mulf %in, %in_32 : f32
        linalg.yield %29 : f32
      }
      shape.assuming_yield %alloc_31 : memref<?x?xf32>
    }
    return %21 : memref<?x?xf32>
  }

What you can notice is that dynamic_broadcast_in_dim is still there. Not only that, but the same situation with dynamic_reshape also has.So I was wondering, can I add a new pass to stablehlo for the two operations of lower?Or if people have other methods. Or add dynamic_reshape and stablehlo.dynamic_broadcast_in_dim to the interpreter?
I hope someone can help me. Thanks.

[stellaraccident]

If coming from the old torch-mlir TorchScript path, then those ops are fairly hard to lower because the frontend makes insufficient guarantees regarding dimension relationships, and it forces the backend to handle a lot of combinations defensively. This is usually pretty bad because it is almost always extremely important to fuse those ops, or performance will be very bad.

This path isn't being developed anymore. Current work is focused on the FX path, which has much stronger semantics, suitable for dynamic shape compilers: https://github.com/llvm/torch-mlir/blob/main/docs/roadmap.md#current-api-paths

Compilers using that path are assuming the stronger semantics, particularly around the program forms that yield these ops.

I don't think stablehlo ever defined its dynamic shape system in a similar fashion, so I'm not sure how to write good lowerings with the existing stablehlo definitions, but I'd advise finding out if you can assume the stronger constraints practically. Doing so will most likely cause all such ops to trivially fuse or be directly representable by linalg.