【npu】fix cast and arange for npu (#1630)

Author: xiaoguoguo626807

backends/npu/kernels/arange_kernel.cc

@@ -114,9 +114,6 @@ void AclopArangeKernel(const Context& dev_ctx,
   int64_t size = 0;
   GetSize(start_value, end_value, step_value, &size);
 
-  out->Resize(phi::make_ddim({size}));
-  dev_ctx.template Alloc<T>(out);
-
   std::vector<T> odata;
   T value = start_value;
   for (int64_t i = 0; i < size; ++i) {
@@ -136,7 +133,16 @@ void ArangeKernel(const Context& dev_ctx,
   DO_COMPATIBILITY(aclnnArange,
                    (custom_kernel::AclopArangeKernel<T, Context>(
                        dev_ctx, start, end, step, out)));
+  T start_value = start.to<T>();
+  T end_value = end.to<T>();
+  T step_value = step.to<T>();
+
+  int64_t size = 0;
+  GetSize(start_value, end_value, step_value, &size);
+
+  out->Resize(phi::make_ddim({size}));
   dev_ctx.template Alloc<T>(out);
+
   EXEC_NPU_CMD(aclnnArange, dev_ctx, start, end, step, *out);
 }
 

backends/npu/kernels/cast_kernel.cc

@@ -93,13 +93,20 @@ void CastKernel(const Context& dev_ctx,
   DO_COMPATIBILITY(
       aclnnCast,
       (custom_kernel::AclopCastKernel<T, Context>(dev_ctx, x, dtype, out)));
-
+  out->Resize((x.dims()));
   if (x.dtype() == dtype) {
+    if (x.dims() == phi::make_ddim({-1})) {
+      *out = x;
+      return;
+    }
     dev_ctx.template Alloc<T>(out);
     TensorCopy(dev_ctx, x, false, out);
     return;
   }
-
+  if (x.dims() == phi::make_ddim({-1})) {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "canot cast tensor with unknown shape for diffrent dtype"));
+  }
   int aclDtype = ConvertToNpuDtype(dtype);
 
   if (dtype == phi::DataType::FLOAT32) {