# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
NPP 3-Channel Canny Edge Detection - Simple Example
====================================================

Detects edges in color images using NVIDIA NPP.
20x faster than OpenCV, detects 60% more edges.

Requirements:
    pip install torch opencv-python numpy

Usage:
    python npp_canny_simple.py image.jpg
"""

import cv2
import torch
import ctypes
from ctypes import c_int, c_int16, c_ubyte, POINTER, c_void_p, sizeof
import sys
import time

class NPPCanny:
    """NPP 3-Channel Canny Edge Detector"""

    # NPP enum constants
    NPP_FILTER_SOBEL = 0
    NPP_MASK_SIZE_3_X_3 = 200
    NPPI_NORM_L2 = 2
    NPP_BORDER_REPLICATE = 2

    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        # Load NPP library
        import platform
        if platform.system() == 'Windows':
            lib = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1\bin\x64\nppif64_13.dll"
        else:
            lib = "/usr/local/cuda/lib64/libnppif.so"

        self.npp = ctypes.cdll.LoadLibrary(lib)

        # Setup functions
        self.get_buffer_size = self.npp.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [NPPCanny.NppiSize, POINTER(c_int)]

        self.canny = self.npp.nppiFilterCannyBorder_8u_C3C1R_Ctx
        self.canny.restype = c_int
        self.canny.argtypes = [
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize, NPPCanny.NppiPoint,
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize,
            c_int, c_int, c_int16, c_int16, c_int, c_int, c_void_p,
            NPPCanny.NppStreamContext
        ]

    def detect(self, image, low=50, high=100):
        """
        Detect edges in RGB image

        Args:
            image: numpy array (H, W, 3) BGR format
            low: low threshold (0-255)
            high: high threshold (0-255)

        Returns:
            edges: numpy array (H, W) binary edge map
        """
        # Convert to tensor (HWC format required by NPP)
        img_tensor = torch.from_numpy(image).cuda().to(torch.uint8).contiguous()

        h, w, c = img_tensor.shape

        # Allocate output
        output = torch.empty((h, w), dtype=torch.uint8, device='cuda')

        # Get buffer size
        roi = NPPCanny.NppiSize(w, h)
        buf_size = c_int(0)
        status = self.get_buffer_size(roi, ctypes.byref(buf_size))
        if status != 0:
            raise RuntimeError(f"NPP buffer size error: {status}")
        buffer = torch.empty(buf_size.value, dtype=torch.uint8, device='cuda')

        # Setup context - properly initialize all fields
        device_id = torch.cuda.current_device()
        device_props = torch.cuda.get_device_properties(device_id)

        ctx = NPPCanny.NppStreamContext()
        ctx.hStream = c_void_p(torch.cuda.current_stream().cuda_stream)
        ctx.nCudaDeviceId = device_id
        ctx.nMultiProcessorCount = device_props.multi_processor_count
        ctx.nMaxThreadsPerMultiProcessor = device_props.max_threads_per_multi_processor
        ctx.nMaxThreadsPerBlock = device_props.max_threads_per_block
        ctx.nSharedMemPerBlock = device_props.total_constant_memory  # Approximation
        ctx.nCudaDevAttrComputeCapabilityMajor = device_props.major
        ctx.nCudaDevAttrComputeCapabilityMinor = device_props.minor
        ctx.nStreamFlags = 0

        # Run Canny
        status = self.canny(
            ctypes.cast(img_tensor.data_ptr(), POINTER(c_ubyte)),
            3 * w * sizeof(c_ubyte),
            NPPCanny.NppiSize(w, h),
            NPPCanny.NppiPoint(0, 0),
            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte),
            roi,
            NPPCanny.NPP_FILTER_SOBEL,
            NPPCanny.NPP_MASK_SIZE_3_X_3,
            c_int16(low),
            c_int16(high),
            NPPCanny.NPPI_NORM_L2,
            NPPCanny.NPP_BORDER_REPLICATE,
            ctypes.cast(buffer.data_ptr(), c_void_p),
            ctx
        )

        # Check NPP status
        if status != 0:
            raise RuntimeError(f"NPP error: {status}")

        return output.cpu().numpy()

def main():
    # Load image
    image_path = sys.argv[1] if len(sys.argv) > 1 else "image.jpg"
    image = cv2.imread(image_path)

    if image is None:
        print(f"Error: Could not read {image_path}")
        return

    print(f"Image: {image.shape[1]}×{image.shape[0]}")

    # Check GPU
    if not torch.cuda.is_available():
        print("Error: CUDA not available")
        return

    print(f"GPU: {torch.cuda.get_device_name(0)}")

    # NPP 3-Channel Canny
    print("\nNPP 3-Channel Canny...")
    detector = NPPCanny()

    # Warmup run (initializes CUDA context)
    _ = detector.detect(image, low=50, high=100)
    torch.cuda.synchronize()

    # Timed run (average of 10 iterations)
    start = time.time()
    for _ in range(10):
        edges_npp = detector.detect(image, low=50, high=100)
    torch.cuda.synchronize()
    npp_time = (time.time() - start) * 1000 / 10

    print(f"  Time: {npp_time:.2f} ms")
    print(f"  Edges: {edges_npp.sum() // 255} pixels")

    # OpenCV Grayscale (comparison)
    print("\nOpenCV Grayscale...")
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Average of 10 iterations for fair comparison
    start = time.time()
    for _ in range(10):
        edges_cv = cv2.Canny(gray, 50, 100, L2gradient=True)
    cv_time = (time.time() - start) * 1000 / 10

    print(f"  Time: {cv_time:.2f} ms")
    print(f"  Edges: {edges_cv.sum() // 255} pixels")

    # OpenCV 3-Channel (naive approach: 3 separate Canny + merge)
    print("\nOpenCV 3-Channel (3× separate)...")
    b, g, r = cv2.split(image)

    # Average of 10 iterations
    start = time.time()
    for _ in range(10):
        edges_b = cv2.Canny(b, 50, 100, L2gradient=True)
        edges_g = cv2.Canny(g, 50, 100, L2gradient=True)
        edges_r = cv2.Canny(r, 50, 100, L2gradient=True)
        edges_cv_3ch = cv2.bitwise_or(edges_r, cv2.bitwise_or(edges_g, edges_b))
    cv_3ch_time = (time.time() - start) * 1000 / 10

    print(f"  Time: {cv_3ch_time:.2f} ms")
    print(f"  Edges: {edges_cv_3ch.sum() // 255} pixels")

    # Results
    print("\n" + "="*50)
    print("PERFORMANCE COMPARISON")
    print("="*50)
    print(f"NPP 3-Channel:       {npp_time:6.2f} ms  ({edges_npp.sum() // 255:5d} edges)")
    print(f"OpenCV Grayscale:    {cv_time:6.2f} ms  ({edges_cv.sum() // 255:5d} edges)")
    print(f"OpenCV 3-Channel:    {cv_3ch_time:6.2f} ms  ({edges_cv_3ch.sum() // 255:5d} edges)")
    print("="*50)
    print(f"NPP vs OpenCV Gray:  {cv_time/npp_time:5.1f}× faster")
    print(f"NPP vs OpenCV 3-Ch:  {cv_3ch_time/npp_time:5.1f}× faster")
    print(f"Extra edges vs Gray: +{((edges_npp.sum() - edges_cv.sum()) / edges_cv.sum() * 100):4.0f}%")
    print("="*50)

    # Save outputs
    cv2.imwrite("edges_npp.png", edges_npp)
    cv2.imwrite("edges_opencv_gray.png", edges_cv)
    cv2.imwrite("edges_opencv_3ch.png", edges_cv_3ch)

    print("\nSaved: edges_npp.png, edges_opencv_gray.png, edges_opencv_3ch.png")

if __name__ == "__main__":
    main()


NPPCanny

====================================================

Detects edges in color images using NVIDIA NPP.
20x faster than OpenCV, detects 60% more edges.

Requirements:
    pip install torch opencv-python numpy

Usage:
    python npp_canny_simple.py image.jpg
"""

import cv2
import torch
import ctypes
from ctypes import c_int, c_int16, c_ubyte, POINTER, c_void_p, sizeof
import sys
import time

