tempstore

benchmarks/ks/auto_bench.py

@@ -10,7 +10,10 @@
 from pathlib import Path
 
 import torch
-
+try:
+      import torch_npu
+except ImportError:
+      pass
 
 class KsCompareError(Exception):
     pass
@@ -423,7 +426,7 @@ def _first_input_device(inputs):
 def _detect_target_device(model, model_new, v0_inputs, v1_inputs):
     """Pick a non-CPU device from models/inputs, or auto-detect one.
 
-    Priority: model device > input device > auto-detect (cuda → npu).
+    Priority: model device > input device > auto-detect (cuda 闂佹剚鍋撻幏锟� npu).
     Raises KsCompareError if no accelerator is available.
     """
     for m in (model, model_new):

dlblas/kernels/ks_competition/QY2026_Autumn_liujie925_Task01ComputeContactProb_Ascend.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-7 -*-
+import torch_npu
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, min_bin: float = 2.3125, max_bin: float = 21.6875, no_bins: int = 64, thres: float = 8.0):
+        super().__init__()
+        no_bins = int(no_bins)
+        self.no_bins = no_bins
+        edges = torch.linspace(min_bin, max_bin, no_bins + 1)
+        bin_centers = 0.5 * (edges[:-1] + edges[1:])
+        thres_idx_val = int((bin_centers < thres).sum().item())
+        self.register_buffer("bin_centers", bin_centers)
+        self.register_buffer("thres_idx", torch.tensor(thres_idx_val, dtype=torch.long))
+
+    def forward(self, distogram_logits: torch.Tensor) -> torch.Tensor:
+        N = distogram_logits.shape[0]
+        # 取上三角索引（不含对角线，对角线也可根据需求保留）
+        triu_idx = torch.triu_indices(N, N, offset=0, device=distogram_logits.device)
+        logits_triu = distogram_logits[triu_idx[0], triu_idx[1], :]   # [M, 64]
+
+        prob_triu = torch.softmax(logits_triu, dim=-1)
+        contact_triu = prob_triu[:, :self.thres_idx].sum(dim=-1)
+
+        # 构造对称矩阵（对角线直接赋值，非对角线同时赋值对称元素）
+        contact_prob = torch.zeros(N, N, device=distogram_logits.device, dtype=contact_triu.dtype)
+        contact_prob[triu_idx[0], triu_idx[1]] = contact_triu
+        contact_prob = contact_prob + contact_prob.T
+        # 如果上三角包含对角线，对角线加了两次，需减半
+        if (triu_idx[0] == triu_idx[1]).any():
+            diag_mask = triu_idx[0] == triu_idx[1]
+            contact_prob[triu_idx[0][diag_mask], triu_idx[1][diag_mask]] = contact_triu[diag_mask]
+        # prob = torch.softmax(distogram_logits, dim=-1)
+        # contact_prob = prob[..., :self.thres_idx.item()].sum(dim=-1)
+
+        return contact_prob
+
+# Hyperparameters
+N_TOKEN = 256
+NO_BINS = 64
+MIN_BIN = 2.3125
+MAX_BIN = 21.6875
+THRES = 8.0
+
+def get_inputs():
+    device = 'npu:0'
+    torch.manual_seed(42)
+    logits = torch.randn(N_TOKEN, N_TOKEN, NO_BINS, device=device)
+    return [logits]
+
+def get_init_inputs():
+    return [MIN_BIN, MAX_BIN, NO_BINS, THRES]
+
+if __name__ == "__main__":
+    torch_npu.npu.set_device(0)
+    device = torch.device("npu:0")
+
+    raw_model = Model(*get_init_inputs()).to(device)
+
+    inputs = get_inputs()
+
+    # 关键替换：用trace而非script，不读取源码
+    traced_model = torch.jit.trace(raw_model, inputs)
+
+    # 混合精度推理
+    with torch.npu.amp.autocast(dtype=torch.float16):
+        res = traced_model(*inputs)
+    print("输出shape:", res.shape)
+    print(res)

dlblas/kernels/ks_competition/torch/ComputeContactProb.py

