Line data Source code
1 : // Copyright 2024, UChicago Argonne, LLC
2 : // All Rights Reserved
3 : // Software Name: NEML2 -- the New Engineering material Model Library, version 2
4 : // By: Argonne National Laboratory
5 : // OPEN SOURCE LICENSE (MIT)
6 : //
7 : // Permission is hereby granted, free of charge, to any person obtaining a copy
8 : // of this software and associated documentation files (the "Software"), to deal
9 : // in the Software without restriction, including without limitation the rights
10 : // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11 : // copies of the Software, and to permit persons to whom the Software is
12 : // furnished to do so, subject to the following conditions:
13 : //
14 : // The above copyright notice and this permission notice shall be included in
15 : // all copies or substantial portions of the Software.
16 : //
17 : // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18 : // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19 : // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20 : // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21 : // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22 : // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
23 : // THE SOFTWARE.
24 :
25 : #include "neml2/models/VariableStore.h"
26 : #include "neml2/models/Model.h"
27 : #include "neml2/misc/assertions.h"
28 : #include "neml2/models/map_types.h"
29 : #include "neml2/models/Variable.h"
30 : #include "neml2/base/LabeledAxis.h"
31 : #include "neml2/tensors/tensors.h"
32 : #include "neml2/tensors/functions/sum.h"
33 :
34 : namespace neml2
35 : {
36 432 : VariableStore::VariableStore(Model * object)
37 432 : : _object(object),
38 864 : _input_axis(declare_axis("input")),
39 432 : _output_axis(declare_axis("output")),
40 864 : _options(default_tensor_options())
41 : {
42 432 : }
43 :
44 : LabeledAxis &
45 864 : VariableStore::declare_axis(const std::string & name)
46 : {
47 864 : neml_assert(!_axes.count(name),
48 : "Trying to declare an axis named ",
49 : name,
50 : ", but an axis with the same name already exists.");
51 :
52 864 : auto axis = std::make_unique<LabeledAxis>();
53 864 : auto [it, success] = _axes.emplace(name, std::move(axis));
54 1728 : return *it->second;
55 864 : }
56 :
57 : void
58 432 : VariableStore::setup_layout()
59 : {
60 432 : input_axis().setup_layout();
61 432 : output_axis().setup_layout();
62 432 : }
63 :
64 : template <typename T>
65 : const Variable<T> &
66 417 : VariableStore::declare_input_variable(const char * name, TensorShapeRef list_shape)
67 : {
68 1251 : if (_object->input_options().contains(name))
69 1251 : return declare_input_variable<T>(_object->input_options().get<VariableName>(name), list_shape);
70 :
71 0 : return declare_input_variable<T>(VariableName(name), list_shape);
72 : }
73 :
74 : template <typename T>
75 : const Variable<T> &
76 730 : VariableStore::declare_input_variable(const VariableName & name, TensorShapeRef list_shape)
77 : {
78 730 : const auto list_sz = utils::storage_size(list_shape);
79 730 : const auto base_sz = T::const_base_storage;
80 730 : const auto sz = list_sz * base_sz;
81 :
82 730 : _input_axis.add_variable(name, sz);
83 730 : return *create_variable<T>(_input_variables, name, list_shape);
84 : }
85 : #define INSTANTIATE_DECLARE_INPUT_VARIABLE(T) \
86 : template const Variable<T> & VariableStore::declare_input_variable<T>(const char *, \
87 : TensorShapeRef); \
88 : template const Variable<T> & VariableStore::declare_input_variable<T>(const VariableName &, \
89 : TensorShapeRef)
90 : FOR_ALL_PRIMITIVETENSOR(INSTANTIATE_DECLARE_INPUT_VARIABLE);
91 :
92 : template <typename T>
93 : Variable<T> &
94 252 : VariableStore::declare_output_variable(const char * name, TensorShapeRef list_shape)
95 : {
96 756 : if (_object->input_options().contains(name))
97 756 : return declare_output_variable<T>(_object->input_options().get<VariableName>(name), list_shape);
98 :
99 0 : return declare_output_variable<T>(VariableName(name), list_shape);
100 : }
101 :
102 : template <typename T>
103 : Variable<T> &
104 372 : VariableStore::declare_output_variable(const VariableName & name, TensorShapeRef list_shape)
105 : {
106 372 : const auto list_sz = utils::storage_size(list_shape);
107 372 : const auto base_sz = T::const_base_storage;
108 372 : const auto sz = list_sz * base_sz;
109 :
110 372 : _output_axis.add_variable(name, sz);
111 372 : return *create_variable<T>(_output_variables, name, list_shape);
112 : }
113 : #define INSTANTIATE_DECLARE_OUTPUT_VARIABLE(T) \
114 : template Variable<T> & VariableStore::declare_output_variable<T>(const char *, TensorShapeRef); \
115 : template Variable<T> & VariableStore::declare_output_variable<T>(const VariableName &, \
116 : TensorShapeRef)
117 : FOR_ALL_PRIMITIVETENSOR(INSTANTIATE_DECLARE_OUTPUT_VARIABLE);
118 :
119 : const VariableBase *
120 359 : VariableStore::clone_input_variable(const VariableBase & var, const VariableName & new_name)
121 : {
122 359 : neml_assert(&var.owner() != _object, "Trying to clone a variable from the same model.");
123 :
124 359 : const auto var_name = new_name.empty() ? var.name() : new_name;
125 359 : neml_assert(
126 359 : !_input_variables.count(var_name), "Input variable '", var_name.str(), "' already exists.");
127 359 : auto var_clone = var.clone(var_name, _object);
128 :
129 359 : _input_axis.add_variable(var_name, var_clone->assembly_storage());
130 359 : auto [it, success] = _input_variables.emplace(var_name, std::move(var_clone));
131 718 : return it->second.get();
132 359 : }
133 :
134 : VariableBase *
135 118 : VariableStore::clone_output_variable(const VariableBase & var, const VariableName & new_name)
136 : {
137 118 : neml_assert(&var.owner() != _object, "Trying to clone a variable from the same model.");
138 :
139 118 : const auto var_name = new_name.empty() ? var.name() : new_name;
140 118 : neml_assert(
141 118 : !_output_variables.count(var_name), "Output variable '", var_name, "' already exists.");
142 118 : auto var_clone = var.clone(var_name, _object);
143 :
144 118 : _output_axis.add_variable(var_name, var_clone->assembly_storage());
145 118 : auto [it, success] = _output_variables.emplace(var_name, std::move(var_clone));
146 236 : return it->second.get();
147 118 : }
148 :
149 : template <typename T>
150 : Variable<T> *
151 1102 : VariableStore::create_variable(VariableStorage & variables,
152 : const VariableName & name,
153 : TensorShapeRef list_shape)
154 : {
155 : // Make sure we don't duplicate variables
156 1102 : neml_assert(!variables.count(name),
157 : "Trying to create variable '",
158 : name,
159 : "', but a variable with the same name already exists.");
160 :
161 : // Allocate
162 1102 : std::unique_ptr<VariableBase> var;
163 1102 : var = std::make_unique<Variable<T>>(name, _object, list_shape);
164 1102 : auto [it, success] = variables.emplace(name, std::move(var));
165 1102 : auto * var_base_ptr = it->second.get();
166 :
167 : // Cast it to the concrete type
168 1102 : auto var_ptr = dynamic_cast<Variable<T> *>(var_base_ptr);
169 1102 : if (!var_ptr)
170 0 : throw NEMLException("Internal error: Failed to cast variable '" + name.str() +
171 : "' to its concrete type.");
172 :
173 1102 : return var_ptr;
174 1102 : }
175 : #define INSTANTIATE_CREATE_VARIABLE(T) \
176 : template Variable<T> * VariableStore::create_variable<T>( \
177 : VariableStorage &, const VariableName &, TensorShapeRef)
178 : FOR_ALL_PRIMITIVETENSOR(INSTANTIATE_CREATE_VARIABLE);
179 :
180 : VariableBase &
181 53860 : VariableStore::input_variable(const VariableName & name)
182 : {
183 53860 : auto it = _input_variables.find(name);
184 53860 : neml_assert(it != _input_variables.end(),
185 : "Input variable ",
186 : name,
187 : " does not exist in model ",
188 53860 : _object->name());
189 107720 : return *it->second;
190 : }
191 :
192 : const VariableBase &
193 36647 : VariableStore::input_variable(const VariableName & name) const
194 : {
195 : // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
196 36647 : return const_cast<VariableStore *>(this)->input_variable(name);
197 : }
198 :
199 : VariableBase &
200 13360 : VariableStore::output_variable(const VariableName & name)
201 : {
202 13360 : auto it = _output_variables.find(name);
203 13360 : neml_assert(it != _output_variables.end(),
204 : "Output variable ",
205 : name,
206 : " does not exist in model ",
207 13360 : _object->name());
208 26720 : return *it->second;
209 : }
210 :
211 : const VariableBase &
212 435 : VariableStore::output_variable(const VariableName & name) const
213 : {
214 : // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
215 435 : return const_cast<VariableStore *>(this)->output_variable(name);
216 : }
217 :
218 : void
219 91 : VariableStore::send_variables_to(const TensorOptions & options)
220 : {
221 91 : _options = options;
222 91 : }
223 :
224 : void
225 8451 : VariableStore::clear_input()
226 : {
227 33119 : for (auto && [name, var] : input_variables())
228 24668 : if (var->owning())
229 10751 : var->clear();
230 8451 : }
231 :
232 : void
233 8451 : VariableStore::clear_output()
234 : {
235 18353 : for (auto && [name, var] : output_variables())
236 9902 : if (var->owning())
237 7564 : var->clear();
238 8451 : }
239 :
240 : void
241 8451 : VariableStore::zero_input()
242 : {
243 33119 : for (auto && [name, var] : input_variables())
244 24668 : if (var->owning())
245 10751 : var->zero(_options);
246 8451 : }
247 :
248 : void
249 8451 : VariableStore::zero_output()
250 : {
251 18353 : for (auto && [name, var] : output_variables())
252 9902 : if (var->owning())
253 7564 : var->zero(_options);
254 8451 : }
255 :
256 : void
257 3264 : VariableStore::assign_input(const ValueMap & vals)
258 : {
259 13631 : for (const auto & [name, val] : vals)
260 10367 : if (input_axis().has_variable(name))
261 10367 : input_variable(name).set(val.clone());
262 3264 : }
263 :
264 : void
265 1565 : VariableStore::assign_input(ValueMap && vals)
266 : {
267 6126 : for (const auto & [name, val] : std::move(vals))
268 4561 : if (input_axis().has_variable(name))
269 4561 : input_variable(name).set(val.clone());
270 1565 : }
271 :
272 : void
273 26 : VariableStore::assign_output(const ValueMap & vals)
274 : {
275 104 : for (const auto & [name, val] : vals)
276 78 : output_variable(name).set(val.clone());
277 26 : }
278 :
279 : void
280 9 : VariableStore::assign_output_derivatives(const DerivMap & derivs)
281 : {
282 30 : for (const auto & [yvar, deriv] : derivs)
283 : {
284 21 : auto & y = output_variable(yvar);
285 60 : for (const auto & [xvar, val] : deriv)
286 39 : y.derivatives().insert_or_assign(xvar, val.clone());
287 : }
288 9 : }
289 :
290 : void
291 394 : VariableStore::assign_input_stack(jit::Stack & stack)
292 : {
293 394 : const auto & vars = input_axis().variable_names();
294 394 : neml_assert_dbg(stack.size() >= vars.size(),
295 : "Number of input variables in the stack (",
296 394 : stack.size(),
297 : ") is smaller than the number of input variables in the model (",
298 394 : vars.size(),
299 : ").");
300 :
301 : // Last n tensors in the stack are the input variables
302 1383 : for (std::size_t i = 0; i < vars.size(); i++)
303 989 : input_variable(vars[i]).set(stack[stack.size() - vars.size() + i].toTensor(), /*force=*/true);
304 :
305 : // Drop the input variables from the stack
306 394 : jit::drop(stack, vars.size());
307 394 : }
308 :
309 : void
310 5972 : VariableStore::assign_output_stack(jit::Stack & stack, bool out, bool dout, bool d2out)
311 : {
312 5972 : neml_assert_dbg(out || dout || d2out,
313 : "At least one of the output/derivative flags must be true.");
314 :
315 5972 : neml_assert_dbg(!stack.empty(), "Empty output stack.");
316 5972 : const auto stacklist = stack.back().toTensorVector();
317 :
318 : // With our protocol, the last tensor in the list is the sparsity tensor
319 5972 : const auto sparsity_tensor = stacklist.back().contiguous();
320 5972 : neml_assert_dbg(at::sum(sparsity_tensor).item<Size>() == Size(stacklist.size()) - 1,
321 : "Sparsity tensor has incorrect size. Got ",
322 11944 : at::sum(sparsity_tensor).item<Size>(),
323 : " expected ",
324 11944 : Size(stacklist.size()) - 1);
325 : const std::vector<Size> sparsity(sparsity_tensor.data_ptr<Size>(),
326 5972 : sparsity_tensor.data_ptr<Size>() + sparsity_tensor.size(0));
327 :
328 5972 : const auto & yvars = output_axis().variable_names();
329 5972 : const auto & xvars = input_axis().variable_names();
330 :
331 5972 : std::size_t sti = 0; // stack counter
332 5972 : std::size_t spi = 0; // sparsity counter
333 :
334 5972 : if (out)
335 : {
336 11043 : for (const auto & yvar : yvars)
337 : {
338 7225 : neml_assert(sparsity[spi++], "Corrupted sparsity tensor.");
339 7225 : output_variable(yvar).set(stacklist[sti++], /*force=*/true);
340 : }
341 : }
342 :
343 5972 : if (dout)
344 : {
345 6635 : for (const auto & yvar : yvars)
346 : {
347 4527 : auto & derivs = output_variable(yvar).derivatives();
348 41746 : for (const auto & xvar : xvars)
349 : {
350 37219 : if (sparsity[spi++])
351 : {
352 14235 : const auto & val = stacklist[sti++];
353 14235 : neml_assert_dbg(val.dim() >= 2,
354 : "Derivative tensor d(",
355 : yvar,
356 : ")/d(",
357 : xvar,
358 : ") must have at least 2 dimensions. Got ",
359 14235 : val.dim(),
360 : ".");
361 14235 : derivs[xvar] = Tensor(val, val.dim() - 2);
362 : }
363 : }
364 : }
365 : }
366 :
367 5972 : if (d2out)
368 : {
369 96 : for (const auto & yvar : yvars)
370 : {
371 48 : auto & derivs = output_variable(yvar).second_derivatives();
372 139 : for (const auto & x1var : xvars)
373 368 : for (const auto & x2var : xvars)
374 : {
375 277 : if (sparsity[spi++])
376 : {
377 204 : const auto & val = stacklist[sti++];
378 204 : neml_assert_dbg(val.dim() >= 3,
379 : "Second derivative tensor d2(",
380 : yvar,
381 : ")/d(",
382 : x1var,
383 : ")d(",
384 : x2var,
385 : ") must have at least 3 dimensions. Got ",
386 204 : val.dim(),
387 : ".");
388 204 : derivs[x1var][x2var] = Tensor(val, val.dim() - 3);
389 : }
390 : }
391 : }
392 : }
393 :
394 5972 : jit::drop(stack, 1);
395 5972 : }
396 :
397 : ValueMap
398 189 : VariableStore::collect_input() const
399 : {
400 189 : ValueMap vals;
401 1826 : for (auto && [name, var] : input_variables())
402 1637 : vals[name] = var->tensor();
403 189 : return vals;
404 0 : }
405 :
406 : ValueMap
407 3918 : VariableStore::collect_output() const
408 : {
409 3918 : ValueMap vals;
410 11514 : for (auto && [name, var] : output_variables())
411 7596 : vals[name] = var->tensor();
412 3918 : return vals;
413 0 : }
414 :
415 : DerivMap
416 2110 : VariableStore::collect_output_derivatives() const
417 : {
418 2110 : DerivMap derivs;
419 6641 : for (auto && [name, var] : output_variables())
420 4531 : derivs[name] = var->derivatives();
421 2110 : return derivs;
422 0 : }
423 :
424 : SecDerivMap
425 48 : VariableStore::collect_output_second_derivatives() const
426 : {
427 48 : SecDerivMap sec_derivs;
428 96 : for (auto && [name, var] : output_variables())
429 48 : sec_derivs[name] = var->second_derivatives();
430 48 : return sec_derivs;
431 0 : }
432 :
433 : jit::Stack
434 6366 : VariableStore::collect_input_stack() const
435 : {
436 6366 : jit::Stack stack;
437 6366 : const auto & vars = input_axis().variable_names();
438 6366 : stack.reserve(vars.size());
439 42258 : for (const auto & name : vars)
440 35892 : stack.emplace_back(input_variable(name).tensor());
441 6366 : return stack;
442 0 : }
443 :
444 : jit::Stack
445 394 : VariableStore::collect_output_stack(bool out, bool dout, bool d2out) const
446 : {
447 394 : neml_assert_dbg(out || dout || d2out,
448 : "At least one of the output/derivative flags must be true.");
449 :
450 394 : const auto & yvars = output_axis().variable_names();
451 394 : const auto & xvars = input_axis().variable_names();
452 :
453 394 : std::vector<ATensor> stacklist;
454 394 : std::vector<Size> sparsity;
455 :
456 394 : if (out)
457 : {
458 177 : sparsity.insert(sparsity.end(), yvars.size(), 1);
459 372 : for (const auto & yvar : yvars)
460 195 : stacklist.push_back(output_variable(yvar).tensor());
461 : }
462 :
463 394 : if (dout)
464 : {
465 363 : for (const auto & yvar : yvars)
466 : {
467 192 : const auto & derivs = output_variable(yvar).derivatives();
468 799 : for (const auto & xvar : xvars)
469 : {
470 607 : const auto & deriv = derivs.find(xvar);
471 607 : sparsity.push_back(deriv == derivs.end() || !input_variable(xvar).is_dependent() ? 0 : 1);
472 607 : if (sparsity.back())
473 478 : stacklist.push_back(deriv->second);
474 : }
475 : }
476 : }
477 :
478 394 : if (d2out)
479 : {
480 96 : for (const auto & yvar : yvars)
481 : {
482 48 : const auto & derivs = output_variable(yvar).second_derivatives();
483 139 : for (const auto & x1var : xvars)
484 : {
485 91 : const auto & x1derivs = derivs.find(x1var);
486 91 : if (x1derivs != derivs.end() && input_variable(x1var).is_dependent())
487 327 : for (const auto & x2var : xvars)
488 : {
489 254 : const auto & x1x2deriv = x1derivs->second.find(x2var);
490 254 : sparsity.push_back(
491 254 : x1x2deriv == x1derivs->second.end() || !input_variable(x2var).is_dependent() ? 0
492 : : 1);
493 254 : if (sparsity.back())
494 204 : stacklist.push_back(x1x2deriv->second);
495 : }
496 : else
497 18 : sparsity.insert(sparsity.end(), xvars.size(), 0);
498 : }
499 : }
500 : }
501 :
502 394 : const auto sparsity_tensor = Tensor::create(sparsity, kInt64);
503 394 : const auto nnz = base_sum(sparsity_tensor).item<Size>();
504 394 : neml_assert_dbg(nnz == Size(stacklist.size()),
505 : "Corrupted sparsity tensor. Got ",
506 : nnz,
507 : " non-zero entries, expected ",
508 394 : Size(stacklist.size()));
509 394 : stacklist.push_back(sparsity_tensor);
510 :
511 1576 : return {stacklist};
512 788 : }
513 :
514 : } // namespace neml2
|