class NPPCanny:
    """NPP 3-Channel Canny Edge Detector"""

    # NPP enum constants
    NPP_FILTER_SOBEL = 0
    NPP_MASK_SIZE_3_X_3 = 200
    NPPI_NORM_L2 = 2
    NPP_BORDER_REPLICATE = 2

    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        # Load NPP library
        import platform
        if platform.system() == 'Windows':
            lib = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1\bin\x64\nppif64_13.dll"
        else:
            lib = "/usr/local/cuda/lib64/libnppif.so"

        self.npp = ctypes.cdll.LoadLibrary(lib)

        # Setup functions
        self.get_buffer_size = self.npp.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [NPPCanny.NppiSize, POINTER(c_int)]

        self.canny = self.npp.nppiFilterCannyBorder_8u_C3C1R_Ctx
        self.canny.restype = c_int
        self.canny.argtypes = [
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize, NPPCanny.NppiPoint,
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize,
            c_int, c_int, c_int16, c_int16, c_int, c_int, c_void_p,
            NPPCanny.NppStreamContext
        ]

    def detect(self, image, low=50, high=100):
        """
        Detect edges in RGB image

        Args:
            image: numpy array (H, W, 3) BGR format
            low: low threshold (0-255)
            high: high threshold (0-255)

        Returns:
            edges: numpy array (H, W) binary edge map
        """
        # Convert to tensor (HWC format required by NPP)
        img_tensor = torch.from_numpy(image).cuda().to(torch.uint8).contiguous()

        h, w, c = img_tensor.shape

        # Allocate output
        output = torch.empty((h, w), dtype=torch.uint8, device='cuda')

        # Get buffer size
        roi = NPPCanny.NppiSize(w, h)
        buf_size = c_int(0)
        status = self.get_buffer_size(roi, ctypes.byref(buf_size))
        if status != 0:
            raise RuntimeError(f"NPP buffer size error: {status}")
        buffer = torch.empty(buf_size.value, dtype=torch.uint8, device='cuda')

        # Setup context - properly initialize all fields
        device_id = torch.cuda.current_device()
        device_props = torch.cuda.get_device_properties(device_id)

        ctx = NPPCanny.NppStreamContext()
        ctx.hStream = c_void_p(torch.cuda.current_stream().cuda_stream)
        ctx.nCudaDeviceId = device_id
        ctx.nMultiProcessorCount = device_props.multi_processor_count
        ctx.nMaxThreadsPerMultiProcessor = device_props.max_threads_per_multi_processor
        ctx.nMaxThreadsPerBlock = device_props.max_threads_per_block
        ctx.nSharedMemPerBlock = device_props.total_constant_memory  # Approximation
        ctx.nCudaDevAttrComputeCapabilityMajor = device_props.major
        ctx.nCudaDevAttrComputeCapabilityMinor = device_props.minor
        ctx.nStreamFlags = 0

        # Run Canny
        status = self.canny(
            ctypes.cast(img_tensor.data_ptr(), POINTER(c_ubyte)),
            3 * w * sizeof(c_ubyte),
            NPPCanny.NppiSize(w, h),
            NPPCanny.NppiPoint(0, 0),
            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte),
            roi,
            NPPCanny.NPP_FILTER_SOBEL,
            NPPCanny.NPP_MASK_SIZE_3_X_3,
            c_int16(low),
            c_int16(high),
            NPPCanny.NPPI_NORM_L2,
            NPPCanny.NPP_BORDER_REPLICATE,
            ctypes.cast(buffer.data_ptr(), c_void_p),
            ctx
        )

        # Check NPP status
        if status != 0:
            raise RuntimeError(f"NPP error: {status}")

        return output.cpu().numpy()

def main():
    # Load image
    image_path = sys.argv[1] if len(sys.argv) > 1 else "image.jpg"
    image = cv2.imread(image_path)

    if image is None:
        print(f"Error: Could not read {image_path}")
        return

    print(f"Image: {image.shape[1]}×{image.shape[0]}")

    # Check GPU
    if not torch.cuda.is_available():
        print("Error: CUDA not available")
        return

    print(f"GPU: {torch.cuda.get_device_name(0)}")


main

            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte),
            roi,
            NPPCanny.NPP_FILTER_SOBEL,
            NPPCanny.NPP_MASK_SIZE_3_X_3,
            c_int16(low),
            c_int16(high),
            NPPCanny.NPPI_NORM_L2,
            NPPCanny.NPP_BORDER_REPLICATE,
            ctypes.cast(buffer.data_ptr(), c_void_p),
            ctx
        )

        # Check NPP status
        if status != 0:
            raise RuntimeError(f"NPP error: {status}")

        return output.cpu().numpy()

def main():
    # Load image
    image_path = sys.argv[1] if len(sys.argv) > 1 else "image.jpg"
    image = cv2.imread(image_path)

    if image is None:
        print(f"Error: Could not read {image_path}")
        return

    print(f"Image: {image.shape[1]}×{image.shape[0]}")

    # Check GPU
    if not torch.cuda.is_available():
        print("Error: CUDA not available")
        return

    print(f"GPU: {torch.cuda.get_device_name(0)}")

    # NPP 3-Channel Canny
    print("\nNPP 3-Channel Canny...")
    detector = NPPCanny()

    # Warmup run (initializes CUDA context)
    _ = detector.detect(image, low=50, high=100)
    torch.cuda.synchronize()

    # Timed run (average of 10 iterations)
    start = time.time()
    for _ in range(10):
        edges_npp = detector.detect(image, low=50, high=100)
    torch.cuda.synchronize()
    npp_time = (time.time() - start) * 1000 / 10

    print(f"  Time: {npp_time:.2f} ms")
    print(f"  Edges: {edges_npp.sum() // 255} pixels")

    # OpenCV Grayscale (comparison)
    print("\nOpenCV Grayscale...")
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Average of 10 iterations for fair comparison
    start = time.time()
    for _ in range(10):
        edges_cv = cv2.Canny(gray, 50, 100, L2gradient=True)
    cv_time = (time.time() - start) * 1000 / 10

    print(f"  Time: {cv_time:.2f} ms")
    print(f"  Edges: {edges_cv.sum() // 255} pixels")

    # OpenCV 3-Channel (naive approach: 3 separate Canny + merge)
    print("\nOpenCV 3-Channel (3× separate)...")
    b, g, r = cv2.split(image)

    # Average of 10 iterations
    start = time.time()
    for _ in range(10):
        edges_b = cv2.Canny(b, 50, 100, L2gradient=True)
        edges_g = cv2.Canny(g, 50, 100, L2gradient=True)
        edges_r = cv2.Canny(r, 50, 100, L2gradient=True)
        edges_cv_3ch = cv2.bitwise_or(edges_r, cv2.bitwise_or(edges_g, edges_b))
    cv_3ch_time = (time.time() - start) * 1000 / 10

    print(f"  Time: {cv_3ch_time:.2f} ms")
    print(f"  Edges: {edges_cv_3ch.sum() // 255} pixels")

    # Results
    print("\n" + "="*50)
    print("PERFORMANCE COMPARISON")
    print("="*50)
    print(f"NPP 3-Channel:       {npp_time:6.2f} ms  ({edges_npp.sum() // 255:5d} edges)")
    print(f"OpenCV Grayscale:    {cv_time:6.2f} ms  ({edges_cv.sum() // 255:5d} edges)")
    print(f"OpenCV 3-Channel:    {cv_3ch_time:6.2f} ms  ({edges_cv_3ch.sum() // 255:5d} edges)")
    print("="*50)
    print(f"NPP vs OpenCV Gray:  {cv_time/npp_time:5.1f}× faster")
    print(f"NPP vs OpenCV 3-Ch:  {cv_3ch_time/npp_time:5.1f}× faster")
    print(f"Extra edges vs Gray: +{((edges_npp.sum() - edges_cv.sum()) / edges_cv.sum() * 100):4.0f}%")
    print("="*50)

    # Save outputs
    cv2.imwrite("edges_npp.png", edges_npp)
    cv2.imwrite("edges_opencv_gray.png", edges_cv)
    cv2.imwrite("edges_opencv_3ch.png", edges_cv_3ch)

    print("\nSaved: edges_npp.png, edges_opencv_gray.png, edges_opencv_3ch.png")

if __name__ == "__main__":
    main()


NppiSize

    python npp_canny_simple.py image.jpg