@@ -0,0 +1,84 @@
+"""
+Compute Contact Probability (distogram logits -> contact probability)
+
+From: protenix/model/sample_confidence.py:compute_contact_prob
+"""
+
+import torch
+import torch.nn as nn
+
+
+def get_bin_centers(min_bin: float, max_bin: float, no_bins: int) -> torch.Tensor:
+    """
+    distogram bins centers锛堝父瑙佸仛娉曪細绾挎€х瓑闂撮殧锛�
+    """
+    edges = torch.linspace(min_bin, max_bin, no_bins + 1)
+    centers = 0.5 * (edges[:-1] + edges[1:])
+    return centers
+
+
+def compute_contact_prob(
+    distogram_logits: torch.Tensor,
+    min_bin: float,
+    max_bin: float,
+    no_bins: int,
+    thres: float = 8.0,
+) -> torch.Tensor:
+    """
+    Args:
+        distogram_logits: [N_token, N_token, no_bins]
+    Returns:
+        contact_prob: [N_token, N_token]
+    """
+    distogram_prob = torch.softmax(distogram_logits, dim=-1)
+    bins = get_bin_centers(min_bin, max_bin, no_bins).to(distogram_logits.device)
+    thres_idx = int((bins < thres).sum().item())
+    return distogram_prob[..., :thres_idx].sum(dim=-1)
+
+
+class Model(nn.Module):
+    def __init__(self, min_bin: float = 2.3125, max_bin: float = 21.6875, no_bins: int = 64, thres: float = 8.0):
+        super().__init__()
+        self.min_bin = float(min_bin)
+        self.max_bin = float(max_bin)
+        self.no_bins = int(no_bins)
+        self.thres = float(thres)
+
+    def forward(self, distogram_logits: torch.Tensor) -> torch.Tensor:
+        return compute_contact_prob(
+            distogram_logits=distogram_logits,
+            min_bin=self.min_bin,
+            max_bin=self.max_bin,
+            no_bins=self.no_bins,
+            thres=self.thres,
+        )
+
+
+# ==========================================
+# Hyperparameters & Data Generation
+# ==========================================
+
+N_TOKEN = 256
+NO_BINS = 64
+MIN_BIN = 2.3125
+MAX_BIN = 21.6875
+THRES = 8.0
+
+
+def get_inputs():
+    device = 'npu'
+    torch.manual_seed(42)
+
+    distogram_logits = torch.randn(N_TOKEN, N_TOKEN, NO_BINS, device=device)
+
+    return [distogram_logits]
+
+
+def get_init_inputs():
+    return [MIN_BIN, MAX_BIN, NO_BINS, THRES]
+
+if __name__ == "__main__":
+    torch.set_default_device("npu")
+    model = Model(*get_init_inputs())
+    inputs = get_inputs()
+    print(model(*inputs))

dlblas/kernels/ks_competition/torch/ComputeContactProb_v1.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+import torch_npu
+import torch
+import torch.nn as nn
+
+class ModelNew(nn.Module):
+    def __init__(self, min_bin: float = 2.3125, max_bin: float = 21.6875, no_bins: int = 64, thres: float = 8.0):
+        super().__init__()
+        self.no_bins = int(no_bins)
+
+        # 一次性预计算bin中心与阈值下标，全局常量
+        edges = torch.linspace(min_bin, max_bin, self.no_bins + 1)
+        bin_centers = 0.5 * (edges[:-1] + edges[1:])
+        self.thres_idx = int((bin_centers < thres).sum().item())
+        self.register_buffer("bin_centers", bin_centers)
+
+    def forward(self, distogram_logits: torch.Tensor) -> torch.Tensor:
+        # 极简内联计算，无多余中间函数
+        prob = torch.softmax(distogram_logits, dim=-1)
+        contact_prob = prob[..., :self.thres_idx].sum(dim=-1)
+        return contact_prob
+
+# Hyperparameters
+N_TOKEN = 256
+NO_BINS = 64
+MIN_BIN = 2.3125
+MAX_BIN = 21.6875
+THRES = 8.0
+
+def get_inputs():
+#    evice = 'npu:0'
+# logits = torch.randn(N_TOKEN, N_TOKEN, NO_BINS, device=device)
+    torch.manual_seed(42)
+    logits = torch.randn(N_TOKEN, N_TOKEN, NO_BINS)
+    return [logits]
+
+def get_init_inputs():
+    return [MIN_BIN, MAX_BIN, NO_BINS, THRES]
+
+if __name__ == "__main__":
+    torch_npu.npu.set_device(0)
+    device = torch.device("npu:0")
+
+    model = ModelNew(*get_init_inputs()).to(device)
+    # JIT编译加速
+    model = torch.jit.script(model)
+
+    inputs = get_inputs()
+    # fp16混合精度推理
+    with torch.npu.amp.autocast(dtype=torch.float16):
+        res = model(*inputs)
+    print(res.shape)
+

