trme 0.1.0__tar.gz → 0.2.0__tar.gz

{trme-0.1.0 → trme-0.2.0}/Cargo.toml

@@ -7,7 +7,7 @@ authors = ["TRME Team"]
 readme = "README.md"
 
 [lib]
-name = "trme"
+name = "_trme_backend"
 crate-type = ["cdylib"]
 
 [dependencies]

{trme-0.1.0 → trme-0.2.0}/Dockerfile

@@ -10,7 +10,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     curl \
     git \
     libomp-dev \
-    patchelf \
     && rm -rf /var/lib/apt/lists/*
 
 # Install Rust (for Maturin)

{trme-0.1.0 → trme-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trme
-Version: 0.1.0
+Version: 0.2.0
 Requires-Dist: torch
 Requires-Dist: numpy
 Requires-Dist: scipy

trme-0.2.0/build.rs

@@ -0,0 +1,40 @@
+fn main() {
+    // 1. Instancia o compilador C++
+    let mut build = cc::Build::new();
+
+    // 2. Adiciona o arquivo fonte C++
+    // CERTIFIQUE-SE que o arquivo "rsr_gemm.cpp" está na mesma pasta que o Cargo.toml
+    build.file("rsr_gemm.cpp");
+
+    // Indica que é C++
+    build.cpp(true);
+
+    // 3. Configura Flags baseadas no Sistema Operacional
+    let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
+
+    if target_os == "windows" {
+        // --- Configuração WINDOWS (MSVC) ---
+        build.flag("/O2");       // Otimização máxima
+        build.flag("/EHsc");     // Tratamento de exceções padrão
+        // build.flag("/openmp"); // Descomente se seu C++ usar OpenMP (#include <omp.h>)
+    } else {
+        // --- Configuração LINUX/MAC ---
+        build.flag("-O3");
+        build.flag("-march=native");
+        build.flag("-fopenmp");
+
+        // Só linka essas libs se NÃO for Windows
+        println!("cargo:rustc-link-lib=gomp");
+        println!("cargo:rustc-link-lib=stdc++");
+    }
+
+    // 4. Compila a biblioteca estática "trme_core"
+    // Isso gera o arquivo .lib que o linker estava reclamando que não existia
+    build.compile("trme_core");
+
+    // 5. Avisa o Cargo para recompilar se o C++ mudar
+    println!("cargo:rerun-if-changed=rsr_gemm.cpp");
+
+    // Avisa onde procurar a lib (na pasta de build atual)
+    println!("cargo:rustc-link-search=native=.");
+}

{trme-0.1.0 → trme-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "trme"
-version = "0.1.0"
+version = "0.2.0"
 description = "Tensor-RNS-Multipole Engine (TRME) Framework"
 requires-python = ">=3.8"
 dependencies = [

{trme-0.1.0 → trme-0.2.0}/src/lib.rs

@@ -50,7 +50,7 @@ fn rsr_gemm(
 }
 
 #[pymodule]
-fn trme(_py: Python, m: &PyModule) -> PyResult<()> {
+fn _trme_backend(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(rsr_gemm, m)?)?;
     Ok(())
 }

trme-0.2.0/trme/__init__.py

@@ -0,0 +1,67 @@
+import torch
+import numpy as np
+import _trme_backend
+
+class TRMEMatmul(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weight, bias=None, block_size=4):
+        """
+        input: (N, K)
+        weight: (M, K)
+        bias: (M,) Optional
+        """
+        if isinstance(bias, int):
+            block_size = bias
+            bias = None
+
+        # Quantization (A -> [0..3], B -> Int8)
+        # Note: We do simple casting here for the prototype.
+        # Ideally calling C++ quantization logic.
+        A_cpu = input.detach().cpu().numpy()
+        A_q = np.clip(np.abs(A_cpu), 0, 3).astype(np.int8)
+
+        B_cpu = weight.detach().cpu().numpy().T
+        B_q = np.ascontiguousarray(np.clip(B_cpu, -127, 127).astype(np.int8))
+
+        N, K = A_q.shape
+        _, M = B_q.shape
+
+        # Flatten for Rust FFI
+        a_bytes = A_q.tobytes()
+        b_bytes = B_q.tobytes()
+
+        # Call Rust Backend
+        # rsr_gemm returns list[int] (flattened C)
+        c_flat = _trme_backend.rsr_gemm(a_bytes, b_bytes, N, K, M, int(block_size))
+
+        # Reshape and Convert to Tensor
+        C_q = np.array(c_flat, dtype=np.int32).reshape(N, M)
+        output = torch.from_numpy(C_q.astype(np.float32))
+
+        if bias is not None:
+             output += bias.cpu()
+
+        ctx.save_for_backward(input, weight, bias)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weight, bias = ctx.saved_tensors
+        grad_input = grad_weight = grad_bias = None
+
+        if ctx.needs_input_grad[0]:
+            grad_input = grad_output.matmul(weight)
+        if ctx.needs_input_grad[1]:
+            grad_weight = grad_output.t().matmul(input)
+        if bias is not None and ctx.needs_input_grad[2]:
+            grad_bias = grad_output.sum(0)
+
+        return grad_input, grad_weight, grad_bias, None
+
+def matmul(input, weight):
+    # Basic matmul interface
+    return TRMEMatmul.apply(input, weight)
+
+def linear(input, weight, bias=None, block_size=4):
+    # nn.Linear interface
+    return TRMEMatmul.apply(input, weight, bias, block_size)

trme-0.1.0/build.rs

@@ -1,40 +0,0 @@
-fn main() {
-    // 1. Instancia o compilador C++
-    let mut build = cc::Build::new();
-
-    // 2. Adiciona o arquivo fonte C++
-    // IMPORTANTE: Verifique se o nome do arquivo na sua pasta é exatamente este.
-    // Baseado no seu README, deve ser "rsr_gemm.cpp" ou "rsr_fused.cpp".
-    // Se tiver os dois, adicione outra linha: build.file("rsr_fused.cpp");
-    build.file("rsr_gemm.cpp");
-
-    // Indica que é C++ (não C puro)
-    build.cpp(true);
-
-    // 3. Configura Flags específicas por Sistema Operacional
-    let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
-
-    if target_os == "windows" {
-        // --- Configuração WINDOWS (MSVC) ---
-        build.flag("/O2");       // Otimização máxima
-        build.flag("/EHsc");     // Tratamento de exceções C++ padrão
-        build.flag("/openmp");   // OpenMP nativo do Visual Studio
-        // No Windows, não precisamos linkar stdc++ ou gomp manualmente
-    } else {
-        // --- Configuração LINUX/MAC ---
-        build.flag("-O3");
-        build.flag("-march=native");
-        build.flag("-fopenmp");
-
-        // No Linux, precisamos pedir para linkar as libs explicitamente no final
-        println!("cargo:rustc-link-lib=gomp");
-        println!("cargo:rustc-link-lib=stdc++");
-    }
-
-    // 4. Compila a biblioteca estática "trme_core.lib"
-    // Isso vai gerar o arquivo que estava faltando no erro anterior
-    build.compile("trme_core");
-
-    // 5. Instrui o Cargo a recompilar se o C++ mudar
-    println!("cargo:rerun-if-changed=rsr_gemm.cpp");
-}

trme-0.1.0/trme_torch.py

@@ -1,95 +0,0 @@
-import torch
-import ctypes
-import numpy as np
-import os
-
-# Load Library
-_lib_path = os.path.abspath('./librsr.so')
-if os.path.exists(_lib_path):
-    _lib = ctypes.CDLL(_lib_path)
-
-    _lib.rsr_gemm_dispatch.argtypes = [
-        np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
-        np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
-        np.ctypeslib.ndpointer(dtype=np.int32, ndim=2, flags='C_CONTIGUOUS'),
-        ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int
-    ]
-else:
-    print("Warning: librsr.so not found. TRME Ops will fail.")
-    _lib = None
-
-class TRMEMatmul(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, input, weight, bias=None, block_size=4):
-        """
-        input: (N, K)
-        weight: (M, K)
-        bias: (M,) Optional
-        """
-        # Check if block_size is passed as positional arg
-        if isinstance(bias, int):
-            block_size = bias
-            bias = None
-
-        # Quantization
-        A_cpu = input.detach().cpu().numpy()
-        A_q = np.clip(np.abs(A_cpu), 0, 3).astype(np.int8)
-
-        B_cpu = weight.detach().cpu().numpy().T
-        B_q = np.ascontiguousarray(np.clip(B_cpu, -127, 127).astype(np.int8))
-
-        N, K = A_q.shape
-        _, M = B_q.shape
-
-        C_q = np.zeros((N, M), dtype=np.int32)
-
-        # Handle Bias (Quantized int32)
-        bias_ptr = None
-        if bias is not None:
-            bias_q = bias.detach().cpu().numpy().astype(np.int32)
-            # Check contiguity
-            if not bias_q.flags['C_CONTIGUOUS']:
-                bias_q = np.ascontiguousarray(bias_q)
-            bias_ptr = bias_q.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
-        else:
-            bias_ptr = ctypes.POINTER(ctypes.c_int32)()
-
-        # Call Kernel
-        if _lib:
-            # Update signature dynamically if needed
-            # For V5, we enable the real call
-            if not hasattr(_lib, 'rsr_gemm_dispatch_configured'):
-                 _lib.rsr_gemm_dispatch.argtypes = [
-                    np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
-                    np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
-                    np.ctypeslib.ndpointer(dtype=np.int32, ndim=2, flags='C_CONTIGUOUS'),
-                    ctypes.POINTER(ctypes.c_int32),
-                    ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int
-                ]
-                 _lib.rsr_gemm_dispatch_configured = True
-
-            _lib.rsr_gemm_dispatch(A_q, B_q, C_q, bias_ptr, N, K, M, int(block_size))
-
-        output = torch.from_numpy(C_q.astype(np.float32))
-        # Note: Bias is added in C++ now, no need to add in Python unless fallback
-
-        ctx.save_for_backward(input, weight, bias)
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        # Backward pass is standard GEMM (simulated here)
-        input, weight, bias = ctx.saved_tensors
-        grad_input = grad_weight = grad_bias = None
-
-        if ctx.needs_input_grad[0]:
-            grad_input = grad_output.matmul(weight)
-        if ctx.needs_input_grad[1]:
-            grad_weight = grad_output.t().matmul(input)
-        if bias is not None and ctx.needs_input_grad[2]:
-            grad_bias = grad_output.sum(0)
-
-        return grad_input, grad_weight, grad_bias, None
-
-def matmul(input, weight, block_size=4):
-    return TRMEMatmul.apply(input, weight, block_size)