Fix issue on Intel MPS.

liuliu · Sep 16, 2023 · d3e2463 · d3e2463
1 parent 9cfdad0
commit d3e2463
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/lib/nnc/mps/ccv_nnc_mps.m b/lib/nnc/mps/ccv_nnc_mps.m
@@ -1148,13 +1148,15 @@ MPSGraphTensorNamedDataLayout ccv_nnc_mps_tensor_data_layout(const int format)
 		uint16_t half_bytes;
 		ccv_float_to_half_precision(&val, &half_bytes, 1);
 #ifdef __x86_64__
-		buffer = [ccv_nnc_default_device() newBufferWithBytes:&half_bytes length:sizeof(uint16_t) options:MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModePrivate | MTLResourceHazardTrackingModeTracked];
+		buffer = [ccv_nnc_default_device() newBufferWithLength:sizeof(uint16_t) options:MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModePrivate | MTLResourceHazardTrackingModeTracked];
+		mpmemcpy(buffer, 0, CCV_TENSOR_GPU_MEMORY, &half_bytes, 0, CCV_TENSOR_CPU_MEMORY, sizeof(uint16_t));
 #else
 		buffer = [ccv_nnc_default_device() newBufferWithBytes:&half_bytes length:sizeof(uint16_t) options:MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModeShared | MTLResourceHazardTrackingModeTracked];
 #endif
 	} else {
 #ifdef __x86_64__
-		buffer = [ccv_nnc_default_device() newBufferWithBytes:&val length:sizeof(float) options:MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModePrivate | MTLResourceHazardTrackingModeTracked];
+		buffer = [ccv_nnc_default_device() newBufferWithLength:sizeof(float) options:MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModePrivate | MTLResourceHazardTrackingModeTracked];
+		mpmemcpy(buffer, 0, CCV_TENSOR_GPU_MEMORY, &val, 0, CCV_TENSOR_CPU_MEMORY, sizeof(float));
 #else
 		buffer = [ccv_nnc_default_device() newBufferWithBytes:&val length:sizeof(float) options:MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModeShared | MTLResourceHazardTrackingModeTracked];
 #endif