"""

import cv2
import torch
import ctypes
from ctypes import c_int, c_int16, c_ubyte, POINTER, c_void_p, sizeof
import sys
import time

class NPPCanny:
    """NPP 3-Channel Canny Edge Detector"""

    # NPP enum constants
    NPP_FILTER_SOBEL = 0
    NPP_MASK_SIZE_3_X_3 = 200
    NPPI_NORM_L2 = 2
    NPP_BORDER_REPLICATE = 2

    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        # Load NPP library
        import platform
        if platform.system() == 'Windows':
            lib = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1\bin\x64\nppif64_13.dll"
        else:


NppiPoint

import cv2
import torch
import ctypes
from ctypes import c_int, c_int16, c_ubyte, POINTER, c_void_p, sizeof
import sys
import time

class NPPCanny:
    """NPP 3-Channel Canny Edge Detector"""

    # NPP enum constants
    NPP_FILTER_SOBEL = 0
    NPP_MASK_SIZE_3_X_3 = 200
    NPPI_NORM_L2 = 2
    NPP_BORDER_REPLICATE = 2

    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        # Load NPP library
        import platform
        if platform.system() == 'Windows':
            lib = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1\bin\x64\nppif64_13.dll"
        else:
            lib = "/usr/local/cuda/lib64/libnppif.so"

        self.npp = ctypes.cdll.LoadLibrary(lib)


NppStreamContext

from ctypes import c_int, c_int16, c_ubyte, POINTER, c_void_p, sizeof
import sys
import time

class NPPCanny:
    """NPP 3-Channel Canny Edge Detector"""

    # NPP enum constants
    NPP_FILTER_SOBEL = 0
    NPP_MASK_SIZE_3_X_3 = 200
    NPPI_NORM_L2 = 2
    NPP_BORDER_REPLICATE = 2

    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        # Load NPP library
        import platform
        if platform.system() == 'Windows':
            lib = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1\bin\x64\nppif64_13.dll"
        else:
            lib = "/usr/local/cuda/lib64/libnppif.so"

        self.npp = ctypes.cdll.LoadLibrary(lib)

        # Setup functions
        self.get_buffer_size = self.npp.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [NPPCanny.NppiSize, POINTER(c_int)]

        self.canny = self.npp.nppiFilterCannyBorder_8u_C3C1R_Ctx
        self.canny.restype = c_int
        self.canny.argtypes = [
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize, NPPCanny.NppiPoint,


    NPP_MASK_SIZE_3_X_3 = 200
    NPPI_NORM_L2 = 2
    NPP_BORDER_REPLICATE = 2

    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        # Load NPP library
        import platform
        if platform.system() == 'Windows':
            lib = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1\bin\x64\nppif64_13.dll"
        else:
            lib = "/usr/local/cuda/lib64/libnppif.so"

        self.npp = ctypes.cdll.LoadLibrary(lib)

        # Setup functions
        self.get_buffer_size = self.npp.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [NPPCanny.NppiSize, POINTER(c_int)]

        self.canny = self.npp.nppiFilterCannyBorder_8u_C3C1R_Ctx
        self.canny.restype = c_int
        self.canny.argtypes = [
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize, NPPCanny.NppiPoint,
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize,
            c_int, c_int, c_int16, c_int16, c_int, c_int, c_void_p,
            NPPCanny.NppStreamContext
        ]

    def detect(self, image, low=50, high=100):
        """
        Detect edges in RGB image

        Args:
            image: numpy array (H, W, 3) BGR format
            low: low threshold (0-255)
            high: high threshold (0-255)

        Returns:
            edges: numpy array (H, W) binary edge map
        """
        # Convert to tensor (HWC format required by NPP)
        img_tensor = torch.from_numpy(image).cuda().to(torch.uint8).contiguous()

        h, w, c = img_tensor.shape

        # Allocate output
        output = torch.empty((h, w), dtype=torch.uint8, device='cuda')


detect

            lib = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1\bin\x64\nppif64_13.dll"
        else:
            lib = "/usr/local/cuda/lib64/libnppif.so"

        self.npp = ctypes.cdll.LoadLibrary(lib)

        # Setup functions
        self.get_buffer_size = self.npp.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [NPPCanny.NppiSize, POINTER(c_int)]

        self.canny = self.npp.nppiFilterCannyBorder_8u_C3C1R_Ctx
        self.canny.restype = c_int
        self.canny.argtypes = [
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize, NPPCanny.NppiPoint,
            POINTER(c_ubyte), c_int, NPPCanny.NppiSize,
            c_int, c_int, c_int16, c_int16, c_int, c_int, c_void_p,
            NPPCanny.NppStreamContext
        ]

    def detect(self, image, low=50, high=100):
        """
        Detect edges in RGB image

        Args:
            image: numpy array (H, W, 3) BGR format
            low: low threshold (0-255)
            high: high threshold (0-255)

        Returns:
            edges: numpy array (H, W) binary edge map
        """
        # Convert to tensor (HWC format required by NPP)
        img_tensor = torch.from_numpy(image).cuda().to(torch.uint8).contiguous()

        h, w, c = img_tensor.shape

        # Allocate output
        output = torch.empty((h, w), dtype=torch.uint8, device='cuda')

        # Get buffer size
        roi = NPPCanny.NppiSize(w, h)
        buf_size = c_int(0)
        status = self.get_buffer_size(roi, ctypes.byref(buf_size))
        if status != 0:
            raise RuntimeError(f"NPP buffer size error: {status}")
        buffer = torch.empty(buf_size.value, dtype=torch.uint8, device='cuda')

        # Setup context - properly initialize all fields
        device_id = torch.cuda.current_device()
        device_props = torch.cuda.get_device_properties(device_id)

        ctx = NPPCanny.NppStreamContext()
        ctx.hStream = c_void_p(torch.cuda.current_stream().cuda_stream)
        ctx.nCudaDeviceId = device_id
        ctx.nMultiProcessorCount = device_props.multi_processor_count
        ctx.nMaxThreadsPerMultiProcessor = device_props.max_threads_per_multi_processor
        ctx.nMaxThreadsPerBlock = device_props.max_threads_per_block
        ctx.nSharedMemPerBlock = device_props.total_constant_memory  # Approximation
        ctx.nCudaDevAttrComputeCapabilityMajor = device_props.major
        ctx.nCudaDevAttrComputeCapabilityMinor = device_props.minor
        ctx.nStreamFlags = 0

        # Run Canny
        status = self.canny(
            ctypes.cast(img_tensor.data_ptr(), POINTER(c_ubyte)),
            3 * w * sizeof(c_ubyte),
            NPPCanny.NppiSize(w, h),
            NPPCanny.NppiPoint(0, 0),
            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte),
            roi,
            NPPCanny.NPP_FILTER_SOBEL,
            NPPCanny.NPP_MASK_SIZE_3_X_3,
            c_int16(low),
            c_int16(high),
            NPPCanny.NPPI_NORM_L2,
            NPPCanny.NPP_BORDER_REPLICATE,
            ctypes.cast(buffer.data_ptr(), c_void_p),
            ctx
        )

        # Check NPP status
        if status != 0:
            raise RuntimeError(f"NPP error: {status}")

        return output.cpu().numpy()

def main():
    # Load image
    image_path = sys.argv[1] if len(sys.argv) > 1 else "image.jpg"
    image = cv2.imread(image_path)

    if image is None:
        print(f"Error: Could not read {image_path}")
        return

    print(f"Image: {image.shape[1]}×{image.shape[0]}")

    # Check GPU
    if not torch.cuda.is_available():
        print("Error: CUDA not available")
        return

    print(f"GPU: {torch.cuda.get_device_name(0)}")


#!/usr/bin/env python3

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import sys

from nvMatmulHeuristics import (
    NvMatmulHeuristicsInterface,
    NvMatmulHeuristicsTarget,
    NvMatmulHeuristicsFlags,
    NvMatmulHeuristicsMatmulLayout,
    NvMatmulHeuristicsNvidiaGpu,
    layoutToStr
)

def parse_args():
    parser = argparse.ArgumentParser(
        description='Get and display GEMM configurations for specified parameters'
    )

    parser.add_argument('-M', '--m-dim',
                        type=int,
                        required=True,
                        help='Output matrix height')

    parser.add_argument('-N', '--n-dim',
                        type=int,
                        required=True,
                        help='Output matrix width')

    parser.add_argument('-K', '--k-dim',
                        type=int,
                        required=True,
                        help='Reduced dimension')

    parser.add_argument('-B', '--batch-size',
                        type=int,
                        default=1,
                        help='Batch size (default: 1)')

    parser.add_argument('--gpu',
                        type=str,
                        choices=[gpu.name for gpu in NvMatmulHeuristicsNvidiaGpu if gpu.name != 'END'],
                        required=True,
                        help='Target GPU')

    parser.add_argument('--layout',
                        type=str,
                        choices=[layout.name for layout in NvMatmulHeuristicsMatmulLayout if layout.name != 'END'],
                        default='NN_ROW_MAJOR',
                        help='Matrix layout (default: NN_ROW_MAJOR)')

    parser.add_argument('--backend',
                        type=str,
                        choices=[backend.name for backend in NvMatmulHeuristicsTarget if backend.name != 'END'],
                        default='CUTLASS3',
                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--count',
                        type=int,
                        default=8,
                        help='Number of configurations to retrieve (default: 8)')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    return parser.parse_args()

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nGetting {args.count} configurations for:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"Batch size: {args.batch_size}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}\n")

    # Initialize nvMatmulHeuristics interface
    try:
        nvMatmulHeuristics = NvMatmulHeuristicsInterface(
            path=args.lib_path,
            backend=backend,
            precision=args.precision,
            flags=NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING
        )
    except OSError as e:
        print(f"Error: Failed to load nvMatmulHeuristics library: {e}")
        print(f"Make sure the library exists at: {args.lib_path}")
        return 1
    except AssertionError:
        print("Error: Version mismatch or unsupported precision")
        return 1

    try:
        # Create and initialize hardware descriptor
        hw_desc = nvMatmulHeuristics.createHardwareDescriptor()
        try:
            # Set the target GPU
            nvMatmulHeuristics.setHardwarePredefinedGpu(hw_desc, gpu)

            # Load internal discovery set for the specified layout
            success = nvMatmulHeuristics.loadInternalDiscoverySet(layout, hw_desc)
            if success:
                print("Successfully loaded internal discovery set.")
            else:
                print("Failed to load internal discovery set. Make sure it is available for selected hardware and layout.")

            # Create problem object
            problem = nvMatmulHeuristics.makeNvMatmulHeuristicsProblem(
                args.m_dim, args.n_dim, args.k_dim, layout, args.batch_size
            )

            # Get configurations using the problem object
            configs = nvMatmulHeuristics.get(problem, args.count, hw_desc)

            # Print results
            print(f"Found {len(configs)} configurations:\n")
            for i, config in enumerate(configs, 1):
                kernel = config["kernel"]
                runtime = config["runtime"]
                print(f"Configuration {i}:")
                print(f"  Kernel: {kernel}")
                print(f"  Estimated runtime: {runtime * 1000:.6f} ms\n")

        finally:
            # Clean up hardware descriptor
            nvMatmulHeuristics.destroyHardwareDescriptor(hw_desc)

    except RuntimeError as e:
        print(f"Error: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())


#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import sys

from nvMatmulHeuristics import (
    NvMatmulHeuristicsInterface,
    NvMatmulHeuristicsTarget,
    NvMatmulHeuristicsFlags,
    NvMatmulHeuristicsMatmulLayout,
    NvMatmulHeuristicsNvidiaGpu,
    layoutToStr
)

def parse_args():
    parser = argparse.ArgumentParser(
        description='Get and display GEMM configurations for specified parameters'
    )

    parser.add_argument('-M', '--m-dim',
                        type=int,
                        required=True,
                        help='Output matrix height')

    parser.add_argument('-N', '--n-dim',
                        type=int,
                        required=True,
                        help='Output matrix width')

    parser.add_argument('-K', '--k-dim',
                        type=int,
                        required=True,
                        help='Reduced dimension')

    parser.add_argument('-B', '--batch-size',
                        type=int,
                        default=1,
                        help='Batch size (default: 1)')

    parser.add_argument('--gpu',
                        type=str,
                        choices=[gpu.name for gpu in NvMatmulHeuristicsNvidiaGpu if gpu.name != 'END'],
                        required=True,
                        help='Target GPU')

    parser.add_argument('--layout',
                        type=str,
                        choices=[layout.name for layout in NvMatmulHeuristicsMatmulLayout if layout.name != 'END'],
                        default='NN_ROW_MAJOR',
                        help='Matrix layout (default: NN_ROW_MAJOR)')

    parser.add_argument('--backend',
                        type=str,
                        choices=[backend.name for backend in NvMatmulHeuristicsTarget if backend.name != 'END'],
                        default='CUTLASS3',
                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--count',
                        type=int,
                        default=8,
                        help='Number of configurations to retrieve (default: 8)')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    return parser.parse_args()

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nGetting {args.count} configurations for:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"Batch size: {args.batch_size}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}\n")

    # Initialize nvMatmulHeuristics interface
    try:


main

                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--count',
                        type=int,
                        default=8,
                        help='Number of configurations to retrieve (default: 8)')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    return parser.parse_args()

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nGetting {args.count} configurations for:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"Batch size: {args.batch_size}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}\n")

    # Initialize nvMatmulHeuristics interface
    try:
        nvMatmulHeuristics = NvMatmulHeuristicsInterface(
            path=args.lib_path,
            backend=backend,
            precision=args.precision,
            flags=NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING
        )
    except OSError as e:
        print(f"Error: Failed to load nvMatmulHeuristics library: {e}")
        print(f"Make sure the library exists at: {args.lib_path}")
        return 1
    except AssertionError:
        print("Error: Version mismatch or unsupported precision")
        return 1

    try:
        # Create and initialize hardware descriptor
        hw_desc = nvMatmulHeuristics.createHardwareDescriptor()
        try:
            # Set the target GPU
            nvMatmulHeuristics.setHardwarePredefinedGpu(hw_desc, gpu)

            # Load internal discovery set for the specified layout
            success = nvMatmulHeuristics.loadInternalDiscoverySet(layout, hw_desc)
            if success:
                print("Successfully loaded internal discovery set.")
            else:
                print("Failed to load internal discovery set. Make sure it is available for selected hardware and layout.")

            # Create problem object
            problem = nvMatmulHeuristics.makeNvMatmulHeuristicsProblem(
                args.m_dim, args.n_dim, args.k_dim, layout, args.batch_size
            )

            # Get configurations using the problem object
            configs = nvMatmulHeuristics.get(problem, args.count, hw_desc)

            # Print results
            print(f"Found {len(configs)} configurations:\n")
            for i, config in enumerate(configs, 1):
                kernel = config["kernel"]
                runtime = config["runtime"]
                print(f"Configuration {i}:")
                print(f"  Kernel: {kernel}")
                print(f"  Estimated runtime: {runtime * 1000:.6f} ms\n")

        finally:
            # Clean up hardware descriptor
            nvMatmulHeuristics.destroyHardwareDescriptor(hw_desc)

    except RuntimeError as e:
        print(f"Error: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())


#!/usr/bin/env python3

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import ctypes
import sys

from nvMatmulHeuristics import (
    NvMatmulHeuristicsInterface,
    NvMatmulHeuristicsTarget,
    NvMatmulHeuristicsFlags,
    NvMatmulHeuristicsMatmulLayout,
    NvMatmulHeuristicsNvidiaGpu,
    NvMatmulHeuristicsBackendProperty,
    layoutToStr
)

def parse_args():
    parser = argparse.ArgumentParser(
        description='Compare GEMM configurations with and without SMEM carveout'
    )

    parser.add_argument('-M', '--m-dim',
                        type=int,
                        required=True,
                        help='Output matrix height')

    parser.add_argument('-N', '--n-dim',
                        type=int,
                        required=True,
                        help='Output matrix width')

    parser.add_argument('-K', '--k-dim',
                        type=int,
                        required=True,
                        help='Reduced dimension')

    parser.add_argument('--gpu',
                        type=str,
                        choices=[gpu.name for gpu in NvMatmulHeuristicsNvidiaGpu if gpu.name != 'END'],
                        required=True,
                        help='Target GPU')

    parser.add_argument('--layout',
                        type=str,
                        choices=[layout.name for layout in NvMatmulHeuristicsMatmulLayout if layout.name != 'END'],
                        default='NN_ROW_MAJOR',
                        help='Matrix layout (default: NN_ROW_MAJOR)')

    parser.add_argument('--backend',
                        type=str,
                        choices=[backend.name for backend in NvMatmulHeuristicsTarget if backend.name != 'END'],
                        default='CUTLASS3',
                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--smem-carveout',
                        type=int,
                        required=True,
                        help='SMEM carveout size in bytes')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    return parser.parse_args()

def print_kernel_config(config, title):
    print(f"\n{title}:")
    print(f"  CTA Tile: {config['kernel'].cta_tile_m}x{config['kernel'].cta_tile_n}x{config['kernel'].cta_tile_k}")
    print(f"  Warp Tile: {config['kernel'].warp_tile_m}x{config['kernel'].warp_tile_n}x{config['kernel'].warp_tile_k}")
    print(f"  Instruction Tile: {config['kernel'].instr_tile_m}x{config['kernel'].instr_tile_n}x{config['kernel'].instr_tile_k}")
    print(f"  Split K: {config['kernel'].split_k}")
    print(f"  Stages: {config['kernel'].stages}")
    print(f"  CTA swizzling: {config['kernel'].swizzle_factor}")
    print(f"  CTA Order: {config['kernel'].cta_order}")
    print(f"  Cluster Config: {config['kernel'].cluster_m}x{config['kernel'].cluster_n}")
    print(f"  Estimated Runtime: {config['runtime'] * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nComparing configurations with and without SMEM carveout:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}")
    print(f"SMEM Carveout: {args.smem_carveout} bytes\n")

    # Initialize nvMatmulHeuristics interface
    try:
        nvMatmulHeuristics = NvMatmulHeuristicsInterface(
            path=args.lib_path,
            backend=backend,
            precision=args.precision,
            flags=NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING
        )
    except OSError as e:
        print(f"Error: Failed to load nvMatmulHeuristics library: {e}")
        print(f"Make sure the library exists at: {args.lib_path}")
        return 1
    except AssertionError:
        print("Error: Version mismatch or unsupported precision")
        return 1

    try:
        # Create and initialize hardware descriptor
        hw_desc = nvMatmulHeuristics.createHardwareDescriptor()
        try:
            # Set the target GPU
            nvMatmulHeuristics.setHardwarePredefinedGpu(hw_desc, gpu)

            # Load internal discovery set for the specified layout
            success = nvMatmulHeuristics.loadInternalDiscoverySet(layout, hw_desc)
            if not success:
                print("Failed to load internal discovery set. Make sure it is available for selected hardware and layout.")
                return 1

            # Create problem object
            problem = nvMatmulHeuristics.makeNvMatmulHeuristicsProblem(
                args.m_dim, args.n_dim, args.k_dim, layout
            )

            # Get configurations without SMEM carveout
            configs_no_carveout = nvMatmulHeuristics.get(problem, 1, hw_desc)
            if not configs_no_carveout:
                print("No configurations found without SMEM carveout")
                return 1

            # Create backend with SMEM carveout
            backend_obj = nvMatmulHeuristics.createBackend(backend)
            try:
                # Set SMEM carveout size as int32_t
                smem_carveout = ctypes.c_int32(args.smem_carveout)
                nvMatmulHeuristics.setBackendValueProperty(
                    backend_obj,
                    NvMatmulHeuristicsBackendProperty.SMEM_CARVEOUT_SIZE,
                    ctypes.byref(smem_carveout),
                    ctypes.sizeof(smem_carveout)
                )

                # Get configurations with SMEM carveout
                configs_with_carveout = nvMatmulHeuristics.getEx(problem, 1, backend_obj, hw_desc)
                if not configs_with_carveout:
                    print("No configurations found with SMEM carveout")
                    return 1

                # Print and compare configurations
                print_kernel_config(configs_no_carveout[0], "Configuration without SMEM carveout")
                print_kernel_config(configs_with_carveout[0], "Configuration with SMEM carveout")

            finally:
                nvMatmulHeuristics.destroyBackend(backend_obj)

        finally:
            nvMatmulHeuristics.destroyHardwareDescriptor(hw_desc)

    except RuntimeError as e:
        print(f"Error: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())


# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import ctypes
import sys

from nvMatmulHeuristics import (
    NvMatmulHeuristicsInterface,
    NvMatmulHeuristicsTarget,
    NvMatmulHeuristicsFlags,
    NvMatmulHeuristicsMatmulLayout,
    NvMatmulHeuristicsNvidiaGpu,
    NvMatmulHeuristicsBackendProperty,
    layoutToStr
)

def parse_args():
    parser = argparse.ArgumentParser(
        description='Compare GEMM configurations with and without SMEM carveout'
    )

    parser.add_argument('-M', '--m-dim',
                        type=int,
                        required=True,
                        help='Output matrix height')

    parser.add_argument('-N', '--n-dim',
                        type=int,
                        required=True,
                        help='Output matrix width')

    parser.add_argument('-K', '--k-dim',
                        type=int,
                        required=True,
                        help='Reduced dimension')

    parser.add_argument('--gpu',
                        type=str,
                        choices=[gpu.name for gpu in NvMatmulHeuristicsNvidiaGpu if gpu.name != 'END'],
                        required=True,
                        help='Target GPU')

    parser.add_argument('--layout',
                        type=str,
                        choices=[layout.name for layout in NvMatmulHeuristicsMatmulLayout if layout.name != 'END'],
                        default='NN_ROW_MAJOR',
                        help='Matrix layout (default: NN_ROW_MAJOR)')

    parser.add_argument('--backend',
                        type=str,
                        choices=[backend.name for backend in NvMatmulHeuristicsTarget if backend.name != 'END'],
                        default='CUTLASS3',
                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--smem-carveout',
                        type=int,
                        required=True,
                        help='SMEM carveout size in bytes')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    return parser.parse_args()

def print_kernel_config(config, title):
    print(f"\n{title}:")
    print(f"  CTA Tile: {config['kernel'].cta_tile_m}x{config['kernel'].cta_tile_n}x{config['kernel'].cta_tile_k}")
    print(f"  Warp Tile: {config['kernel'].warp_tile_m}x{config['kernel'].warp_tile_n}x{config['kernel'].warp_tile_k}")
    print(f"  Instruction Tile: {config['kernel'].instr_tile_m}x{config['kernel'].instr_tile_n}x{config['kernel'].instr_tile_k}")
    print(f"  Split K: {config['kernel'].split_k}")
    print(f"  Stages: {config['kernel'].stages}")
    print(f"  CTA swizzling: {config['kernel'].swizzle_factor}")
    print(f"  CTA Order: {config['kernel'].cta_order}")
    print(f"  Cluster Config: {config['kernel'].cluster_m}x{config['kernel'].cluster_n}")
    print(f"  Estimated Runtime: {config['runtime'] * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]


                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--smem-carveout',
                        type=int,
                        required=True,
                        help='SMEM carveout size in bytes')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    return parser.parse_args()

def print_kernel_config(config, title):
    print(f"\n{title}:")
    print(f"  CTA Tile: {config['kernel'].cta_tile_m}x{config['kernel'].cta_tile_n}x{config['kernel'].cta_tile_k}")
    print(f"  Warp Tile: {config['kernel'].warp_tile_m}x{config['kernel'].warp_tile_n}x{config['kernel'].warp_tile_k}")
    print(f"  Instruction Tile: {config['kernel'].instr_tile_m}x{config['kernel'].instr_tile_n}x{config['kernel'].instr_tile_k}")
    print(f"  Split K: {config['kernel'].split_k}")
    print(f"  Stages: {config['kernel'].stages}")
    print(f"  CTA swizzling: {config['kernel'].swizzle_factor}")
    print(f"  CTA Order: {config['kernel'].cta_order}")
    print(f"  Cluster Config: {config['kernel'].cluster_m}x{config['kernel'].cluster_n}")
    print(f"  Estimated Runtime: {config['runtime'] * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nComparing configurations with and without SMEM carveout:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}")
    print(f"SMEM Carveout: {args.smem_carveout} bytes\n")

    # Initialize nvMatmulHeuristics interface
    try:


                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    return parser.parse_args()

def print_kernel_config(config, title):
    print(f"\n{title}:")
    print(f"  CTA Tile: {config['kernel'].cta_tile_m}x{config['kernel'].cta_tile_n}x{config['kernel'].cta_tile_k}")
    print(f"  Warp Tile: {config['kernel'].warp_tile_m}x{config['kernel'].warp_tile_n}x{config['kernel'].warp_tile_k}")
    print(f"  Instruction Tile: {config['kernel'].instr_tile_m}x{config['kernel'].instr_tile_n}x{config['kernel'].instr_tile_k}")
    print(f"  Split K: {config['kernel'].split_k}")
    print(f"  Stages: {config['kernel'].stages}")
    print(f"  CTA swizzling: {config['kernel'].swizzle_factor}")
    print(f"  CTA Order: {config['kernel'].cta_order}")
    print(f"  Cluster Config: {config['kernel'].cluster_m}x{config['kernel'].cluster_n}")
    print(f"  Estimated Runtime: {config['runtime'] * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nComparing configurations with and without SMEM carveout:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}")
    print(f"SMEM Carveout: {args.smem_carveout} bytes\n")

    # Initialize nvMatmulHeuristics interface
    try:
        nvMatmulHeuristics = NvMatmulHeuristicsInterface(
            path=args.lib_path,
            backend=backend,
            precision=args.precision,
            flags=NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING
        )
    except OSError as e:
        print(f"Error: Failed to load nvMatmulHeuristics library: {e}")
        print(f"Make sure the library exists at: {args.lib_path}")
        return 1
    except AssertionError:
        print("Error: Version mismatch or unsupported precision")
        return 1

    try:
        # Create and initialize hardware descriptor
        hw_desc = nvMatmulHeuristics.createHardwareDescriptor()
        try:
            # Set the target GPU
            nvMatmulHeuristics.setHardwarePredefinedGpu(hw_desc, gpu)

            # Load internal discovery set for the specified layout
            success = nvMatmulHeuristics.loadInternalDiscoverySet(layout, hw_desc)
            if not success:
                print("Failed to load internal discovery set. Make sure it is available for selected hardware and layout.")
                return 1

            # Create problem object
            problem = nvMatmulHeuristics.makeNvMatmulHeuristicsProblem(
                args.m_dim, args.n_dim, args.k_dim, layout
            )

            # Get configurations without SMEM carveout
            configs_no_carveout = nvMatmulHeuristics.get(problem, 1, hw_desc)
            if not configs_no_carveout:
                print("No configurations found without SMEM carveout")
                return 1

            # Create backend with SMEM carveout
            backend_obj = nvMatmulHeuristics.createBackend(backend)
            try:
                # Set SMEM carveout size as int32_t
                smem_carveout = ctypes.c_int32(args.smem_carveout)
                nvMatmulHeuristics.setBackendValueProperty(
                    backend_obj,
                    NvMatmulHeuristicsBackendProperty.SMEM_CARVEOUT_SIZE,
                    ctypes.byref(smem_carveout),
                    ctypes.sizeof(smem_carveout)
                )

                # Get configurations with SMEM carveout
                configs_with_carveout = nvMatmulHeuristics.getEx(problem, 1, backend_obj, hw_desc)
                if not configs_with_carveout:
                    print("No configurations found with SMEM carveout")
                    return 1

                # Print and compare configurations
                print_kernel_config(configs_no_carveout[0], "Configuration without SMEM carveout")
                print_kernel_config(configs_with_carveout[0], "Configuration with SMEM carveout")

            finally:
                nvMatmulHeuristics.destroyBackend(backend_obj)

        finally:
            nvMatmulHeuristics.destroyHardwareDescriptor(hw_desc)

    except RuntimeError as e:
        print(f"Error: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())


#!/usr/bin/env python3

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import sys

from nvMatmulHeuristics import (
    NvMatmulHeuristicsInterfaceEx,
    NvMatmulHeuristicsTarget,
    NvMatmulHeuristicsFlags,
    NvMatmulHeuristicsMatmulLayout,
    NvMatmulHeuristicsNvidiaGpu,
    layoutToStr
)

def parse_args():
    parser = argparse.ArgumentParser(
        description='Get and display GEMM configurations using NvMatmulHeuristicsInterfaceEx'
    )

    parser.add_argument('-M', '--m-dim',
                        type=int,
                        required=True,
                        help='Output matrix height')

    parser.add_argument('-N', '--n-dim',
                        type=int,
                        required=True,
                        help='Output matrix width')

    parser.add_argument('-K', '--k-dim',
                        type=int,
                        required=True,
                        help='Reduced dimension')

    parser.add_argument('-B', '--batch-size',
                        type=int,
                        default=1,
                        help='Batch size (default: 1)')

    parser.add_argument('--gpu',
                        type=str,
                        choices=[gpu.name for gpu in NvMatmulHeuristicsNvidiaGpu if gpu.name != 'END'],
                        required=True,
                        help='Target GPU')

    parser.add_argument('--layout',
                        type=str,
                        choices=[layout.name for layout in NvMatmulHeuristicsMatmulLayout if layout.name != 'END'],
                        default='NN_ROW_MAJOR',
                        help='Matrix layout (default: NN_ROW_MAJOR)')

    parser.add_argument('--backend',
                        type=str,
                        choices=[backend.name for backend in NvMatmulHeuristicsTarget if backend.name != 'END'],
                        default='CUTLASS3',
                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--count',
                        type=int,
                        default=8,
                        help='Number of configurations to retrieve (default: 8)')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    parser.add_argument('--no-auto-load',
                        action='store_true',
                        help='Disable automatic loading of discovery sets')

    return parser.parse_args()

def print_kernel_config(config, title):
    """Print kernel configuration details."""
    kernel = config["kernel"]
    runtime = config["runtime"]
    print(f"\n{title}:")
    print(f"  Kernel: {kernel}")
    print(f"  Estimated runtime: {runtime * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nGetting {args.count} configurations for:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"Batch size: {args.batch_size}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}")
    print(f"Auto-load discovery sets: {not args.no_auto_load}\n")

    # Initialize nvMatmulHeuristics interface
    try:
        # Create interface with the GPU
        nvMatmulHeuristics = NvMatmulHeuristicsInterfaceEx(
            path=args.lib_path,
            backend=backend,
            flags=NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING,
            load_discovery_implicitly=not args.no_auto_load,
            gpu=gpu
        )

        # Create problem object
        problem = nvMatmulHeuristics.makeNvMatmulHeuristicsProblem(
            args.m_dim, args.n_dim, args.k_dim, layout, args.batch_size
        )

        # Set precision
        precision = args.precision
        print(f"\nGetting configurations with precision ({precision})...")

        # Track loaded discovery sets before the call
        loaded_before = set(nvMatmulHeuristics._loaded_discovery_sets.keys())

        # Get configurations
        configs = nvMatmulHeuristics.get(problem, args.count, precision=precision)

        # Check which discovery sets were loaded
        loaded_after = set(nvMatmulHeuristics._loaded_discovery_sets.keys())
        newly_loaded = loaded_after - loaded_before

        if newly_loaded:
            print("\nImplicitly loaded discovery sets:")
            for key in newly_loaded:
                target, prec, layout = key
                print(f"  - Target: {target.name}, Precision: {prec}, Layout: {layoutToStr(layout)}")
        else:
            print("\nNo new discovery sets were loaded")

        print(f"\nFound {len(configs)} configurations with {precision} precision:")
        for i, config in enumerate(configs, 1):
            print_kernel_config(config, f"Configuration {i}")

    except OSError as e:
        print(f"Error: Failed to load nvMatmulHeuristics library: {e}")
        print(f"Make sure the library exists at: {args.lib_path}")
        return 1
    except AssertionError:
        print("Error: Version mismatch or unsupported precision")
        return 1
    except RuntimeError as e:
        print(f"Error: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())


#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import sys

from nvMatmulHeuristics import (
    NvMatmulHeuristicsInterfaceEx,
    NvMatmulHeuristicsTarget,
    NvMatmulHeuristicsFlags,
    NvMatmulHeuristicsMatmulLayout,
    NvMatmulHeuristicsNvidiaGpu,
    layoutToStr
)

def parse_args():
    parser = argparse.ArgumentParser(
        description='Get and display GEMM configurations using NvMatmulHeuristicsInterfaceEx'
    )

    parser.add_argument('-M', '--m-dim',
                        type=int,
                        required=True,
                        help='Output matrix height')

    parser.add_argument('-N', '--n-dim',
                        type=int,
                        required=True,
                        help='Output matrix width')

    parser.add_argument('-K', '--k-dim',
                        type=int,
                        required=True,
                        help='Reduced dimension')

    parser.add_argument('-B', '--batch-size',
                        type=int,
                        default=1,
                        help='Batch size (default: 1)')

    parser.add_argument('--gpu',
                        type=str,
                        choices=[gpu.name for gpu in NvMatmulHeuristicsNvidiaGpu if gpu.name != 'END'],
                        required=True,
                        help='Target GPU')

    parser.add_argument('--layout',
                        type=str,
                        choices=[layout.name for layout in NvMatmulHeuristicsMatmulLayout if layout.name != 'END'],
                        default='NN_ROW_MAJOR',
                        help='Matrix layout (default: NN_ROW_MAJOR)')

    parser.add_argument('--backend',
                        type=str,
                        choices=[backend.name for backend in NvMatmulHeuristicsTarget if backend.name != 'END'],
                        default='CUTLASS3',
                        help='Target backend (default: CUTLASS3)')

    parser.add_argument('--precision',
                        type=str,
                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--count',
                        type=int,
                        default=8,
                        help='Number of configurations to retrieve (default: 8)')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    parser.add_argument('--no-auto-load',
                        action='store_true',
                        help='Disable automatic loading of discovery sets')

    return parser.parse_args()

def print_kernel_config(config, title):
    """Print kernel configuration details."""
    kernel = config["kernel"]
    runtime = config["runtime"]
    print(f"\n{title}:")
    print(f"  Kernel: {kernel}")
    print(f"  Estimated runtime: {runtime * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nGetting {args.count} configurations for:")


                        default='HSH',
                        help='Precision string (e.g. HSS, TST) (default: HSH)')

    parser.add_argument('--count',
                        type=int,
                        default=8,
                        help='Number of configurations to retrieve (default: 8)')

    parser.add_argument('--lib-path',
                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    parser.add_argument('--no-auto-load',
                        action='store_true',
                        help='Disable automatic loading of discovery sets')

    return parser.parse_args()

def print_kernel_config(config, title):
    """Print kernel configuration details."""
    kernel = config["kernel"]
    runtime = config["runtime"]
    print(f"\n{title}:")
    print(f"  Kernel: {kernel}")
    print(f"  Estimated runtime: {runtime * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nGetting {args.count} configurations for:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"Batch size: {args.batch_size}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}")
    print(f"Auto-load discovery sets: {not args.no_auto_load}\n")

    # Initialize nvMatmulHeuristics interface


                        type=str,
                        default=None,
                        help='Path to nvMatmulHeuristics shared library (default: uses the library from the wheel)')

    parser.add_argument('--no-auto-load',
                        action='store_true',
                        help='Disable automatic loading of discovery sets')

    return parser.parse_args()

def print_kernel_config(config, title):
    """Print kernel configuration details."""
    kernel = config["kernel"]
    runtime = config["runtime"]
    print(f"\n{title}:")
    print(f"  Kernel: {kernel}")
    print(f"  Estimated runtime: {runtime * 1000:.6f} ms")

def main():
    args = parse_args()

    # Convert string arguments to enum values
    gpu = NvMatmulHeuristicsNvidiaGpu[args.gpu]
    layout = NvMatmulHeuristicsMatmulLayout[args.layout]
    backend = NvMatmulHeuristicsTarget[args.backend]

    print(f"\nGetting {args.count} configurations for:")
    print(f"Problem size: M={args.m_dim}, N={args.n_dim}, K={args.k_dim}")
    print(f"Batch size: {args.batch_size}")
    print(f"GPU: {gpu.name}")
    print(f"Layout: {layoutToStr(layout)}")
    print(f"Backend: {backend.name}")
    print(f"Precision: {args.precision}")
    print(f"Auto-load discovery sets: {not args.no_auto_load}\n")

    # Initialize nvMatmulHeuristics interface
    try:
        # Create interface with the GPU
        nvMatmulHeuristics = NvMatmulHeuristicsInterfaceEx(
            path=args.lib_path,
            backend=backend,
            flags=NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING,
            load_discovery_implicitly=not args.no_auto_load,
            gpu=gpu
        )

        # Create problem object
        problem = nvMatmulHeuristics.makeNvMatmulHeuristicsProblem(
            args.m_dim, args.n_dim, args.k_dim, layout, args.batch_size
        )

        # Set precision
        precision = args.precision
        print(f"\nGetting configurations with precision ({precision})...")

        # Track loaded discovery sets before the call
        loaded_before = set(nvMatmulHeuristics._loaded_discovery_sets.keys())

        # Get configurations
        configs = nvMatmulHeuristics.get(problem, args.count, precision=precision)

        # Check which discovery sets were loaded
        loaded_after = set(nvMatmulHeuristics._loaded_discovery_sets.keys())
        newly_loaded = loaded_after - loaded_before

        if newly_loaded:
            print("\nImplicitly loaded discovery sets:")
            for key in newly_loaded:
                target, prec, layout = key
                print(f"  - Target: {target.name}, Precision: {prec}, Layout: {layoutToStr(layout)}")
        else:
            print("\nNo new discovery sets were loaded")

        print(f"\nFound {len(configs)} configurations with {precision} precision:")
        for i, config in enumerate(configs, 1):
            print_kernel_config(config, f"Configuration {i}")

    except OSError as e:
        print(f"Error: Failed to load nvMatmulHeuristics library: {e}")
        print(f"Make sure the library exists at: {args.lib_path}")
        return 1
    except AssertionError:
        print("Error: Version mismatch or unsupported precision")
        return 1
    except RuntimeError as e:
        print(f"Error: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())


# SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import time
import cv2
import torch
import ctypes
import numpy as np
import pandas as pd
import torchvision
from tabulate import tabulate
from ctypes import c_int, c_int16, c_ubyte, POINTER, c_void_p, sizeof

# Settings
INPUT_IMAGE = "Teapot.jpg"
OUTPUT_DIR = "Teapot_resolutions"
WARMUP_ITERATIONS = 5
MEASURE_ITERATIONS = 1000
THRESH_WEAK = 72
THRESH_STRONG = 256

RESOLUTIONS = [
    (320, 180, "320x180"),
    (640, 360, "640x360"),
    (800, 600, "800x600"),
    (1280, 720, "1280x720"),
    (1920, 1080, "1920x1080"),
    (2560, 1440, "2560x1440"),
    (3840, 2160, "3840x2160"),
    (5120, 2880, "5120x2880"),
]

os.makedirs(OUTPUT_DIR, exist_ok=True)

class CannyEdgeDetector:
    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()

    def _setup_functions(self):
        self.get_buffer_size = self.npp_lib.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [self.NppiSize, POINTER(c_int)]

        self.canny_func = self.npp_lib.nppiFilterCannyBorder_8u_C1R_Ctx
        self.canny_func.restype = c_int
        self.canny_func.argtypes = [
            POINTER(c_ubyte), c_int, self.NppiSize, self.NppiPoint,
            POINTER(c_ubyte), c_int, self.NppiSize,
            c_int, c_int, c_int16, c_int16,
            c_int, c_int, c_void_p, self.NppStreamContext
        ]

    def __call__(self, img_tensor, low_thresh, high_thresh):
        if img_tensor.dtype != torch.uint8:
            img_tensor = (img_tensor * 255).byte()
        if not img_tensor.is_cuda:
            img_tensor = img_tensor.cuda()

        h, w = img_tensor.shape
        output = torch.empty_like(img_tensor)
        scratch_size = c_int()
        self.get_buffer_size(self.NppiSize(w, h), ctypes.byref(scratch_size))
        scratch_buffer = torch.empty(scratch_size.value, dtype=torch.uint8, device='cuda')

        stream_ctx = self.NppStreamContext()
        stream_ctx.hStream = c_void_p(torch.cuda.current_stream().cuda_stream)

        status = self.canny_func(
            ctypes.cast(img_tensor.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte), self.NppiSize(w, h), self.NppiPoint(0, 0),
            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte), self.NppiSize(w, h),
            0, 200, c_int16(low_thresh), c_int16(high_thresh),
            2, 2,
            ctypes.cast(scratch_buffer.data_ptr(), c_void_p), stream_ctx
        )

        if status != 0:
            raise RuntimeError(f"NPP Canny edge detection failed with status {status}")

        return output

# Load input image
img = cv2.imread(INPUT_IMAGE)
if img is None:
    raise FileNotFoundError(f"Image not found: {INPUT_IMAGE}")

if not torch.cuda.is_available():
    raise RuntimeError("CUDA device not available!")

print(f"Input image size: {img.shape}")
print(f"Running benchmark... {MEASURE_ITERATIONS} iterations")

detector = CannyEdgeDetector()
results = []

for width, height, label in RESOLUTIONS:
    print(f"\nProcessing {label}")
    resized = cv2.resize(img, (width, height))
    img_tensor = torch.from_numpy(resized).cuda().permute(2, 0, 1)
    img_tensor = torchvision.transforms.functional.rgb_to_grayscale(img_tensor).squeeze(0)

    for _ in range(WARMUP_ITERATIONS):
        detector(img_tensor, THRESH_WEAK, THRESH_STRONG)

    timings = []
    for _ in range(MEASURE_ITERATIONS):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        detector(img_tensor, THRESH_WEAK, THRESH_STRONG)
        end.record()
        torch.cuda.synchronize()
        timings.append(start.elapsed_time(end))

    output_image = detector(img_tensor, THRESH_WEAK, THRESH_STRONG).cpu().numpy()
    cv2.imwrite(f"{OUTPUT_DIR}/out_npp_{label}.png", output_image)

    results.append({
        "Resolution": label,
        "Megapixels": (width * height) / 1_000_000,
        "NPP Time (ms)": np.mean(timings)
    })

# Save performance summary
df = pd.DataFrame(results)
print("\n--- Performance Summary ---")
print(tabulate(df, headers='keys', tablefmt='pretty', floatfmt='.3f'))

csv_path = os.path.join(OUTPUT_DIR, "performance_results.csv")
df.to_csv(csv_path, index=False)
print(f"\nResults saved to {csv_path}")


CannyEdgeDetector

INPUT_IMAGE = "Teapot.jpg"
OUTPUT_DIR = "Teapot_resolutions"
WARMUP_ITERATIONS = 5
MEASURE_ITERATIONS = 1000
THRESH_WEAK = 72
THRESH_STRONG = 256

RESOLUTIONS = [
    (320, 180, "320x180"),
    (640, 360, "640x360"),
    (800, 600, "800x600"),
    (1280, 720, "1280x720"),
    (1920, 1080, "1920x1080"),
    (2560, 1440, "2560x1440"),
    (3840, 2160, "3840x2160"),
    (5120, 2880, "5120x2880"),
]

os.makedirs(OUTPUT_DIR, exist_ok=True)

class CannyEdgeDetector:
    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()

    def _setup_functions(self):
        self.get_buffer_size = self.npp_lib.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [self.NppiSize, POINTER(c_int)]

        self.canny_func = self.npp_lib.nppiFilterCannyBorder_8u_C1R_Ctx
        self.canny_func.restype = c_int
        self.canny_func.argtypes = [
            POINTER(c_ubyte), c_int, self.NppiSize, self.NppiPoint,
            POINTER(c_ubyte), c_int, self.NppiSize,
            c_int, c_int, c_int16, c_int16,
            c_int, c_int, c_void_p, self.NppStreamContext
        ]

    def __call__(self, img_tensor, low_thresh, high_thresh):
        if img_tensor.dtype != torch.uint8:
            img_tensor = (img_tensor * 255).byte()
        if not img_tensor.is_cuda:
            img_tensor = img_tensor.cuda()

        h, w = img_tensor.shape
        output = torch.empty_like(img_tensor)
        scratch_size = c_int()
        self.get_buffer_size(self.NppiSize(w, h), ctypes.byref(scratch_size))
        scratch_buffer = torch.empty(scratch_size.value, dtype=torch.uint8, device='cuda')

        stream_ctx = self.NppStreamContext()
        stream_ctx.hStream = c_void_p(torch.cuda.current_stream().cuda_stream)

        status = self.canny_func(
            ctypes.cast(img_tensor.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte), self.NppiSize(w, h), self.NppiPoint(0, 0),
            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte), self.NppiSize(w, h),
            0, 200, c_int16(low_thresh), c_int16(high_thresh),
            2, 2,
            ctypes.cast(scratch_buffer.data_ptr(), c_void_p), stream_ctx
        )

        if status != 0:
            raise RuntimeError(f"NPP Canny edge detection failed with status {status}")

        return output

# Load input image
img = cv2.imread(INPUT_IMAGE)
if img is None:
    raise FileNotFoundError(f"Image not found: {INPUT_IMAGE}")

if not torch.cuda.is_available():
    raise RuntimeError("CUDA device not available!")

print(f"Input image size: {img.shape}")
print(f"Running benchmark... {MEASURE_ITERATIONS} iterations")

detector = CannyEdgeDetector()
results = []

for width, height, label in RESOLUTIONS:
    print(f"\nProcessing {label}")
    resized = cv2.resize(img, (width, height))
    img_tensor = torch.from_numpy(resized).cuda().permute(2, 0, 1)
    img_tensor = torchvision.transforms.functional.rgb_to_grayscale(img_tensor).squeeze(0)


NppiSize

OUTPUT_DIR = "Teapot_resolutions"
WARMUP_ITERATIONS = 5
MEASURE_ITERATIONS = 1000
THRESH_WEAK = 72
THRESH_STRONG = 256

RESOLUTIONS = [
    (320, 180, "320x180"),
    (640, 360, "640x360"),
    (800, 600, "800x600"),
    (1280, 720, "1280x720"),
    (1920, 1080, "1920x1080"),
    (2560, 1440, "2560x1440"),
    (3840, 2160, "3840x2160"),
    (5120, 2880, "5120x2880"),
]

os.makedirs(OUTPUT_DIR, exist_ok=True)

class CannyEdgeDetector:
    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()


NppiPoint

THRESH_WEAK = 72
THRESH_STRONG = 256

RESOLUTIONS = [
    (320, 180, "320x180"),
    (640, 360, "640x360"),
    (800, 600, "800x600"),
    (1280, 720, "1280x720"),
    (1920, 1080, "1920x1080"),
    (2560, 1440, "2560x1440"),
    (3840, 2160, "3840x2160"),
    (5120, 2880, "5120x2880"),
]

os.makedirs(OUTPUT_DIR, exist_ok=True)

class CannyEdgeDetector:
    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()

    def _setup_functions(self):
        self.get_buffer_size = self.npp_lib.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int


NppStreamContext

RESOLUTIONS = [
    (320, 180, "320x180"),
    (640, 360, "640x360"),
    (800, 600, "800x600"),
    (1280, 720, "1280x720"),
    (1920, 1080, "1920x1080"),
    (2560, 1440, "2560x1440"),
    (3840, 2160, "3840x2160"),
    (5120, 2880, "5120x2880"),
]

os.makedirs(OUTPUT_DIR, exist_ok=True)

class CannyEdgeDetector:
    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()

    def _setup_functions(self):
        self.get_buffer_size = self.npp_lib.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [self.NppiSize, POINTER(c_int)]

        self.canny_func = self.npp_lib.nppiFilterCannyBorder_8u_C1R_Ctx
        self.canny_func.restype = c_int
        self.canny_func.argtypes = [
            POINTER(c_ubyte), c_int, self.NppiSize, self.NppiPoint,
            POINTER(c_ubyte), c_int, self.NppiSize,
            c_int, c_int, c_int16, c_int16,
            c_int, c_int, c_void_p, self.NppStreamContext
        ]


os.makedirs(OUTPUT_DIR, exist_ok=True)

class CannyEdgeDetector:
    class NppiSize(ctypes.Structure):
        _fields_ = [("width", c_int), ("height", c_int)]

    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()

    def _setup_functions(self):
        self.get_buffer_size = self.npp_lib.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [self.NppiSize, POINTER(c_int)]

        self.canny_func = self.npp_lib.nppiFilterCannyBorder_8u_C1R_Ctx
        self.canny_func.restype = c_int
        self.canny_func.argtypes = [
            POINTER(c_ubyte), c_int, self.NppiSize, self.NppiPoint,
            POINTER(c_ubyte), c_int, self.NppiSize,
            c_int, c_int, c_int16, c_int16,
            c_int, c_int, c_void_p, self.NppStreamContext
        ]

    def __call__(self, img_tensor, low_thresh, high_thresh):
        if img_tensor.dtype != torch.uint8:
            img_tensor = (img_tensor * 255).byte()
        if not img_tensor.is_cuda:
            img_tensor = img_tensor.cuda()


    class NppiPoint(ctypes.Structure):
        _fields_ = [("x", c_int), ("y", c_int)]

    class NppStreamContext(ctypes.Structure):
        _fields_ = [
            ("hStream", c_void_p), ("nCudaDeviceId", c_int),
            ("nMultiProcessorCount", c_int), ("nMaxThreadsPerMultiProcessor", c_int),
            ("nMaxThreadsPerBlock", c_int), ("nSharedMemPerBlock", c_int),
            ("nCudaDevAttrComputeCapabilityMajor", c_int),
            ("nCudaDevAttrComputeCapabilityMinor", c_int),
            ("nStreamFlags", c_int)
        ]

    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()

    def _setup_functions(self):
        self.get_buffer_size = self.npp_lib.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [self.NppiSize, POINTER(c_int)]

        self.canny_func = self.npp_lib.nppiFilterCannyBorder_8u_C1R_Ctx
        self.canny_func.restype = c_int
        self.canny_func.argtypes = [
            POINTER(c_ubyte), c_int, self.NppiSize, self.NppiPoint,
            POINTER(c_ubyte), c_int, self.NppiSize,
            c_int, c_int, c_int16, c_int16,
            c_int, c_int, c_void_p, self.NppStreamContext
        ]

    def __call__(self, img_tensor, low_thresh, high_thresh):
        if img_tensor.dtype != torch.uint8:
            img_tensor = (img_tensor * 255).byte()
        if not img_tensor.is_cuda:
            img_tensor = img_tensor.cuda()

        h, w = img_tensor.shape
        output = torch.empty_like(img_tensor)
        scratch_size = c_int()
        self.get_buffer_size(self.NppiSize(w, h), ctypes.byref(scratch_size))
        scratch_buffer = torch.empty(scratch_size.value, dtype=torch.uint8, device='cuda')

        stream_ctx = self.NppStreamContext()
        stream_ctx.hStream = c_void_p(torch.cuda.current_stream().cuda_stream)

        status = self.canny_func(
            ctypes.cast(img_tensor.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte), self.NppiSize(w, h), self.NppiPoint(0, 0),
            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),


    def __init__(self):
        #npp_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\npp_plus_if64_12.dll" # for windows user
        #self.npp_lib = ctypes.cdll.LoadLibrary(npp_path)
        self.npp_lib = ctypes.CDLL('libnpp_plus_if.so')  # for linux user
        self._setup_functions()

    def _setup_functions(self):
        self.get_buffer_size = self.npp_lib.nppiFilterCannyBorderGetBufferSize
        self.get_buffer_size.restype = c_int
        self.get_buffer_size.argtypes = [self.NppiSize, POINTER(c_int)]

        self.canny_func = self.npp_lib.nppiFilterCannyBorder_8u_C1R_Ctx
        self.canny_func.restype = c_int
        self.canny_func.argtypes = [
            POINTER(c_ubyte), c_int, self.NppiSize, self.NppiPoint,
            POINTER(c_ubyte), c_int, self.NppiSize,
            c_int, c_int, c_int16, c_int16,
            c_int, c_int, c_void_p, self.NppStreamContext
        ]

    def __call__(self, img_tensor, low_thresh, high_thresh):
        if img_tensor.dtype != torch.uint8:
            img_tensor = (img_tensor * 255).byte()
        if not img_tensor.is_cuda:
            img_tensor = img_tensor.cuda()

        h, w = img_tensor.shape
        output = torch.empty_like(img_tensor)
        scratch_size = c_int()
        self.get_buffer_size(self.NppiSize(w, h), ctypes.byref(scratch_size))
        scratch_buffer = torch.empty(scratch_size.value, dtype=torch.uint8, device='cuda')

        stream_ctx = self.NppStreamContext()
        stream_ctx.hStream = c_void_p(torch.cuda.current_stream().cuda_stream)

        status = self.canny_func(
            ctypes.cast(img_tensor.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte), self.NppiSize(w, h), self.NppiPoint(0, 0),
            ctypes.cast(output.data_ptr(), POINTER(c_ubyte)),
            w * sizeof(c_ubyte), self.NppiSize(w, h),
            0, 200, c_int16(low_thresh), c_int16(high_thresh),
            2, 2,
            ctypes.cast(scratch_buffer.data_ptr(), c_void_p), stream_ctx
        )

        if status != 0:
            raise RuntimeError(f"NPP Canny edge detection failed with status {status}")

        return output

# Load input image
img = cv2.imread(INPUT_IMAGE)
if img is None:
    raise FileNotFoundError(f"Image not found: {INPUT_IMAGE}")

if not torch.cuda.is_available():
    raise RuntimeError("CUDA device not available!")

print(f"Input image size: {img.shape}")
print(f"Running benchmark... {MEASURE_ITERATIONS} iterations")

detector = CannyEdgeDetector()
results = []

for width, height, label in RESOLUTIONS:
    print(f"\nProcessing {label}")
    resized = cv2.resize(img, (width, height))
    img_tensor = torch.from_numpy(resized).cuda().permute(2, 0, 1)
    img_tensor = torchvision.transforms.functional.rgb_to_grayscale(img_tensor).squeeze(0)


cuTENSOR.python.setup

#! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from setuptools import setup, find_packages

from cutensor.package_info import __version__
from cutensor.package_info import __package_name__
from cutensor.package_info import __homepage__
from cutensor.package_info import __download_url__
from cutensor.package_info import __description__
from cutensor.package_info import __license__

from cutensor.c_extensions import CustomExtension

setup(name=__package_name__,
      version=__version__,
      description=__description__,
      url=__homepage__,
      download_url=__download_url__,
      license=__license__,
      packages=find_packages(),
      ext_modules=CustomExtension.modules)


cuTENSOR.python.cutensor.common

# ! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

def normalize_subscript(subscript):
    if '->' in subscript:
        subscript = subscript.split('->')
        lhs = subscript[0]
        rhs = subscript[1]
    else:
        lhs = subscript
        rhs = ''.join(sorted([s for s in set(subscript) if s != ',' and subscript.count(s) == 1]))
    if '...' in lhs:
        raise RuntimeError('Elipsis is currently unsupported')
    return lhs + '->' + rhs, ',' in lhs