"git@developer.sourcefind.cn:zhaoyu6/sglang.git" did not exist on "198974cd1a805a7fab2d81fe9e6b5fbd73d03fb8"
bm_rasterize_meshes.py 3.55 KB
Newer Older
1
# Copyright (c) Meta Platforms, Inc. and affiliates.
Patrick Labatut's avatar
Patrick Labatut committed
2
3
4
5
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
facebook-github-bot's avatar
facebook-github-bot committed
6

7
import os
facebook-github-bot's avatar
facebook-github-bot committed
8
from itertools import product
9

facebook-github-bot's avatar
facebook-github-bot committed
10
11
import torch
from fvcore.common.benchmark import benchmark
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
12
from tests.test_rasterize_meshes import TestRasterizeMeshes
facebook-github-bot's avatar
facebook-github-bot committed
13

14
15
BM_RASTERIZE_MESHES_N_THREADS = os.getenv("BM_RASTERIZE_MESHES_N_THREADS", 1)
torch.set_num_threads(int(BM_RASTERIZE_MESHES_N_THREADS))
16

facebook-github-bot's avatar
facebook-github-bot committed
17
18
19
20
21
# ico levels:
# 0: (12 verts, 20 faces)
# 1: (42 verts, 80 faces)
# 3: (642 verts, 1280 faces)
# 4: (2562 verts, 5120 faces)
Nikhila Ravi's avatar
Nikhila Ravi committed
22
23
# 5: (10242 verts, 20480 faces)
# 6: (40962 verts, 81920 faces)
facebook-github-bot's avatar
facebook-github-bot committed
24
25
26
27
28
29
30
31
32


def bm_rasterize_meshes() -> None:
    kwargs_list = [
        {
            "num_meshes": 1,
            "ico_level": 0,
            "image_size": 10,  # very slow with large image size
            "blur_radius": 0.0,
Nikhila Ravi's avatar
Nikhila Ravi committed
33
            "faces_per_pixel": 3,
facebook-github-bot's avatar
facebook-github-bot committed
34
35
36
37
38
39
40
41
42
43
44
45
        }
    ]
    benchmark(
        TestRasterizeMeshes.rasterize_meshes_python_with_init,
        "RASTERIZE_MESHES",
        kwargs_list,
        warmup_iters=1,
    )

    kwargs_list = []
    num_meshes = [1]
    ico_level = [1]
46
    image_size = [64, 128, 512]
Nikhila Ravi's avatar
Nikhila Ravi committed
47
48
49
    blur = [1e-6]
    faces_per_pixel = [3, 50]
    test_cases = product(num_meshes, ico_level, image_size, blur, faces_per_pixel)
facebook-github-bot's avatar
facebook-github-bot committed
50
    for case in test_cases:
Nikhila Ravi's avatar
Nikhila Ravi committed
51
        n, ic, im, b, f = case
facebook-github-bot's avatar
facebook-github-bot committed
52
        kwargs_list.append(
Nikhila Ravi's avatar
Nikhila Ravi committed
53
54
55
56
57
58
59
            {
                "num_meshes": n,
                "ico_level": ic,
                "image_size": im,
                "blur_radius": b,
                "faces_per_pixel": f,
            }
facebook-github-bot's avatar
facebook-github-bot committed
60
61
62
63
64
65
66
67
68
69
        )
    benchmark(
        TestRasterizeMeshes.rasterize_meshes_cpu_with_init,
        "RASTERIZE_MESHES",
        kwargs_list,
        warmup_iters=1,
    )

    if torch.cuda.is_available():
        kwargs_list = []
Nikhila Ravi's avatar
Nikhila Ravi committed
70
71
        num_meshes = [8, 16]
        ico_level = [4, 5, 6]
72
73
        # Square and non square cases
        image_size = [64, 128, 512, (512, 256), (256, 512)]
Nikhila Ravi's avatar
Nikhila Ravi committed
74
        blur = [1e-6]
75
        faces_per_pixel = [40]
Nikhila Ravi's avatar
Nikhila Ravi committed
76
77
        test_cases = product(num_meshes, ico_level, image_size, blur, faces_per_pixel)

facebook-github-bot's avatar
facebook-github-bot committed
78
        for case in test_cases:
Nikhila Ravi's avatar
Nikhila Ravi committed
79
            n, ic, im, b, f = case
facebook-github-bot's avatar
facebook-github-bot committed
80
81
82
83
84
85
            kwargs_list.append(
                {
                    "num_meshes": n,
                    "ico_level": ic,
                    "image_size": im,
                    "blur_radius": b,
Nikhila Ravi's avatar
Nikhila Ravi committed
86
                    "faces_per_pixel": f,
facebook-github-bot's avatar
facebook-github-bot committed
87
88
89
90
91
92
93
94
                }
            )
        benchmark(
            TestRasterizeMeshes.rasterize_meshes_cuda_with_init,
            "RASTERIZE_MESHES_CUDA",
            kwargs_list,
            warmup_iters=1,
        )
Christoph Lassner's avatar
Christoph Lassner committed
95

96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
        # Test a subset of the cases with the
        # image plane intersecting the mesh.
        kwargs_list = []
        num_meshes = [8, 16]
        # Square and non square cases
        image_size = [64, 128, 512, (512, 256), (256, 512)]
        dist = [3, 0.8, 0.5]
        test_cases = product(num_meshes, dist, image_size)

        for case in test_cases:
            n, d, im = case
            kwargs_list.append(
                {
                    "num_meshes": n,
                    "ico_level": 4,
                    "image_size": im,
                    "blur_radius": 1e-6,
                    "faces_per_pixel": 40,
                    "dist": d,
                }
            )

        benchmark(
            TestRasterizeMeshes.bm_rasterize_meshes_with_clipping,
            "RASTERIZE_MESHES_CUDA_CLIPPING",
            kwargs_list,
            warmup_iters=1,
        )

Christoph Lassner's avatar
Christoph Lassner committed
125
126
127

if __name__ == "__main__":
    bm_rasterize_meshes()