Merge branch 'main' into add-aqt-docs

Author: jerryzh168

README.md

@@ -29,9 +29,23 @@ git clone https://github.com/pytorch/ao
 cd ao
 pip install -r requirements.txt
 pip install -r dev-requirements.txt
-pip install .
 ```
 
+There are two options;
+-If you plan to be developing the library run:
+```Shell
+python setup.py develop
+```
+
+If you want to install from source run
+```Shell
+python setup.py install
+```
+
+** Note:
+Since we are building pytorch c++/cuda extensions by default, running `pip install .` will
+not work.
+
 ### Quantization
 
 ```python

test/test_ops.py

@@ -30,21 +30,6 @@ def _create_tensors_with_iou(self, N, iou_thresh):
         scores = torch.rand(N)
         return boxes, scores
 
-    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
-    @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "skipping when torch verion is 2.3 or lower")
-    def test_nms(self):
-        iou = 0.2
-        boxes, scores = self._create_tensors_with_iou(1000, iou)
-        boxes = boxes.cuda()
-        scores = scores.cuda()
-
-        # smoke test
-        _ = torchao.ops.nms(boxes, scores, iou)
-
-        # comprehensive testing
-        test_utils = ["test_schema", "test_autograd_registration", "test_faketensor", "test_aot_dispatch_dynamic"]
-        opcheck(torch.ops.torchao.nms, (boxes, scores, iou), test_utils=test_utils)
-
     def _create_fp6_inputs(self, BS: int, OC: int, IC: int):
         # Randomly initialize each bytes. The highest value for randint() is set the the max value of uint32_t.
         fp6_weight = torch.randint(4294967295, (OC, IC // 16 * 3)).to(torch.int)

torchao/csrc/cuda/nms.cu

@@ -10,26 +10,6 @@ def decorator(func):
             return torch.library.impl_abstract(f"{name}")(func)
     return decorator
 
-def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
-    """
-    See https://pytorch.org/vision/main/generated/torchvision.ops.nms.html
-    """
-    return torch.ops.torchao.nms.default(boxes, scores, iou_threshold)
-
-
-# Defines the meta kernel / fake kernel / abstract impl
-@register_custom_op("torchao::nms")
-def _(dets, scores, iou_threshold):
-    torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
-    torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
-    torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}")
-    torch._check(
-        dets.size(0) == scores.size(0),
-        lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}",
-    )
-    ctx = torch._custom_ops.get_ctx()
-    num_to_keep = ctx.create_unbacked_symint()
-    return dets.new_empty(num_to_keep, dtype=torch.long)
 
 
 def prepack_fp6_weight(fp6_weight: Tensor) -> Tensor:
@@ -45,6 +25,7 @@ def prepack_fp6_weight(fp6_weight: Tensor) -> Tensor:
     return torch.ops.torchao.prepack_fp6_weight.default(fp6_weight)
 
 
+# Defines the meta kernel / fake kernel / abstract impl
 @register_custom_op("torchao::prepack_fp6_weight")
 def _(fp6_weight):
     torch._check(fp6_weight.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp6_weight.dim()}D")