From 63871fc9a284d6567bd9638b8ba2e4201b04499a Mon Sep 17 00:00:00 2001 From: Michele Scuttari Date: Mon, 7 Oct 2024 18:08:46 +0200 Subject: [PATCH] Ignore read-only variables when applying using the IDA solver --- lib/Dialect/BaseModelica/Transforms/IDA.cpp | 114 ++++++++++++++------ 1 file changed, 80 insertions(+), 34 deletions(-) diff --git a/lib/Dialect/BaseModelica/Transforms/IDA.cpp b/lib/Dialect/BaseModelica/Transforms/IDA.cpp index 780941ce3..007861476 100644 --- a/lib/Dialect/BaseModelica/Transforms/IDA.cpp +++ b/lib/Dialect/BaseModelica/Transforms/IDA.cpp @@ -1337,6 +1337,11 @@ mlir::LogicalResult IDAInstance::getIndependentVariablesForAD( symbolTableCollection->lookupSymbolIn( modelOp, access.getVariable()); + if (variableOp.isReadOnly()) { + // Treat read-only variables as if they were just numbers. + continue; + } + result.insert(variableOp); if (auto derivative = getDerivative(access.getVariable())) { @@ -1445,8 +1450,11 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( // Start the body of the function. rewriter.setInsertionPointToStart(functionOp.getBody()); + // Keep track of all the variables that have been declared for creating the body of the function to be derived. + llvm::DenseSet allLocalVariables; + // Replicate the original independent variables inside the function. - llvm::StringMap localVariableOps; + llvm::StringMap mappedVariableOps; size_t independentVariableIndex = 0; for (VariableOp variableOp : variableOps) { @@ -1460,7 +1468,8 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( auto clonedVariableOp = rewriter.create( variableOp.getLoc(), variableOp.getSymName(), variableType); - localVariableOps[variableOp.getSymName()] = clonedVariableOp; + allLocalVariables.insert(clonedVariableOp.getSymName()); + mappedVariableOps[variableOp.getSymName()] = clonedVariableOp; independentVariablesPos[variableOp] = independentVariableIndex++; } @@ -1472,7 +1481,7 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( continue; } - auto mappedVariable = localVariableOps[variableOp.getSymName()]; + auto mappedVariable = mappedVariableOps[variableOp.getSymName()]; llvm::dbgs() << variableOp.getSymName() << " -> " << mappedVariable.getSymName() << "\n"; @@ -1495,6 +1504,7 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( auto variableOp = rewriter.create( loc, variableName, variableType); + allLocalVariables.insert(variableOp.getSymName()); inductionVariablesOps.push_back(variableOp); } @@ -1515,6 +1525,8 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( VariabilityProperty::none, IOProperty::output)); + allLocalVariables.insert(outputVariableOp.getSymName()); + // Create the body of the function. auto algorithmOp = rewriter.create(loc); @@ -1577,7 +1589,7 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( } if (auto globalGetOp = mlir::dyn_cast(op)) { - VariableOp variableOp = localVariableOps[globalGetOp.getVariable()]; + VariableOp variableOp = mappedVariableOps[globalGetOp.getVariable()]; auto getOp = rewriter.create( globalGetOp.getLoc(), variableOp); @@ -1598,6 +1610,19 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( rewriter.create(loc, outputVariableOp, result); + // Use the qualified accesses for the non-independent variables. + llvm::SmallVector getOpsToQualify; + + functionOp.walk([&](VariableGetOp getOp) { + if (!allLocalVariables.contains(getOp.getVariable())) { + getOpsToQualify.push_back(getOp); + } + }); + + if (mlir::failed(replaceVariableGetOps(rewriter, modelOp, getOpsToQualify))) { + return nullptr; + } + // Create the derivative template function. LLVM_DEBUG({ llvm::dbgs() << "Function being derived:\n" << functionOp << "\n"; @@ -1614,11 +1639,11 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( rewriter.eraseOp(functionOp); - // Replace the local variables with the global ones. + // Replace the mapped variables with qualified accesses. llvm::DenseSet variablesToBeReplaced; for (VariableOp variableOp : derTemplate->getVariables()) { - if (localVariableOps.count(variableOp.getSymName()) != 0) { + if (mappedVariableOps.contains(variableOp.getSymName())) { variablesToBeReplaced.insert(variableOp); } } @@ -1626,7 +1651,7 @@ FunctionOp IDAInstance::createPartialDerTemplateFromEquation( llvm::SmallVector variableGetOps; derTemplate->walk([&](VariableGetOp getOp) { - if (localVariableOps.count(getOp.getVariable()) != 0) { + if (mappedVariableOps.contains(getOp.getVariable())) { variableGetOps.push_back(getOp); } }); @@ -1747,14 +1772,16 @@ mlir::LogicalResult IDAInstance::createJacobianFunction( if (jacobianOneSweep) { // Perform just one call to the template function. - assert(independentVariablesPos.count(independentVariable) != 0); + std::optional oneSeedPosition = std::nullopt; - // Set the seed of the variable to one. - size_t oneSeedPosition = - independentVariablesPos.lookup(independentVariable); + if (independentVariablesPos.contains(independentVariable)) { + // Set the seed of the variable to one. + oneSeedPosition = + independentVariablesPos.lookup(independentVariable); - setGlobalADSeed(builder, loc, varSeeds[oneSeedPosition], - jacobianFunction.getVariableIndices(), one); + setGlobalADSeed(builder, loc, varSeeds[*oneSeedPosition], + jacobianFunction.getVariableIndices(), one); + } // Set the seed of the derivative to alpha. std::optional alphaSeedPosition = std::nullopt; @@ -1764,8 +1791,9 @@ mlir::LogicalResult IDAInstance::createJacobianFunction( auto derVariableOp = symbolTableCollection->lookupSymbolIn(modelOp, *derivative); - assert(independentVariablesPos.count(derVariableOp) != 0); - alphaSeedPosition = independentVariablesPos.lookup(derVariableOp); + if (independentVariablesPos.contains(derVariableOp)) { + alphaSeedPosition = independentVariablesPos.lookup(derVariableOp); + } } if (alphaSeedPosition) { @@ -1791,8 +1819,10 @@ mlir::LogicalResult IDAInstance::createJacobianFunction( mlir::Value result = templateCall.getResult(0); // Reset the seeds. - setGlobalADSeed(builder, loc, varSeeds[oneSeedPosition], - jacobianFunction.getVariableIndices(), zero); + if (oneSeedPosition) { + setGlobalADSeed(builder, loc, varSeeds[*oneSeedPosition], + jacobianFunction.getVariableIndices(), zero); + } if (alphaSeedPosition) { setGlobalADSeed(builder, loc, varSeeds[*alphaSeedPosition], @@ -1804,16 +1834,19 @@ mlir::LogicalResult IDAInstance::createJacobianFunction( builder.create(loc, result); } else { llvm::SmallVector args; + std::optional oneSeedPosition = std::nullopt; // Perform the first call to the template function. - assert(independentVariablesPos.count(independentVariable) != 0); - - // Set the seed of the variable to one. - size_t oneSeedPosition = - independentVariablesPos.lookup(independentVariable); + if (independentVariablesPos.contains(independentVariable)) { + // Set the seed of the variable to one. + oneSeedPosition = + independentVariablesPos.lookup(independentVariable); + } - setGlobalADSeed(builder, loc, varSeeds[oneSeedPosition], - jacobianFunction.getVariableIndices(), one); + if (oneSeedPosition) { + setGlobalADSeed(builder, loc, varSeeds[*oneSeedPosition], + jacobianFunction.getVariableIndices(), one); + } // Call the template function. args.clear(); @@ -1827,23 +1860,28 @@ mlir::LogicalResult IDAInstance::createJacobianFunction( mlir::Value result = firstTemplateCall.getResult(0); - // Reset the seed of the variable. - setGlobalADSeed(builder, loc, varSeeds[oneSeedPosition], - jacobianFunction.getVariableIndices(), zero); + if (oneSeedPosition) { + // Reset the seed of the variable. + setGlobalADSeed(builder, loc, varSeeds[*oneSeedPosition], + jacobianFunction.getVariableIndices(), zero); + } if (auto derivative = getDerivative(mlir::SymbolRefAttr::get( independentVariable.getSymNameAttr()))) { auto globalDerivativeOp = symbolTableCollection->lookupSymbolIn(modelOp, *derivative); - assert(independentVariablesPos.count(globalDerivativeOp) != 0); + std::optional derSeedPosition = std::nullopt; - size_t derSeedPosition = - independentVariablesPos.lookup(globalDerivativeOp); + if (independentVariablesPos.contains(globalDerivativeOp)) { + derSeedPosition = independentVariablesPos.lookup(globalDerivativeOp); + } - // Set the seed of the derivative to one. - setGlobalADSeed(builder, loc, varSeeds[derSeedPosition], - jacobianFunction.getVariableIndices(), one); + if (derSeedPosition) { + // Set the seed of the derivative to one. + setGlobalADSeed(builder, loc, varSeeds[*derSeedPosition], + jacobianFunction.getVariableIndices(), one); + } // Call the template function. args.clear(); @@ -1855,6 +1893,12 @@ mlir::LogicalResult IDAInstance::createJacobianFunction( RealType::get(builder.getContext()), args); + if (derSeedPosition) { + // Reset the seed of the variable. + setGlobalADSeed(builder, loc, varSeeds[*derSeedPosition], + jacobianFunction.getVariableIndices(), zero); + } + mlir::Value secondResult = secondTemplateCall.getResult(0); mlir::Value secondResultTimesAlpha = builder.create( @@ -2267,7 +2311,7 @@ mlir::LogicalResult IDAPass::solveMainModel( } else { LLVM_DEBUG(llvm::dbgs() << "Reduced system feature disabled\n"); - // Add all the variables to IDA. + // Add all the non-read-only variables to IDA. for (VariableOp variable : variables) { auto variableName = mlir::SymbolRefAttr::get(variable.getSymNameAttr()); @@ -2394,6 +2438,8 @@ mlir::LogicalResult IDAPass::addMainModelEquation( symbolTableCollection.lookupSymbolIn( modelOp, writtenVariable); + assert(!writtenVariableOp.isReadOnly()); + if (derivativesMap.getDerivedVariable(writtenVariable)) { LLVM_DEBUG(llvm::dbgs() << "Add derivative variable: " << writtenVariableOp.getSymName() << "\n");