Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

SDNQ add NPU support and default to NPU on CPUs with OpenVINO

Browse Source
This commit is contained in:
Disty0
2026-06-19 00:12:22 +03:00
parent 8a9f5c0761
commit 38820cbbd9
1 changed files with 35 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/kernels/openvino_mm.py
+35 -16
View File
@@ -4,19 +4,40 @@ import openvino as ov
from openvino import opset16 as ov_ops
def build_openvino_int_mm(ov_device: str = "CPU"):
core = ov.Core()
input_a = ov_ops.parameter(ov.PartialShape([-1, -1]), ov.Type.i8, name="A")
input_b = ov_ops.parameter(ov.PartialShape([-1, -1]), ov.Type.i8, name="B")
a = ov_ops.fake_quantize(ov_ops.convert(input_a, ov.Type.f32), -128.0, 127.0, -128.0, 127.0, 256)
b = ov_ops.fake_quantize(ov_ops.convert(input_b, ov.Type.f32), -128.0, 127.0, -128.0, 127.0, 256)
core = ov.Core()
OV_DEVICE: str = os.environ.get("SDNQ_OPENVINO_DEVICE", "NPU" if "NPU" in core.get_available_devices() else "CPU")
OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, tuple[int,int] | None], ov.Model] = {}
@torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=())
def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor:
global OV_COMPILED_CACHE, OV_DEVICE # pylint: disable=global-variable-not-assigned
if OV_DEVICE in {"NPU", "CPU"}:
cache_key = (OV_DEVICE, Tensor_A.shape, Tensor_B.shape)
else:
cache_key = (OV_DEVICE, None, None)
ov_int_mm = OV_COMPILED_CACHE.get(cache_key, None)
if ov_int_mm is not None:
return ov_int_mm(Tensor_A, Tensor_B)
if OV_DEVICE in {"NPU", "CPU"}:
shape_a = ov.Shape(Tensor_A.shape)
shape_b = ov.Shape(Tensor_B.shape)
else:
shape_a = ov.PartialShape([-1,-1])
shape_b = ov.PartialShape([-1,-1])
input_a = ov_ops.parameter(shape_a, ov.Type.i8, name="A")
input_b = ov_ops.parameter(shape_b, ov.Type.i8, name="B")
low = ov_ops.constant(-128.0, dtype=ov.Type.f32)
high = ov_ops.constant(127.0, dtype=ov.Type.f32)
a = ov_ops.fake_quantize(ov_ops.convert(input_a, ov.Type.f32), low, high, low, high, 256)
b = ov_ops.fake_quantize(ov_ops.convert(input_b, ov.Type.f32), low, high, low, high, 256)
ov_model = ov.Model([ov_ops.matmul(a, b, False, False)], [input_a, input_b], "ov_int8_mm")
ov_model = core.compile_model(ov_model, ov_device)
ov_model = core.compile_model(ov_model, OV_DEVICE)
infer_request = ov_model.create_infer_request()
out_name = ov_model.outputs[0]
@torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=())
def openvino_int_mm(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
def ov_int_mm(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32)
infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True))
infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True))
@@ -24,11 +45,9 @@ def build_openvino_int_mm(ov_device: str = "CPU"):
infer_request.infer()
return C.to(A.device)
@openvino_int_mm.register_fake
def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))
OV_COMPILED_CACHE[cache_key] = ov_int_mm
return ov_int_mm(Tensor_A, Tensor_B)
return openvino_int_mm
openvino_int_mm = build_openvino_int_mm(ov_device=os.environ.get("SDNQ_OPENVINO_DEVICE", "CPU"))
@openvino_int_mm.register_fake
def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))