diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc index c4aed1a237dd..57ee62e152db 100644 --- a/src/target/llvm/codegen_cpu.cc +++ b/src/target/llvm/codegen_cpu.cc @@ -642,7 +642,8 @@ CodeGenLLVM::TypedPointer CodeGenCPU::PackClosureData(const Array& vfields, } llvm::StructType* ctype = struct_name.size() ? llvm::StructType::create(fields, struct_name) : llvm::StructType::create(fields); - llvm::Value* cvalue = builder_->CreateAlloca(ctype, ConstInt32(1)); + llvm::AllocaInst* cvalue = + WithFunctionEntry([&]() { return builder_->CreateAlloca(ctype, ConstInt32(1)); }); llvm::Value* zero = ConstInt32(0); for (size_t i = 0; i < vfields.size(); ++i) { builder_->CreateStore(var_map_.at(vfields[i].get()), diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index b641ba1fdb13..2fe40ae2f88e 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -1909,6 +1909,26 @@ def test_cumprod(target, dev, executor_kind): @tvm.testing.parametrize_targets def test_scatter_nd(target, dev, executor_kind): + def test_scatter_nd_large_shape(): + def before(): + data = relay.const(np.zeros((1, 900, 300), dtype="float32"), dtype="float32") + indices = relay.const(np.ones((3, 1, 900, 300), dtype="int64"), dtype="int64") + update = relay.const(np.ones((1, 900, 300), dtype="float32"), dtype="float32") + b = relay.op.scatter_nd(data, indices, update) + return relay.Function(relay.analysis.free_vars(b), b) + + passes = tvm.transform.Sequential( + [ + relay.transform.InferType(), + relay.transform.FoldConstant(), + ] + ) + before_mod = tvm.IRModule.from_expr(before()) + with tvm.transform.PassContext(opt_level=3): + after_mod = passes(before_mod) + + test_scatter_nd_large_shape() + def verify_scatter_nd( data_np, indices_np, updates_np, ref_res, mode="add", rtol=1e-5, atol=1e-5 ):