@@ -103,8 +103,14 @@ class CinnJitInstruction::FnPtrImpl {
103
103
104
104
// Pass real tensor data to cinn_buffer_t func args placeholder
105
105
for (size_t i = 0 ; i < kernel_tensor_args.size (); ++i) {
106
- cinn_pod_value_to_buffer_p (&(func_args_[i]))->memory =
107
- reinterpret_cast <uint8_t *>(kernel_tensor_args[i]->data ());
106
+ if (!kernel_tensor_args[i]->has_allocation ()) {
107
+ VLOG (2 ) << " WARNING! Access DenseTensor::data() without allocation, "
108
+ " return nullptr!" ;
109
+ cinn_pod_value_to_buffer_p (&(func_args_[i]))->memory = nullptr ;
110
+ } else {
111
+ cinn_pod_value_to_buffer_p (&(func_args_[i]))->memory =
112
+ reinterpret_cast <uint8_t *>(kernel_tensor_args[i]->data ());
113
+ }
108
114
}
109
115
110
116
// Launch host kernel
@@ -297,6 +303,7 @@ CinnJitInstruction::CinnJitInstruction(
297
303
ir_dims_.push_back (
298
304
result.type ().dyn_cast <paddle::dialect::DenseTensorType>().dims ());
299
305
tensor_args_.push_back (tensor);
306
+ alloc_tensors_.push_back (tensor);
300
307
auto alloc_tensor_type =
301
308
result.type ().dyn_cast <paddle::dialect::AllocatedDenseTensorType>();
302
309
tensor->set_type (
@@ -321,6 +328,7 @@ CinnJitInstruction::CinnJitInstruction(
321
328
}
322
329
for (auto & tensor : temp_space_tensors_) {
323
330
tensor_args_.push_back (&tensor);
331
+ alloc_tensors_.push_back (&tensor);
324
332
}
325
333
output_tensor_size += temp_space_tensors_.size ();
326
334
}
@@ -343,8 +351,8 @@ void CinnJitInstruction::Run() {
343
351
fn_ptr_impl_->InferShape (
344
352
tensor_args_, ir_dims_, input_tensor_size, output_tensor_size);
345
353
}
346
- for (size_t i = 0 ; i < tensor_args_ .size (); ++i) {
347
- dev_ctx_->Alloc (tensor_args_ [i], tensor_args_ [i]->dtype ());
354
+ for (size_t i = 0 ; i < alloc_tensors_ .size (); ++i) {
355
+ dev_ctx_->Alloc (alloc_tensors_ [i], alloc_tensors_ [i]->dtype ());
348
356
}
349
357
350
358
// 2. execute kernel
0 commit comments