dlblas/kernels/ks_competition/torch/Grid.py

@@ -0,0 +1,107 @@
+import torch
+
+class Model(torch.nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, pos, size, start = None, end = None):
+        if pos.dim() != 2:
+          raise ValueError(f"pos should be 2-dimensional, got {pos.dim()}-dimensional")
+
+        if size.dim() != 1:
+          raise ValueError(f"size should be 1-dimensional, got {size.dim()}-dimensional")
+
+        if pos.size(1) != size.size(0):
+          raise ValueError(f"Dimension mismatch: pos has {pos.size(1)} dimensions, "
+                           f"but size has {size.size(0)} dimensions")
+
+        N, D = pos.shape
+        device = pos.device
+
+        # 处理可选参数
+        if start is None:
+            start = torch.zeros(D, device=device)
+        else:
+            if start.dim() != 1 or start.size(0) != D:
+                raise ValueError(f"start should have shape [{D}], got {start.shape}")
+
+        if end is None:
+          # 如果没有提供end，则使用点的最大坐标
+          end = torch.max(pos, dim=0)[0] + size
+        else:
+          if end.dim() != 1 or end.size(0) != D:
+              raise ValueError(f"end should have shape [{D}], got {end.shape}")
+
+        # 将点坐标转换为网格索引
+        grid_indices = ((pos - start.unsqueeze(0)) / size.unsqueeze(0)).long()
+
+        # 确保网格索引在有效范围内
+        grid_indices = torch.clamp(grid_indices, min=0)
+
+        # 计算每个维度上的网格数量
+        grid_counts = ((end - start) / size).long() + 1
+
+        # 计算每个点的唯一网格ID
+        cluster_ids = torch.zeros(N, dtype=torch.long, device=device)
+
+        # 使用多维网格索引计算唯一ID
+        for d in range(D):
+          if d == 0:
+              cluster_ids = grid_indices[:, d]
+          else:
+              cluster_ids = cluster_ids * grid_counts[d] + grid_indices[:, d]
+
+        # 重新映射聚类ID为连续的整数
+        unique_ids, inverse_indices = torch.unique(cluster_ids, return_inverse=True)
+
+        return inverse_indices
+
+def get_inputs():
+    # pos = torch.tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
+    # size = torch.tensor([5, 5])
+    # end = torch.tensor([19, 19])
+    pos = torch.tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]], dtype=torch.float32)
+    size = torch.tensor([5, 5], dtype=torch.float32)
+    end = torch.tensor([19, 19], dtype=torch.float32)
+    N, D = 100000, 2
+    torch.manual_seed(42)
+    pos = torch.rand(N, D, dtype=torch.float32) * 100.0
+    size = torch.full((D,), 5.0, dtype=torch.float32)
+    end = torch.full((D,), 100.0, dtype=torch.float32)
+
+    return [pos, size, end]
+
+def get_init_inputs():
+    return []
+
+if __name__ == "__main__":
+    try:
+        import torch_npu
+        use_npu = True
+        dev = torch.device("npu:0")
+        print("Running on NPU device")
+    except ImportError:
+        use_npu = False
+        dev = torch.device("cpu")
+        print("torch_npu not found, running on CPU")
+
+    model = Model().to(dev)
+    model.eval()
+
+    inputs = get_inputs()
+    pos, size, end = [x.to(dev) for x in inputs]
+
+    with torch.no_grad():
+        out = model(pos, size, end=end)
+
+    print("Output result:")
+    print(out.cpu())
+
+    # 核对CPU基准结果
+    pos_cpu, size_cpu, end_cpu = get_inputs()
+    cpu_model = Model().cpu()
+    with torch.no_grad():
+        out_cpu = cpu_model(pos_cpu, size_cpu, end=end_cpu)
+
+    is_equal = torch.equal(out.cpu(), out_cpu)
+    print(f"\nNPU result matches CPU result: {is_equal}")

dlblas/kernels/ks_competition/torch/layer_norm.py

@@ -20,5 +20,5 @@ def get_init_inputs():
 
 
 torch.manual_seed(42)
-out = Model(*get_init_inputs()).forward(*get_inputs())
+out = Model(*get_init_inputs()).forward(*get_inputs()).to("npu")
 print(out)