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 481 : VariableStore::VariableStore(Model * object)
37 481 : : _object(object),
38 962 : _input_axis(declare_axis("input")),
39 481 : _output_axis(declare_axis("output")),
40 962 : _options(default_tensor_options())
41 : {
42 481 : }
43 :
44 : LabeledAxis &
45 962 : VariableStore::declare_axis(const std::string & name)
46 : {
47 962 : 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 962 : auto axis = std::make_unique<LabeledAxis>();
53 962 : auto [it, success] = _axes.emplace(name, std::move(axis));
54 1924 : return *it->second;
55 962 : }
56 :
57 : void
58 481 : VariableStore::setup_layout()
59 : {
60 481 : input_axis().setup_layout();
61 481 : output_axis().setup_layout();
62 481 : }
63 :
64 : template <typename T>
65 : const Variable<T> &
66 462 : VariableStore::declare_input_variable(const char * name,
67 : TensorShapeRef list_shape,
68 : bool allow_duplicate)
69 : {
70 1386 : if (_object->input_options().contains(name))
71 1386 : return declare_input_variable<T>(
72 924 : _object->input_options().get<VariableName>(name), list_shape, allow_duplicate);
73 :
74 0 : return declare_input_variable<T>(VariableName(name), list_shape, allow_duplicate);
75 : }
76 :
77 : template <typename T>
78 : const Variable<T> &
79 806 : VariableStore::declare_input_variable(const VariableName & name,
80 : TensorShapeRef list_shape,
81 : bool allow_duplicate)
82 : {
83 806 : const auto list_sz = utils::storage_size(list_shape);
84 806 : const auto base_sz = T::const_base_storage;
85 806 : const auto sz = list_sz * base_sz;
86 :
87 806 : if (!allow_duplicate || (allow_duplicate && !_input_axis.has_variable(name)))
88 804 : _input_axis.add_variable(name, sz);
89 806 : return *create_variable<T>(_input_variables, name, list_shape, allow_duplicate);
90 : }
91 : #define INSTANTIATE_DECLARE_INPUT_VARIABLE(T) \
92 : template const Variable<T> & VariableStore::declare_input_variable<T>( \
93 : const char *, TensorShapeRef, bool); \
94 : template const Variable<T> & VariableStore::declare_input_variable<T>( \
95 : const VariableName &, TensorShapeRef, bool)
96 : FOR_ALL_PRIMITIVETENSOR(INSTANTIATE_DECLARE_INPUT_VARIABLE);
97 :
98 : template <typename T>
99 : Variable<T> &
100 289 : VariableStore::declare_output_variable(const char * name, TensorShapeRef list_shape)
101 : {
102 867 : if (_object->input_options().contains(name))
103 867 : return declare_output_variable<T>(_object->input_options().get<VariableName>(name), list_shape);
104 :
105 0 : return declare_output_variable<T>(VariableName(name), list_shape);
106 : }
107 :
108 : template <typename T>
109 : Variable<T> &
110 422 : VariableStore::declare_output_variable(const VariableName & name, TensorShapeRef list_shape)
111 : {
112 422 : const auto list_sz = utils::storage_size(list_shape);
113 422 : const auto base_sz = T::const_base_storage;
114 422 : const auto sz = list_sz * base_sz;
115 :
116 422 : _output_axis.add_variable(name, sz);
117 422 : return *create_variable<T>(_output_variables, name, list_shape);
118 : }
119 : #define INSTANTIATE_DECLARE_OUTPUT_VARIABLE(T) \
120 : template Variable<T> & VariableStore::declare_output_variable<T>(const char *, TensorShapeRef); \
121 : template Variable<T> & VariableStore::declare_output_variable<T>(const VariableName &, \
122 : TensorShapeRef)
123 : FOR_ALL_PRIMITIVETENSOR(INSTANTIATE_DECLARE_OUTPUT_VARIABLE);
124 :
125 : const VariableBase *
126 405 : VariableStore::clone_input_variable(const VariableBase & var, const VariableName & new_name)
127 : {
128 405 : neml_assert(&var.owner() != _object, "Trying to clone a variable from the same model.");
129 :
130 405 : const auto var_name = new_name.empty() ? var.name() : new_name;
131 405 : neml_assert(
132 405 : !_input_variables.count(var_name), "Input variable '", var_name.str(), "' already exists.");
133 405 : auto var_clone = var.clone(var_name, _object);
134 :
135 405 : _input_axis.add_variable(var_name, var_clone->assembly_storage());
136 405 : auto [it, success] = _input_variables.emplace(var_name, std::move(var_clone));
137 810 : return it->second.get();
138 405 : }
139 :
140 : VariableBase *
141 133 : VariableStore::clone_output_variable(const VariableBase & var, const VariableName & new_name)
142 : {
143 133 : neml_assert(&var.owner() != _object, "Trying to clone a variable from the same model.");
144 :
145 133 : const auto var_name = new_name.empty() ? var.name() : new_name;
146 133 : neml_assert(
147 133 : !_output_variables.count(var_name), "Output variable '", var_name, "' already exists.");
148 133 : auto var_clone = var.clone(var_name, _object);
149 :
150 133 : _output_axis.add_variable(var_name, var_clone->assembly_storage());
151 133 : auto [it, success] = _output_variables.emplace(var_name, std::move(var_clone));
152 266 : return it->second.get();
153 133 : }
154 :
155 : template <typename T>
156 : Variable<T> *
157 1228 : VariableStore::create_variable(VariableStorage & variables,
158 : const VariableName & name,
159 : TensorShapeRef list_shape,
160 : bool allow_duplicate)
161 : {
162 : // Make sure we don't duplicate variables
163 1228 : if (!allow_duplicate)
164 1150 : neml_assert(!variables.count(name),
165 : "Trying to create variable '",
166 : name,
167 : "', but a variable with the same name already exists.");
168 :
169 1228 : VariableBase * var_base_ptr = nullptr;
170 :
171 1228 : if (allow_duplicate && variables.count(name))
172 2 : var_base_ptr = variables[name].get();
173 : else
174 : {
175 : // Allocate
176 1226 : std::unique_ptr<VariableBase> var;
177 1226 : var = std::make_unique<Variable<T>>(name, _object, list_shape);
178 1226 : auto [it, success] = variables.emplace(name, std::move(var));
179 1226 : var_base_ptr = it->second.get();
180 1226 : }
181 :
182 : // Cast it to the concrete type
183 1228 : auto var_ptr = dynamic_cast<Variable<T> *>(var_base_ptr);
184 1228 : if (!var_ptr)
185 0 : throw NEMLException("Internal error: Failed to cast variable '" + name.str() +
186 : "' to its concrete type.");
187 :
188 1228 : return var_ptr;
189 : }
190 : #define INSTANTIATE_CREATE_VARIABLE(T) \
191 : template Variable<T> * VariableStore::create_variable<T>( \
192 : VariableStorage &, const VariableName &, TensorShapeRef, bool)
193 : FOR_ALL_PRIMITIVETENSOR(INSTANTIATE_CREATE_VARIABLE);
194 :
195 : VariableBase &
196 56138 : VariableStore::input_variable(const VariableName & name)
197 : {
198 56138 : auto it = _input_variables.find(name);
199 56138 : neml_assert(it != _input_variables.end(),
200 : "Input variable ",
201 : name,
202 : " does not exist in model ",
203 56138 : _object->name());
204 112276 : return *it->second;
205 : }
206 :
207 : const VariableBase &
208 37818 : VariableStore::input_variable(const VariableName & name) const
209 : {
210 : // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
211 37818 : return const_cast<VariableStore *>(this)->input_variable(name);
212 : }
213 :
214 : VariableBase &
215 14490 : VariableStore::output_variable(const VariableName & name)
216 : {
217 14490 : auto it = _output_variables.find(name);
218 14490 : neml_assert(it != _output_variables.end(),
219 : "Output variable ",
220 : name,
221 : " does not exist in model ",
222 14490 : _object->name());
223 28980 : return *it->second;
224 : }
225 :
226 : const VariableBase &
227 529 : VariableStore::output_variable(const VariableName & name) const
228 : {
229 : // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
230 529 : return const_cast<VariableStore *>(this)->output_variable(name);
231 : }
232 :
233 : void
234 91 : VariableStore::send_variables_to(const TensorOptions & options)
235 : {
236 91 : _options = options;
237 91 : }
238 :
239 : void
240 9219 : VariableStore::clear_input()
241 : {
242 35265 : for (auto && [name, var] : input_variables())
243 26046 : if (var->owning())
244 11570 : var->clear();
245 9219 : }
246 :
247 : void
248 9219 : VariableStore::clear_output()
249 : {
250 20283 : for (auto && [name, var] : output_variables())
251 11064 : if (var->owning())
252 8605 : var->clear();
253 9219 : }
254 :
255 : void
256 547 : VariableStore::clear_derivatives()
257 : {
258 1153 : for (auto && [name, var] : output_variables())
259 606 : var->clear_derivatives();
260 547 : }
261 :
262 : void
263 9221 : VariableStore::zero_input()
264 : {
265 35269 : for (auto && [name, var] : input_variables())
266 26048 : if (var->owning())
267 11572 : var->zero(_options);
268 9221 : }
269 :
270 : void
271 9221 : VariableStore::zero_output()
272 : {
273 20287 : for (auto && [name, var] : output_variables())
274 11066 : if (var->owning())
275 8607 : var->zero(_options);
276 9221 : }
277 :
278 : void
279 3792 : VariableStore::assign_input(const ValueMap & vals)
280 : {
281 14975 : for (const auto & [name, val] : vals)
282 11183 : if (input_axis().has_variable(name))
283 11183 : input_variable(name).set(val.clone());
284 3792 : }
285 :
286 : void
287 1574 : VariableStore::assign_input(ValueMap && vals)
288 : {
289 6162 : for (const auto & [name, val] : std::move(vals))
290 4588 : if (input_axis().has_variable(name))
291 4588 : input_variable(name).set(val.clone());
292 1574 : }
293 :
294 : void
295 26 : VariableStore::assign_output(const ValueMap & vals)
296 : {
297 104 : for (const auto & [name, val] : vals)
298 78 : output_variable(name).set(val.clone());
299 26 : }
300 :
301 : void
302 9 : VariableStore::assign_output_derivatives(const DerivMap & derivs)
303 : {
304 30 : for (const auto & [yvar, deriv] : derivs)
305 : {
306 21 : auto & y = output_variable(yvar);
307 60 : for (const auto & [xvar, val] : deriv)
308 39 : y.derivatives().insert_or_assign(xvar, val.clone());
309 : }
310 9 : }
311 :
312 : void
313 454 : VariableStore::assign_input_stack(jit::Stack & stack)
314 : {
315 454 : const auto & vars = input_axis().variable_names();
316 454 : neml_assert_dbg(stack.size() >= vars.size(),
317 : "Number of input variables in the stack (",
318 454 : stack.size(),
319 : ") is smaller than the number of input variables in the model (",
320 454 : vars.size(),
321 : ").");
322 :
323 : // Last n tensors in the stack are the input variables
324 1557 : for (std::size_t i = 0; i < vars.size(); i++)
325 1103 : input_variable(vars[i]).set(stack[stack.size() - vars.size() + i].toTensor(), /*force=*/true);
326 :
327 : // Drop the input variables from the stack
328 454 : jit::drop(stack, vars.size());
329 454 : }
330 :
331 : void
332 6512 : VariableStore::assign_output_stack(jit::Stack & stack, bool out, bool dout, bool d2out)
333 : {
334 6512 : neml_assert_dbg(out || dout || d2out,
335 : "At least one of the output/derivative flags must be true.");
336 :
337 6512 : neml_assert_dbg(!stack.empty(), "Empty output stack.");
338 6512 : const auto stacklist = stack.back().toTensorVector();
339 :
340 : // With our protocol, the last tensor in the list is the sparsity tensor
341 6512 : const auto sparsity_tensor = stacklist.back().contiguous();
342 6512 : neml_assert_dbg(at::sum(sparsity_tensor).item<Size>() == Size(stacklist.size()) - 1,
343 : "Sparsity tensor has incorrect size. Got ",
344 13024 : at::sum(sparsity_tensor).item<Size>(),
345 : " expected ",
346 13024 : Size(stacklist.size()) - 1);
347 : const std::vector<Size> sparsity(sparsity_tensor.data_ptr<Size>(),
348 6512 : sparsity_tensor.data_ptr<Size>() + sparsity_tensor.size(0));
349 :
350 6512 : const auto & yvars = output_axis().variable_names();
351 6512 : const auto & xvars = input_axis().variable_names();
352 :
353 6512 : std::size_t sti = 0; // stack counter
354 6512 : std::size_t spi = 0; // sparsity counter
355 :
356 6512 : if (out)
357 : {
358 11943 : for (const auto & yvar : yvars)
359 : {
360 7795 : neml_assert(sparsity[spi++], "Corrupted sparsity tensor.");
361 7795 : output_variable(yvar).set(stacklist[sti++], /*force=*/true);
362 : }
363 : }
364 :
365 6512 : if (dout)
366 : {
367 7191 : for (const auto & yvar : yvars)
368 : {
369 4887 : auto & derivs = output_variable(yvar).derivatives();
370 42683 : for (const auto & xvar : xvars)
371 : {
372 37796 : if (sparsity[spi++])
373 : {
374 14670 : const auto & val = stacklist[sti++];
375 14670 : neml_assert_dbg(val.dim() >= 2,
376 : "Derivative tensor d(",
377 : yvar,
378 : ")/d(",
379 : xvar,
380 : ") must have at least 2 dimensions. Got ",
381 14670 : val.dim(),
382 : ".");
383 14670 : derivs[xvar] = Tensor(val, val.dim() - 2);
384 : }
385 : }
386 : }
387 : }
388 :
389 6512 : if (d2out)
390 : {
391 134 : for (const auto & yvar : yvars)
392 : {
393 72 : auto & derivs = output_variable(yvar).second_derivatives();
394 191 : for (const auto & x1var : xvars)
395 432 : for (const auto & x2var : xvars)
396 : {
397 313 : if (sparsity[spi++])
398 : {
399 231 : const auto & val = stacklist[sti++];
400 231 : neml_assert_dbg(val.dim() >= 3,
401 : "Second derivative tensor d2(",
402 : yvar,
403 : ")/d(",
404 : x1var,
405 : ")d(",
406 : x2var,
407 : ") must have at least 3 dimensions. Got ",
408 231 : val.dim(),
409 : ".");
410 231 : derivs[x1var][x2var] = Tensor(val, val.dim() - 3);
411 : }
412 : }
413 : }
414 : }
415 :
416 6512 : jit::drop(stack, 1);
417 6512 : }
418 :
419 : ValueMap
420 191 : VariableStore::collect_input() const
421 : {
422 191 : ValueMap vals;
423 1846 : for (auto && [name, var] : input_variables())
424 1655 : vals[name] = var->tensor();
425 191 : return vals;
426 0 : }
427 :
428 : ValueMap
429 4250 : VariableStore::collect_output() const
430 : {
431 4250 : ValueMap vals;
432 12422 : for (auto && [name, var] : output_variables())
433 8172 : vals[name] = var->tensor();
434 4250 : return vals;
435 0 : }
436 :
437 : DerivMap
438 2306 : VariableStore::collect_output_derivatives() const
439 : {
440 2306 : DerivMap derivs;
441 7197 : for (auto && [name, var] : output_variables())
442 4891 : derivs[name] = var->derivatives();
443 2306 : return derivs;
444 0 : }
445 :
446 : SecDerivMap
447 62 : VariableStore::collect_output_second_derivatives() const
448 : {
449 62 : SecDerivMap sec_derivs;
450 134 : for (auto && [name, var] : output_variables())
451 72 : sec_derivs[name] = var->second_derivatives();
452 62 : return sec_derivs;
453 0 : }
454 :
455 : jit::Stack
456 6968 : VariableStore::collect_input_stack() const
457 : {
458 6968 : jit::Stack stack;
459 6968 : const auto & vars = input_axis().variable_names();
460 6968 : stack.reserve(vars.size());
461 43921 : for (const auto & name : vars)
462 36953 : stack.emplace_back(input_variable(name).tensor());
463 6968 : return stack;
464 0 : }
465 :
466 : jit::Stack
467 454 : VariableStore::collect_output_stack(bool out, bool dout, bool d2out) const
468 : {
469 454 : neml_assert_dbg(out || dout || d2out,
470 : "At least one of the output/derivative flags must be true.");
471 :
472 454 : const auto & yvars = output_axis().variable_names();
473 454 : const auto & xvars = input_axis().variable_names();
474 :
475 454 : std::vector<ATensor> stacklist;
476 454 : std::vector<Size> sparsity;
477 :
478 454 : if (out)
479 : {
480 200 : sparsity.insert(sparsity.end(), yvars.size(), 1);
481 430 : for (const auto & yvar : yvars)
482 230 : stacklist.push_back(output_variable(yvar).tensor());
483 : }
484 :
485 454 : if (dout)
486 : {
487 421 : for (const auto & yvar : yvars)
488 : {
489 227 : const auto & derivs = output_variable(yvar).derivatives();
490 910 : for (const auto & xvar : xvars)
491 : {
492 683 : const auto & deriv = derivs.find(xvar);
493 683 : sparsity.push_back(deriv == derivs.end() || !input_variable(xvar).is_dependent() ? 0 : 1);
494 683 : if (sparsity.back())
495 534 : stacklist.push_back(deriv->second);
496 : }
497 : }
498 : }
499 :
500 454 : if (d2out)
501 : {
502 134 : for (const auto & yvar : yvars)
503 : {
504 72 : const auto & derivs = output_variable(yvar).second_derivatives();
505 191 : for (const auto & x1var : xvars)
506 : {
507 119 : const auto & x1derivs = derivs.find(x1var);
508 119 : if (x1derivs != derivs.end() && input_variable(x1var).is_dependent())
509 389 : for (const auto & x2var : xvars)
510 : {
511 289 : const auto & x1x2deriv = x1derivs->second.find(x2var);
512 289 : sparsity.push_back(
513 289 : x1x2deriv == x1derivs->second.end() || !input_variable(x2var).is_dependent() ? 0
514 : : 1);
515 289 : if (sparsity.back())
516 231 : stacklist.push_back(x1x2deriv->second);
517 : }
518 : else
519 19 : sparsity.insert(sparsity.end(), xvars.size(), 0);
520 : }
521 : }
522 : }
523 :
524 454 : const auto sparsity_tensor = Tensor::create(sparsity, kInt64);
525 454 : const auto nnz = base_sum(sparsity_tensor).item<Size>();
526 454 : neml_assert_dbg(nnz == Size(stacklist.size()),
527 : "Corrupted sparsity tensor. Got ",
528 : nnz,
529 : " non-zero entries, expected ",
530 454 : Size(stacklist.size()));
531 454 : stacklist.push_back(sparsity_tensor);
532 :
533 1816 : return {stacklist};
534 908 : }
535 :
536 : } // namespace neml2
|
