yueyericardo

from torchani.datasets import AniH5Dataset
import os
import cudf

try:
    path = os.path.dirname(os.path.realpath(__file__))
except NameError:
    path = os.getcwd()
h5file = os.path.join(path, '../dataset/ani-1x/sample.h5')

yueyericardo

from openmm.app import *
from openmm import *
from openmm.unit import *
from simtk import openmm, unit
from openmmtools import testsystems

water_box = testsystems.WaterBox(box_edge=20.0*unit.nanometer)
positions = water_box.positions
atoms = list(water_box.topology.atoms())
elements = [a.element.symbol for a in atoms]

yueyericardo

#!/usr/bin/env python

total = 10

def info(text):
    print('\033[32m{}\33[0m'.format(text))  # green

def cal(points):
    eq = points.replace(' ', ' + ')
    eq = f'{total} - ({eq})'

yueyericardo

Compiler	Toolset Versions Currently Tested
XCode	13.2.1 & 12.5.1
GCC	111 & 10[^2]
Clang	12[^2] & 11[^2]
Visual Studio with MSVC	VS2022[^3] & VS2019[^4]
Visual Studio with LLVM	VS2022[^3] & VS2019[^4]

Footnotes

1. Precise version may be found in the latest CI results. ↩

yueyericardo

import time
import argparse
from functorch import vmap, jacrev, jacfwd
import torch
import torch.nn as nn

torch.backends.cuda.matmul.allow_tf32 = False


_ = torch.manual_seed(0)

yueyericardo

====================================================================== erf ======================================================================
/opt/conda/lib/python3.8/site-packages/torch/profiler/profiler.py:395: UserWarning: use_cuda is deprecated, use activities argument instead
  warn("use_cuda is deprecated, use activities argument instead")
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
void at::native::vectorized_elementwise_kernel<4, at...         0.00%       0.000us

yueyericardo

300k atoms 8 GPUs

gpu/ware off

Loop time of 30.3401 on 8 procs for 2000 steps with 300003 atoms

Performance: 2.848 ns/day, 8.428 hours/ns, 65.919 timesteps/s
77.5% CPU use with 8 MPI tasks x 1 OpenMP threads

yueyericardo

CUAEV benchmark: https://github.com/roitberg-group/torchani_sandbox/blob/ed90fa65a7f07e59a95e75c962371a37ffb69ba0/tools/aev-benchmark-size.py

intrinsics on: python setup.py develop --ext --cuaev-opt use_fast_math need nvcc args: -use_fast_math

A100 Result

Turn on use_fast_math, Turn on intrinsics

File: small.pdb, Molecule size: 264 / 264, Species: [1, 6, 7, 8]