Skip to content

Commit f7ee061

Browse files
ezyangpytorchmergebot
ezyang
authored and
pytorchmergebot
committed
Wconstab/reland pysymint (pytorch#79795)
rebased pytorch#79617 to see if issues are reproducible. Pull Request resolved: pytorch#79795 Approved by: https://github.com/malfet
1 parent a6b783e commit f7ee061

31 files changed

+879
-59
lines changed

‎aten/src/ATen/NestedTensorImpl.cpp

+7
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,13 @@ bool NestedTensorImpl::is_contiguous_custom(MemoryFormat) const {
7676
IntArrayRef NestedTensorImpl::sizes_custom() const {
7777
TORCH_CHECK(false, "Internal error: NestedTensorImpl doesn't support sizes. Please file an issue on https://github.com/pytorch/nestedtensor");
7878
}
79+
c10::SymIntArrayRef NestedTensorImpl::sym_sizes_custom() const {
80+
TORCH_CHECK(false, "Internal error: NestedTensorImpl doesn't support sizes. Please file an issue on https://github.com/pytorch/nestedtensor");
81+
}
82+
83+
c10::SymIntArrayRef NestedTensorImpl::sym_sizes() const {
84+
return sym_sizes_custom();
85+
}
7986

8087
IntArrayRef NestedTensorImpl::strides_custom() const {
8188
TORCH_CHECK(false, "Internal error: NestedTensorImpl doesn't support strides. Please file an issue on https://github.com/pytorch/nestedtensor");

‎aten/src/ATen/NestedTensorImpl.h

+2
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ struct TORCH_API NestedTensorImpl : public c10::TensorImpl {
4242
int64_t numel_custom() const override;
4343
bool is_contiguous_custom(MemoryFormat) const override;
4444
IntArrayRef sizes_custom() const override;
45+
c10::SymIntArrayRef sym_sizes_custom() const override;
46+
c10::SymIntArrayRef sym_sizes() const override;
4547
IntArrayRef strides_custom() const override;
4648

4749
// this one is real

‎aten/src/ATen/core/NamedRegistrations.cpp

+1
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,7 @@ TORCH_LIBRARY_IMPL(aten, Named, m) {
179179
m.impl("exp.out", CppFunction::makeFallthrough());
180180
m.impl("exp_", CppFunction::makeFallthrough());
181181
m.impl("expand", CppFunction::makeFallthrough());
182+
m.impl("expand.SymInt", CppFunction::makeFallthrough());
182183
m.impl("expm1", CppFunction::makeFallthrough());
183184
m.impl("expm1.out", CppFunction::makeFallthrough());
184185
m.impl("expm1_", CppFunction::makeFallthrough());

‎aten/src/ATen/core/TensorBase.h

+8
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,14 @@ class TORCH_API TensorBase {
156156
return at::isSignedType(this->scalar_type());
157157
}
158158

159+
c10::SymInt sym_size(int64_t dim) const {
160+
const auto sizes = this->sym_sizes();
161+
const auto ndim = static_cast<int64_t>(sizes.size());
162+
// false is passed to maybe_wrap_dim so behavior is identical to array access (but with wrapping)
163+
return sizes[c10::maybe_wrap_dim(dim, ndim, /*wrap_scalar=*/false)];
164+
165+
}
166+
159167
int64_t size(int64_t dim) const {
160168
const auto sizes = this->sizes();
161169
const auto ndim = static_cast<int64_t>(sizes.size());

‎aten/src/ATen/templates/TensorBody.h

+1
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ class TORCH_API Tensor: public TensorBase {
132132

133133
// Aliased by Dimname overloads, so need explicit using
134134
using TensorBase::size;
135+
using TensorBase::sym_size;
135136
using TensorBase::stride;
136137

137138
/// Should be used if *this can reasonably be expected to be contiguous and

‎c10/core/SymIntTable.cpp

+1
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,5 @@ SymIntTable& getSymIntTable() {
2525
static SymIntTable sit;
2626
return sit;
2727
}
28+
2829
} // namespace c10

‎c10/core/SymbolicIntNode.h

+47-1
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,53 @@ class C10_API SymbolicIntNode
1313
public:
1414
c10::SymInt toSymInt();
1515
virtual ~SymbolicIntNode(){};
16-
virtual std::ostream& operator<<(std::ostream& os) {
16+
// these could be pure virtual when we implement LTC versions
17+
virtual std::shared_ptr<SymbolicIntNode> add(
18+
const std::shared_ptr<SymbolicIntNode>& other) {
19+
TORCH_CHECK(false, "NYI");
20+
};
21+
virtual std::shared_ptr<SymbolicIntNode> sub(
22+
const std::shared_ptr<SymbolicIntNode>& other) {
23+
TORCH_CHECK(false, "NYI");
24+
};
25+
virtual std::shared_ptr<SymbolicIntNode> mul(
26+
const std::shared_ptr<SymbolicIntNode>& other) {
27+
TORCH_CHECK(false, "NYI");
28+
};
29+
virtual std::shared_ptr<SymbolicIntNode> div(
30+
const std::shared_ptr<SymbolicIntNode>& other) {
31+
TORCH_CHECK(false, "NYI");
32+
};
33+
virtual std::shared_ptr<SymbolicIntNode> mod(
34+
const std::shared_ptr<SymbolicIntNode>& other) {
35+
TORCH_CHECK(false, "NYI");
36+
};
37+
virtual std::shared_ptr<SymbolicIntNode> eq(
38+
const std::shared_ptr<SymbolicIntNode>& other) {
39+
TORCH_CHECK(false, "NYI");
40+
};
41+
virtual std::shared_ptr<SymbolicIntNode> gt(
42+
const std::shared_ptr<SymbolicIntNode>& other) {
43+
TORCH_CHECK(false, "NYI");
44+
};
45+
virtual std::shared_ptr<SymbolicIntNode> lt(
46+
const std::shared_ptr<SymbolicIntNode>& other) {
47+
TORCH_CHECK(false, "NYI");
48+
};
49+
virtual std::shared_ptr<SymbolicIntNode> wrap(int64_t num) {
50+
TORCH_CHECK(false, "NYI");
51+
};
52+
virtual bool bool_() {
53+
TORCH_CHECK(false, "NYI");
54+
};
55+
virtual int64_t int_() {
56+
TORCH_CHECK(false, "NYI");
57+
}
58+
virtual std::string str() {
59+
TORCH_CHECK(false, "NYI");
60+
};
61+
std::ostream& operator<<(std::ostream& os) {
62+
os << str();
1763
return os;
1864
};
1965
};

‎c10/core/TensorImpl.cpp

+9
Original file line numberDiff line numberDiff line change
@@ -811,6 +811,15 @@ void TensorImpl::ShareExternalPointer(
811811
}
812812
}
813813

814+
void TensorImpl::set_sym_sizes_and_strides(
815+
c10::SymIntArrayRef sizes,
816+
c10::SymIntArrayRef strides) {
817+
has_symbolic_sizes_strides_ = true;
818+
sizes_strides_policy_ = static_cast<uint8_t>(SizesStridesPolicy::CustomSizes);
819+
sizes_and_strides_.set_sizes(sizes);
820+
sizes_and_strides_.set_strides(strides);
821+
}
822+
814823
namespace impl {
815824

816825
namespace {

‎c10/core/TensorImpl.h

+9-7
Original file line numberDiff line numberDiff line change
@@ -552,12 +552,7 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
552552
return sizes_default();
553553
}
554554

555-
c10::SymIntArrayRef sym_sizes() const {
556-
if (C10_UNLIKELY(
557-
sizes_strides_policy_ >=
558-
static_cast<uint8_t>(SizesStridesPolicy::CustomSizes))) {
559-
return sym_sizes_custom();
560-
}
555+
virtual c10::SymIntArrayRef sym_sizes() const {
561556
return sym_sizes_default();
562557
}
563558

@@ -1312,6 +1307,12 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
13121307
return numel() == 0;
13131308
}
13141309

1310+
// if we are going to use sym sizes, we should be setting sym strides at the
1311+
// same time, otherwise it's very easy to misuse this API
1312+
void set_sym_sizes_and_strides(
1313+
c10::SymIntArrayRef sizes,
1314+
c10::SymIntArrayRef strides);
1315+
13151316
/**
13161317
* Change the size at some dimension. This DOES NOT update strides;
13171318
* thus, most changes to size will not preserve contiguity. You probably
@@ -2326,7 +2327,7 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
23262327
// Customizable sizes behavior, e.g., nested tensor
23272328
//
23282329
// Can override: strides(), is_contiguous(), sizes(), dim(), numel()
2329-
CustomSizes = 2,
2330+
CustomSizes = 2
23302331
};
23312332

23322333
void set_sizes_strides_policy(SizesStridesPolicy policy) {
@@ -2337,6 +2338,7 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
23372338
custom_device_ = custom_device;
23382339
}
23392340

2341+
protected:
23402342
Storage storage_;
23412343

23422344
private:

‎c10/core/impl/SizesAndStrides.h

+5
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,11 @@ class C10_API SizesAndStrides {
170170
std::copy(newSizes.begin(), newSizes.end(), sizes_begin());
171171
}
172172

173+
void set_strides(SymIntArrayRef strides) {
174+
TORCH_INTERNAL_ASSERT(strides.size() == size());
175+
std::copy(strides.begin(), strides.end(), strides_begin());
176+
}
177+
173178
void set_sizes(IntArrayRef newSizes) {
174179
set_sizes(SymIntArrayRef::fromIntArrayRef(newSizes));
175180
}

‎docs/source/conf.py

+1
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,7 @@
320320
"Quantize",
321321
# torch.utils.backcompat
322322
"Warning",
323+
"SymbolicIntNode"
323324
]
324325

325326
# The suffix(es) of source filenames.

‎requirements.txt

+1
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,4 @@ setuptools
1010
six
1111
types-dataclasses
1212
typing_extensions
13+
sympy

‎test/lazy/test_reuse_ir.py

+9-4
Original file line numberDiff line numberDiff line change
@@ -104,15 +104,20 @@ def testAddSubFallback(self):
104104
def testBatchNorm(self):
105105
device = get_test_device()
106106
x = torch.randn(16, 3, 224, 224, device=device)
107-
bn = torch.nn.BatchNorm2d(3).to(device=device)
107+
weight = torch.randn(3, device=device)
108+
bias = torch.randn(3, device=device)
109+
108110
for i in range(10):
109-
z = bn(x)
111+
# BatchNorm2d does extra checks on dimensions which SymInts don't support yet
112+
# so we call `torch.ops.aten.native_batch_norm` to bypass the checks.
113+
z, _, _ = torch.ops.aten.native_batch_norm(x, weight, bias, None, None, True, 0.1, 1e-5)
110114

111115
device = "lazy"
112116
x_lazy = x.detach().clone().to(device=device)
113-
bn = bn.to(device=device)
117+
weight_lazy = weight.detach().clone().to(device=device)
118+
bias_lazy = bias.detach().clone().to(device=device)
114119
for i in range(10):
115-
z_lazy = bn(x_lazy)
120+
z_lazy, _, _ = torch.ops.aten.native_batch_norm(x_lazy, weight_lazy, bias_lazy, None, None, True, 0.1, 1e-5)
116121
torch._lazy.mark_step()
117122

118123
torch.testing.assert_close(z.cpu(), z_lazy.cpu())

‎test/lazy/test_ts_opinfo.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import yaml
1818
import os
1919
import pathlib
20+
from unittest import skip
2021

2122
torch._lazy.ts_backend.init()
2223

@@ -66,6 +67,9 @@ def clone_move(t):
6667
return copy_t
6768

6869
class TestLazyTensor(JitTestCase):
70+
71+
72+
@skip("Disable until autograd supports symints")
6973
def testConvolutionBackward(self):
7074
test_device = get_test_device()
7175
inp = torch.rand(1, 3, 128, 128, device=test_device, requires_grad=True)
@@ -220,8 +224,9 @@ def test_nonzero_dynamic(self):
220224
x1 = torch.tensor([[0, 1.0, 2.0], [3.0, 0, 0]], device=test_device, requires_grad=True)
221225
x1_lazy = clone_move(x1)
222226
x2_lazy = torch.nonzero(x1_lazy)
223-
print(x2_lazy.size())
224-
self.assertEqual(tuple(x2_lazy.size()), (6, 2))
227+
228+
# FIXME: Add bindings to get upper bounds
229+
# self.assertEqual(tuple(x2_lazy.size()), (6, 2))
225230

226231
# We should still be able to instantiate it and get the actual result
227232
x2_eager = x2_lazy.cpu()

0 commit comments

Comments
 (0)