Skip to content

Commit

Permalink
Update package manager to poetry.
Browse files Browse the repository at this point in the history
  • Loading branch information
tatsy committed Jul 13, 2024
1 parent 3aef8af commit a837f2c
Show file tree
Hide file tree
Showing 12 changed files with 2,322 additions and 1,192 deletions.
28 changes: 28 additions & 0 deletions .clang-format
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
BasedOnStyle: Google

ColumnLimit: 120
IndentWidth: 4
Standard: Cpp11
AccessModifierOffset: -4
IndentCaseLabels: false

AlignAfterOpenBracket: Align
AlignConsecutiveAssignments: false
AlignTrailingComments: true
AlignOperands: true
AllowAllParametersOfDeclarationOnNextLine: false
AllowAllConstructorInitializersOnNextLine: false
AllowShortBlocksOnASingleLine: false
AllowShortFunctionsOnASingleLine: false
BinPackArguments: true
BreakConstructorInitializers: BeforeColon
BreakConstructorInitializersBeforeComma: true
Cpp11BracedListStyle: false
ConstructorInitializerAllOnOneLineOrOnePerLine: false
DerivePointerAlignment: false
MaxEmptyLinesToKeep: 1
PointerAlignment: Right
ReflowComments: false
SortIncludes: false

UseTab: Never
104 changes: 104 additions & 0 deletions build_cxx.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
import platform

from setuptools import setup
from setuptools.errors import CCompilerError, PackageDiscoveryError
from torch.utils.cpp_extension import BuildExtension


class MyBuildExtension(BuildExtension):
def run(self):
try:
super(MyBuildExtension, self).run()
except FileNotFoundError:
raise Exception("File not found. Could not compile C extension")

def build_extension(self, ext):
# common settings
for e in self.extensions:
pass

# OS specific settings
if platform.system() == "Darwin":
for e in self.extensions:
e.extra_compile_args.extend([
"-Xpreprocessor",
"-fopenmp",
"-mmacosx-version-min=10.15",
])
e.extra_link_args.extend([
"-lomp",
])

elif platform.system() == "Linux":
for e in self.extensions:
e.extra_compile_args.extend([
"-fopenmp",
])
e.extra_link_args.extend([
"-fopenmp",
])

# compiler specific settings
if self.compiler.compiler_type == "unix":
for e in self.extensions:
e.extra_compile_args.extend([
"-std=c++17",
"-pthread",
])

elif self.compiler.compiler_type == "msvc":
for e in self.extensions:
e.extra_compile_args.extend(["/utf-8", "/std:c++17", "/openmp"])
e.define_macros.extend([
("_CRT_SECURE_NO_WARNINGS", 1),
("_SILENCE_EXPERIMENTAL_FILESYSTEM_DEPRECATION_WARNING", 1),
])

# building
try:
super(MyBuildExtension, self).build_extension(ext)
except (CCompilerError, PackageDiscoveryError, ValueError):
raise Exception("Could not compile C extension")


def build(setup_kwargs):
from torch.utils.cpp_extension import CUDAExtension

try:
setup_kwargs.update({
"ext_modules": [
CUDAExtension(
'torchmcubes_module',
[
'cxx/pscan.cu',
'cxx/mcubes.cpp',
'cxx/mcubes_cpu.cpp',
'cxx/mcubes_cuda.cu',
'cxx/grid_interp_cpu.cpp',
'cxx/grid_interp_cuda.cu',
],
extra_compile_args=['-DWITH_CUDA'],
)
],
"cmdclass": {
"build_ext": BuildExtension
}
})
except:
from torch.utils.cpp_extension import CppExtension

print('CUDA environment was not successfully loaded!')
print('Build only CPU module!')

setup_kwargs.update({
'ext_modules': [
CppExtension('torchmcubes_module', [
'cxx/mcubes.cpp',
'cxx/mcubes_cpu.cpp',
'cxx/grid_interp_cpu.cpp',
])
],
'cmdclass': {
'build_ext': MyBuildExtension
}
})
13 changes: 6 additions & 7 deletions cxx/grid_interp_cpu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,16 +25,15 @@ torch::Tensor grid_interp_cpu(torch::Tensor vol, torch::Tensor points) {
const int C = vol.size(0);
const int Np = points.size(0);

torch::Tensor output = torch::zeros({Np, C},
torch::TensorOptions().dtype(torch::kFloat32).device(vol.device()));
torch::Tensor output = torch::zeros({ Np, C }, torch::TensorOptions().dtype(torch::kFloat32).device(vol.device()));

auto vol_ascr = vol.accessor<float, 4>();
auto pts_ascr = points.accessor<float, 2>();
auto out_ascr = output.accessor<float, 2>();

#ifdef _OPENMP
#pragma omp parallel for
#endif
#ifdef _OPENMP
#pragma omp parallel for
#endif
for (int i = 0; i < Np; i++) {
const float x = pts_ascr[i][0];
const float y = pts_ascr[i][1];
Expand All @@ -59,12 +58,12 @@ torch::Tensor grid_interp_cpu(torch::Tensor vol, torch::Tensor points) {
const float v01 = (1.0 - fx) * vol_ascr[c][z0][y1][x0] + fx * vol_ascr[c][z0][y1][x1];
const float v10 = (1.0 - fx) * vol_ascr[c][z1][y0][x0] + fx * vol_ascr[c][z1][y0][x1];
const float v11 = (1.0 - fx) * vol_ascr[c][z1][y1][x0] + fx * vol_ascr[c][z1][y1][x1];

const float v0 = (1.0 - fy) * v00 + fy * v01;
const float v1 = (1.0 - fy) * v10 + fy * v11;

out_ascr[i][c] = (1.0 - fz) * v0 + fz * v1;
}
}
}

return output;
Expand Down
Loading

0 comments on commit a837f2c

Please sign in to comment.