Assert when lowering 1D Convolution from StableHLO to LinAlg

[pstarkcdpr]

What happened?

I found this when reporting the following IREE issue: iree-org/iree#21230 (comment)

When compiling a model in StableHLO using IREE with iree-compile in Debug mode, the following assert triggers.

~/iree-test$ iree-org/iree-compiler/tools/iree-compile --iree-hal-target-device=vulkan --iree-vulkan-target=sm_75 -o model1.vmfb model1.mlir
iree-compile: /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/IR/OperationSupport.h:1082: void mlir::OperationState::addAttribute(mlir::StringAttr, mlir::Attribute): Assertion `attr && "attribute cannot be null"' failed.
Please report issues to https://github.com/iree-org/iree/issues and include the crash backtrace.
Stack dump:
0.      Program arguments: iree-org/iree-compiler/tools/iree-compile --iree-hal-target-device=vulkan --iree-vulkan-target=sm_75 -o model1.vmfb model1.mlir
 #0 0x0000742db96831e0 ___interceptor_backtrace.part.0 /build/gcc-14-ig5ci0/gcc-14-14.2.0/build/x86_64-linux-gnu/libsanitizer/asan/../../../../src/libsanitizer/sanitizer_common/sanitizer_common_interceptors.inc:4358:28
 #1 0x0000742d746eca69 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:804:22
 #2 0x0000742d746edb50 PrintStackTraceSignalHandler(void*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:888:1
 #3 0x0000742d746e4d40 llvm::sys::RunSignalHandlers() /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/lib/Support/Signals.cpp:105:20
 #4 0x0000742d746eb07a SignalHandler(int, siginfo_t*, void*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:418:13
 #5 0x0000742d68645330 (/lib/x86_64-linux-gnu/libc.so.6+0x45330)
 #6 0x0000742d6869eb2c __pthread_kill_implementation ./nptl/pthread_kill.c:44:76
 #7 0x0000742d6869eb2c __pthread_kill_internal ./nptl/pthread_kill.c:78:10
 #8 0x0000742d6869eb2c pthread_kill ./nptl/pthread_kill.c:89:10
 #9 0x0000742d6864527e raise ./signal/../sysdeps/posix/raise.c:27:6
#10 0x0000742d686288ff abort ./stdlib/abort.c:81:7
#11 0x0000742d6862881b _nl_load_domain ./intl/loadmsgcat.c:1177:9
#12 0x0000742d6863b517 (/lib/x86_64-linux-gnu/libc.so.6+0x3b517)
#13 0x0000742d74205cb0 mlir::OperationState::addAttribute(mlir::StringAttr, mlir::Attribute) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/IR/OperationSupport.h:1083:22
#14 0x0000742d791e37e5 mlir::OperationState::addAttribute(llvm::StringRef, mlir::Attribute) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/IR/OperationSupport.h:1076:3
#15 0x0000742d8fd6290e mlir::linalg::Conv1DNwcWcfOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::TypeRange, mlir::ValueRange, mlir::ValueRange, mlir::Attribute, mlir::Attribute, llvm::ArrayRef<mlir::NamedAttribute>) /home/paulstark/iree-test/iree-org/iree-compiler/llvm-project/tools/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc:4413:26
#16 0x0000742d784c2ca8 mlir::linalg::Conv1DNwcWcfOp mlir::OpBuilder::create<mlir::linalg::Conv1DNwcWcfOp, mlir::ShapedType&, mlir::ValueRange, mlir::ValueRange, mlir::Attribute&, mlir::Attribute&, llvm::SmallVector<mlir::NamedAttribute, 3u>>(mlir::Location, mlir::ShapedType&, mlir::ValueRange&&, mlir::ValueRange&&, mlir::Attribute&, mlir::Attribute&, llvm::SmallVector<mlir::NamedAttribute, 3u>&&) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/IR/Builders.h:506:16
#17 0x0000742d784aae0d mlir::stablehlo::(anonymous namespace)::NormalConvolutionOpConversion::matchAndRewrite(mlir::stablehlo::ConvolutionOp, mlir::stablehlo::ConvolutionOpAdaptor, mlir::ConversionPatternRewriter&) const /home/paulstark/iree-test/iree-org/iree/third_party/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgConvolution.cpp:269:13
#18 0x0000742d773081f3 mlir::OpConversionPattern<mlir::stablehlo::ConvolutionOp>::matchAndRewrite(mlir::stablehlo::ConvolutionOp, mlir::stablehlo::ConvolutionOpGenericAdaptor<llvm::ArrayRef<mlir::ValueRange>>, mlir::ConversionPatternRewriter&) const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h:624:27
#19 0x0000742d77307c12 mlir::OpConversionPattern<mlir::stablehlo::ConvolutionOp>::matchAndRewrite(mlir::Operation*, llvm::ArrayRef<mlir::ValueRange>, mlir::ConversionPatternRewriter&) const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h:608:27
#20 0x0000742d918ad131 mlir::ConversionPattern::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:1886:25
#21 0x0000742d919f7322 mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>)::'lambda'()::operator()() const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Rewrite/PatternApplicator.cpp:212:46
#22 0x0000742d919f91a7 void llvm::function_ref<void ()>::callback_fn<mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>)::'lambda'()>(long) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:47:40
#23 0x0000742d74491fe0 llvm::function_ref<void ()>::operator()() const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:62
#24 0x0000742d91a00a93 void mlir::MLIRContext::executeAction<mlir::ApplyPatternAction, mlir::Pattern const&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pattern const&) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/IR/MLIRContext.h:281:3
#25 0x0000742d919f8cab mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Rewrite/PatternApplicator.cpp:196:9
#26 0x0000742d918b0f49 (anonymous namespace)::OperationLegalizer::legalizeWithPattern(mlir::Operation*, mlir::ConversionPatternRewriter&) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:2197:36
#27 0x0000742d918ae453 (anonymous namespace)::OperationLegalizer::legalize(mlir::Operation*, mlir::ConversionPatternRewriter&) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:2086:16
#28 0x0000742d918b6dcf mlir::OperationConverter::convert(mlir::ConversionPatternRewriter&, mlir::Operation*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:2596:13
#29 0x0000742d918b9321 mlir::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation*>) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:2696:15
#30 0x0000742d918c69d2 mlir::applyPartialConversion(llvm::ArrayRef<mlir::Operation*>, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, mlir::ConversionConfig) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:3388:39
#31 0x0000742d918c6c21 mlir::applyPartialConversion(mlir::Operation*, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, mlir::ConversionConfig) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:3394:32
#32 0x0000742d76f80e39 mlir::iree_compiler::stablehlo::(anonymous namespace)::ConvertStableHloToIreeInputDialects::runOnOperation() /home/paulstark/iree-test/iree-org/iree/compiler/plugins/input/StableHLO/Conversion/StableHLOToIREEInputDialects.cpp:615:15
#33 0x0000742d75711a2a mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda'()::operator()() const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:529:22
#34 0x0000742d7571decc void llvm::function_ref<void ()>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda'()>(long) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:47:40
#35 0x0000742d74491fe0 llvm::function_ref<void ()>::operator()() const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:62
#36 0x0000742d757301e7 void mlir::MLIRContext::executeAction<mlir::PassExecutionAction, mlir::Pass&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pass&) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/IR/MLIRContext.h:281:3
#37 0x0000742d757128e0 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:523:7
#38 0x0000742d757130ed mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:594:15
#39 0x0000742d757117bb mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda'(mlir::OpPassManager&, mlir::Operation*)::operator()(mlir::OpPassManager&, mlir::Operation*) const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:511:42
#40 0x0000742d757255e7 llvm::LogicalResult llvm::function_ref<llvm::LogicalResult (mlir::OpPassManager&, mlir::Operation*)>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda'(mlir::OpPassManager&, mlir::Operation*)>(long, mlir::OpPassManager&, mlir::Operation*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:47:40
#41 0x0000742d74fe1a82 llvm::function_ref<llvm::LogicalResult (mlir::OpPassManager&, mlir::Operation*)>::operator()(mlir::OpPassManager&, mlir::Operation*) const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:62
#42 0x0000742d74fdf35f mlir::Pass::runPipeline(mlir::OpPassManager&, mlir::Operation*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/Pass/Pass.h:201:3
#43 0x0000742d7db4b342 mlir::iree_compiler::InputConversion::(anonymous namespace)::AutoInputConversionPipelinePass::runOnOperation() /home/paulstark/iree-test/iree-org/iree/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp:72:13
#44 0x0000742d75711a2a mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda'()::operator()() const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:529:22
#45 0x0000742d7571decc void llvm::function_ref<void ()>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda'()>(long) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:47:40
#46 0x0000742d74491fe0 llvm::function_ref<void ()>::operator()() const /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:62
#47 0x0000742d757301e7 void mlir::MLIRContext::executeAction<mlir::PassExecutionAction, mlir::Pass&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pass&) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/include/mlir/IR/MLIRContext.h:281:3
#48 0x0000742d757128e0 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:523:7
#49 0x0000742d757130ed mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:594:15
#50 0x0000742d757197e4 mlir::PassManager::runPasses(mlir::Operation*, mlir::AnalysisManager) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:907:40
#51 0x0000742d7571930a mlir::PassManager::run(mlir::Operation*) /home/paulstark/iree-test/iree-org/iree/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:887:69
#52 0x0000742d742c2ac5 mlir::iree_compiler::embed::(anonymous namespace)::Invocation::runPipeline(iree_compiler_pipeline_t) /home/paulstark/iree-test/iree-org/iree/compiler/src/iree/compiler/API/Internal/CompilerDriver.cpp:1034:0
#53 0x0000742d742c72c8 ireeCompilerInvocationPipeline /home/paulstark/iree-test/iree-org/iree/compiler/src/iree/compiler/API/Internal/CompilerDriver.cpp:1474:0
#54 0x0000742d7512cfb2 mlir::iree_compiler::runIreecMain(int, char**)::'lambda'(iree_compiler_source_t*)::operator()(iree_compiler_source_t*) const /home/paulstark/iree-test/iree-org/iree/compiler/src/iree/compiler/Tools/iree_compile_lib.cc:254:11
#55 0x0000742d75132e69 mlir::iree_compiler::runIreecMain(int, char**) /home/paulstark/iree-test/iree-org/iree/compiler/src/iree/compiler/Tools/iree_compile_lib.cc:355:9
#56 0x0000742d74370039 ireeCompilerRunMain /home/paulstark/iree-test/iree-org/iree/compiler/src/iree/compiler/API/Internal/IREECompileToolEntryPoint.cpp:13:1
#57 0x00005dcc51aad1ad main /home/paulstark/iree-test/iree-org/iree/tools/iree-compile-main.cc:9:72
#58 0x0000742d6862a1ca __libc_start_call_main ./csu/../sysdeps/nptl/libc_start_call_main.h:74:3
#59 0x0000742d6862a28b call_init ./csu/../csu/libc-start.c:128:20
#60 0x0000742d6862a28b __libc_start_main ./csu/../csu/libc-start.c:347:5
#61 0x00005dcc51aad0c5 _start (iree-org/iree-compiler/tools/iree-compile+0x10c5)
Aborted (core dumped)

While this is 100% reproducible, I am not at liberty to share the models. I'll look into making a minimal repro case, but if there is other information I can add to help track this down, please let me know.

Steps to reproduce your issue

1. Run iree-compile --iree-hal-target-device=vulkan --iree-vulkan-target=sm_75 -o model1.vmfb model1.mlir in Debug mode.

Version information

I tested with IREE 3.5.0.

[pstarkcdpr]

Here's a simple repro

module @conv_assert {
  func.func public @main(%arg0: tensor<50x112x15xf32>, %arg1: tensor<96x112x1xf32>) -> tensor<50x96x15xf32> {
    %318 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, f, 0]x[o, i, 0]->[b, f, 0], window = {} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<50x112x15xf32>, tensor<96x112x1xf32>) -> tensor<50x96x15xf32>
    return %318 : tensor<50x96x15xf32>
  }
}

What's happening here is simply that mlir::stablehlo::NormalConvolutionOpConversion::matchAndRewrite() generates the underlying linalg convolutions, for example

        res = rewriter.create<linalg::Conv1DNwcWcfOp>(
            loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor},
            strides, dilations, linalg::getPrunedAttributeList(op))

The problem is that strides and dilations can be null attributes, which triggers an assert. Everything works fine in Release though because there's an default implementation of these, so it's fine for them to be null. The create call doesn't seem to know that, however, so causes a spurious assert.

I tried to fix it by not passing the attributes in if they're null, but I couldn't find a good way to do that. There's a constructor that just takes a list of attributes, so I could just go

        auto attrs = linalg::getPrunedAttributeList(op);
        if (strides) {
          attrs.emplace_back("strides", strides);
        }
        if (dilations) {
          attrs.emplace_back("dilations", dilations);
        }
        res = rewriter.create<linalg::Conv1DNwcWcfOp>(
            loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor},
            attrs);

but I don't want to hardcode the attribute names. There appear to accessors to get these names, like getDilationsAttrName(), but I couldn't figure out how to get those statically without creating the class first.

I wonder if someone who understands this code better can chime in with the right way to fix this assert. It seems to be harmless, but it gets in the way of debugging other issues, so it would be nice to fix.

[jpienaar]

I'd suggest just checking before that create is called and instead call the one that doesn't have them as args, e.g.,

static Conv1DNwcWcfOp create(::mlir::OpBuilder& builder,
                             ::mlir::Location location,
                             TypeRange resultTensorTypes, ValueRange operands,
                             ArrayRef<NamedAttribute> attributes = {});

The alternative would be to teach ODS that when it generates these, that nulls are allowed.

[pstarkcdpr]

Thanks @jpienaar . That was my plan, just add them straight to attributes if they're not null. Is it okay for me to use constant strings, like

attrs.emplace_back("dilations", dilations);

The string "dilations" seems to be available as a static constant, but I could find a way to access it. The generated code for Conv1DNwcWcfOp goes

  static ::llvm::ArrayRef<::llvm::StringRef> getAttributeNames() {
    static ::llvm::StringRef attrNames[] = {::llvm::StringRef("dilations"), ::llvm::StringRef("strides"), ::llvm::StringRef("operandSegmentSizes")};
    return ::llvm::ArrayRef(attrNames);
  }

  ::mlir::StringAttr getDilationsAttrName() {
    return getAttributeNameForIndex(0);
  }

  static ::mlir::StringAttr getDilationsAttrName(::mlir::OperationName name) {
    return getAttributeNameForIndex(name, 0);
  }

The static version of getDilationsAttrName() needs a parameter for some reason. I'm fine with just using constant strings though if that's okay with you.

[jpienaar]

The static version of getDilationsAttrName() needs a parameter for some reason.

Yes, that is as it needs the MLIRContext at least to make an attr.

I was thinking something more basic

if (strides & dilations)
  res = rewriter.create<linalg::Conv1DNwcWcfOp>(
            loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor},
            strides, dilations, linalg::getPrunedAttributeList(op));
else
  res = rewriter.create<linalg::Conv1DNwcWcfOp>(
            loc, resultType, ValueRange{input, filter}, ValueRange{zeroTensor},
            linalg::getPrunedAttributeList(op));

But indeed you can do the else always and then just append if set. Alternatively you can look at the generated Conv1DNwcWcfOp::getStrides() and dilation pair, and then replace strides & dilations and always call what's in the if.

I think a change to ODS would be the nicest answer here, so changing upstream to allow it to be empty. One part is a little bit unfortunate, and that's that optional ones aren't passed in as std::optional (which could be changed too of course, I think the implicit constructor may even make that a non-breaking change).

[pstarkcdpr]

Oh, I see. Your simple solution works fine for me. I had been thinking of the case where, for example, there were strides but not dilations, but that's not the case for us in practice. I agree that changing the code gen would be the nicest solution, but this is good enough for my purposes. Would you like me to make a PR with this change?