iree: linalg.init_tensor with dynamic shapes and shapex dialect not working together

Reproduce with MLIR file here: https://gist.github.com/silvasean/683adcb9daa95de3700f999777fd61ea

Reproducible with iree-opt --iree-convert-to-hal repro.mlir

Assigning to Ben initially, as this looks related to the buffer stuff he was recently redesigning.

This IR is the result of lowering this IR from npcomp (for the IREE analog of https://github.com/llvm/mlir-npcomp/pull/205), which looks correct to me (not an invalid IREE input):

module  {                                                                                       
  func @forward(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> attributes {iree.module.export} {                                                                                              
    %c0 = constant 0 : index                                                                    
    %c1 = constant 1 : index                                                                                                                                                                                      
    %cst = constant 0.000000e+00 : f32                                                          
    %0 = memref.dim %arg0, %c0 : tensor<?x?xf32>                                                                                                                                                                                                              
    %3 = memref.dim %arg1, %c1 : tensor<?x?xf32>                                                                     
    %5 = linalg.init_tensor [%0, %3] : tensor<?x?xf32>                                                                              
    %6 = linalg.fill(%5, %cst) : tensor<?x?xf32>, f32 -> tensor<?x?xf32>                                                 
    %7 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%6 : tensor<?x?xf32>) -> tensor<?x?xf32>                                                                                         
    return %7 : tensor<?x?xf32>                                                                                                                              
  }                                                                                                                              
}

(the original IR from npcomp has an std.assert op checking the contracting dimensions match, but I removed that here as part of reducing the bug). (the origina l

About this issue

Original URL
State: closed
Created 3 years ago
Comments: 37 (25 by maintainers)

Most upvoted comments

FWIW - finally popped enough things off the stack to work directly on the shape related items.

Removing the following from Flow/Transforms/Passes.cpp allows your sample to compile:

  // SmallVector<std::string> doNotRecurseOpNames = {"flow.dispatch.workgroups"};
  // passManager.addNestedPass<FuncOp>(
  //     Shape::createTieDynamicShapesPass(doNotRecurseOpNames));
  // passManager.addNestedPass<FuncOp>(
  //     Shape::createMaterializeShapeCalculationsPass());
  // passManager.addNestedPass<FuncOp>(Shape::createHoistShapeCalculationsPass());

However, I can’t just delete those yet and need to patch the other input pipelines to have better/more modern dynamic shape support. I’d be open to adding a flag that disables those passes for the more pure-path you are running up against. It will probably be ~1-2 weeks for me to re-engineering everything else.

stellaraccident on May 14, 2021

We should stop relying on shapex. We may still want portions of it in the short term as a crutch for MHLO, but it shouldn’t be getting to IREE proper.

As part of my ABI work, I did get a very simple dynamic case compiled (not the one Sean is doing) and convinced myself that the lower level mechanisms are at least somewhat working. My priority right now is stripping out/sequestering the legacy stuff so this is easier to manipulate. That includes shapex.

stellaraccident on May 3, 2021

That snippet is already too late - the linalg op that is at the very bottom there cannot exist: https://gist.github.com/silvasean/683adcb9daa95de3700f999777fd61ea#file-repro-mlir-L614

So we need the original input IR prior to any IREE flow work (what you have inlined)

benvanik on Apr 12, 2021