Hitting Assertion 'succeeded(range) && "element type cannot be iterated"' failed. in Scatter operation

[giordano]

IR:

module {
  func.func @main(%arg0: tensor<1xi64>) -> tensor<1xi64> {
    %0 = mhlo.constant dense<2> : tensor<1xi64>
    %1 = mhlo.constant dense<0> : tensor<1x1xi64>
    %2 = "mhlo.scatter"(%arg0, %1, %0) <{scatter_dimension_numbers = #mhlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>}> ({
    ^bb0(%arg1: tensor<i64>, %arg2: tensor<i64>):
      mhlo.return %arg2 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi64>, tensor<1xi64>) -> tensor<1xi64>
    return %2 : tensor<1xi64>
  }
}

resulting in

external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h:423: auto mlir::DenseElementsAttr::getValues() const [with T = unsigned int]: Assertion `succeeded(range) && "element type cannot be iterated"' failed.

[1679803] signal 6 (-6): Aborted
unknown function (ip: 0xffff8ff46f80)
gsignal at /lib/aarch64-linux-gnu/libc.so.6 (unknown line)
abort at /lib/aarch64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0xffff8feef8a3)
__assert_fail at /lib/aarch64-linux-gnu/libc.so.6 (unknown line)
getValues<unsigned int> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h:423
value_begin<unsigned int> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h:430
getSplatValue<unsigned int> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h:390
fold at /proc/self/cwd/external/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc:6099
foldHook<mlir::mhlo::ScatterOp> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/OpDefinition.h:1932
operator() at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/OpDefinition.h:1890
CallImpl<const mlir::Op<mlir::mhlo::ScatterOp, mlir::OpTrait::OneRegion, mlir::OpTrait::VariadicResults, mlir::OpTrait::ZeroSuccessors, mlir::OpTrait::AtLeastNOperands<1>::Impl, mlir::OpTrait::OpInvariants, mlir::BytecodeOpInterface::Trait, mlir::OpTrait::HasRecursiveMemoryEffects, mlir::InferTypeOpInterface::Trait>::getFoldHookFn()::<lambda(mlir::Operation*, llvm::ArrayRef<mlir::Attribute>, llvm::SmallVectorImpl<mlir::OpFoldResult>&)> > at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/FunctionExtras.h:222
operator() at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/FunctionExtras.h:413
foldHook at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/OperationSupport.h:536
foldHook at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/OperationSupport.h:263
fold at /proc/self/cwd/external/llvm-project/mlir/lib/IR/Operation.cpp:636
fold at /proc/self/cwd/external/llvm-project/mlir/lib/IR/Operation.cpp:666
processWorklist at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:493
operator() at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:874
callback_fn<(anonymous namespace)::RegionPatternRewriteDriver::simplify(bool*) &&::<lambda()> > at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46
operator() at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69
executeAction<(anonymous namespace)::GreedyPatternRewriteIteration, long int&> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/MLIRContext.h:280
simplify at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:872
applyPatternsGreedily at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:919
applyPatternsGreedily at /proc/self/cwd/external/llvm-project/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h:174
runOnOperation at /proc/self/cwd/external/llvm-project/mlir/lib/Transforms/Canonicalizer.cpp:63
operator() at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:526
callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::<lambda()> > at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46
operator() at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69
executeAction<mlir::PassExecutionAction, mlir::Pass&> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/MLIRContext.h:280
run at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:520
runPipeline at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:592
runOnOperationImpl at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:733
runOnOperation at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:715
operator() at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:524
callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::<lambda()> > at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46
operator() at /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69
executeAction<mlir::PassExecutionAction, mlir::Pass&> at /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/MLIRContext.h:280
run at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:520
runPipeline at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:592
runPasses at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:905
run at /proc/self/cwd/external/llvm-project/mlir/lib/Pass/Pass.cpp:885
MlirToXlaComputation at /proc/self/cwd/external/xla/xla/pjrt/mlir_to_hlo.cc:96
Compile at /proc/self/cwd/external/xla/xla/pjrt/cpu/cpu_client.cc:783
[...]

Maybe the issue is at

xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc

Line 6099 in c2a9a2d

index.getSplatValue<uint32_t>() == 0 &&

where index is expected to be 32-bit, but the spec suggests it should be 64-bit.

CC: @wsmoses.