Skip to content

Commit

Permalink
[DSLX:TS][NFC] Break out deduce_spawn translation unit.
Browse files Browse the repository at this point in the history
Similar to instantiations has high enough complexity that it makes some sense
to encapsulate.

PiperOrigin-RevId: 621505425
  • Loading branch information
cdleary authored and copybara-github committed Apr 3, 2024
1 parent 708bf48 commit af3a0e3
Show file tree
Hide file tree
Showing 4 changed files with 252 additions and 158 deletions.
2 changes: 2 additions & 0 deletions xls/dslx/type_system/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,8 @@ cc_library(
"deduce_expr.h",
"deduce_invocation.cc",
"deduce_invocation.h",
"deduce_spawn.cc",
"deduce_spawn.h",
"deduce_struct_def.cc",
"deduce_struct_def.h",
],
Expand Down
159 changes: 1 addition & 158 deletions xls/dslx/type_system/deduce.cc
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
#include "xls/dslx/type_system/deduce_enum_def.h"
#include "xls/dslx/type_system/deduce_expr.h"
#include "xls/dslx/type_system/deduce_invocation.h"
#include "xls/dslx/type_system/deduce_spawn.h"
#include "xls/dslx/type_system/deduce_struct_def.h"
#include "xls/dslx/type_system/deduce_utils.h"
#include "xls/dslx/type_system/parametric_constraint.h"
Expand Down Expand Up @@ -134,18 +135,6 @@ absl::StatusOr<InterpValue> InterpretExpr(
/*args=*/{});
}

// Resolves "ref" to an AST proc.
absl::StatusOr<Proc*> ResolveColonRefToProc(const ColonRef* ref,
DeduceCtx* ctx) {
std::optional<Import*> import = ref->ResolveImportSubject();
XLS_RET_CHECK(import.has_value())
<< "ColonRef did not refer to an import: " << ref->ToString();
std::optional<const ImportedInfo*> imported_info =
ctx->type_info()->GetImported(*import);
return GetMemberOrTypeInferenceError<Proc>(imported_info.value()->module,
ref->attr(), ref->span());
}

absl::StatusOr<std::unique_ptr<Type>> DeduceProcMember(const ProcMember* node,
DeduceCtx* ctx) {
XLS_ASSIGN_OR_RETURN(auto type, ctx->Deduce(node->type_annotation()));
Expand Down Expand Up @@ -2151,152 +2140,6 @@ static absl::Status TypecheckProcConstAsserts(const Proc& p, DeduceCtx* ctx) {
return absl::OkStatus();
}

absl::StatusOr<std::unique_ptr<Type>> DeduceSpawn(const Spawn* node,
DeduceCtx* ctx) {
const ParametricEnv caller_parametric_env =
ctx->fn_stack().back().parametric_env();
VLOG(5) << "Deducing type for spawn: `" << node->ToString()
<< "` caller symbolic bindings: " << caller_parametric_env;

auto resolve_proc = [](const Instantiation* node,
DeduceCtx* ctx) -> absl::StatusOr<Proc*> {
Expr* callee = node->callee();
Proc* proc;
if (auto* colon_ref = dynamic_cast<ColonRef*>(callee)) {
XLS_ASSIGN_OR_RETURN(proc, ResolveColonRefToProc(colon_ref, ctx));
} else {
auto* name_ref = dynamic_cast<NameRef*>(callee);
XLS_RET_CHECK(name_ref != nullptr);
const std::string& callee_name = name_ref->identifier();
XLS_ASSIGN_OR_RETURN(
proc, GetMemberOrTypeInferenceError<Proc>(ctx->module(), callee_name,
name_ref->span()));
}
return proc;
};

// Resolve the proc AST node that's being instantiated.
//
// Note that this can be from a different module than the spawn / spawner.
XLS_ASSIGN_OR_RETURN(Proc * proc, resolve_proc(node, ctx));

auto resolve_config = [proc](const Instantiation* node,
DeduceCtx* ctx) -> absl::StatusOr<Function*> {
return &proc->config();
};
auto resolve_next = [proc](const Instantiation* node,
DeduceCtx* ctx) -> absl::StatusOr<Function*> {
return &proc->next();
};

auto resolve_init = [proc](const Instantiation* node,
DeduceCtx* ctx) -> absl::StatusOr<Function*> {
return &proc->init();
};

XLS_RETURN_IF_ERROR(
DeduceInstantiation(ctx, down_cast<Invocation*>(node->next()->args()[0]),
/*args=*/{}, /*resolve_fn=*/resolve_init,
/*constexpr_env=*/{})
.status());

// Gather up the type of all the (actual) arguments.
std::vector<InstantiateArg> config_args;
XLS_RETURN_IF_ERROR(AppendArgsForInstantiation(
node, node->callee(), node->config()->args(), ctx, &config_args));

std::vector<InstantiateArg> next_args;
next_args.push_back(
InstantiateArg{std::make_unique<TokenType>(), node->span()});
XLS_RETURN_IF_ERROR(AppendArgsForInstantiation(
node, node->callee(), node->next()->args(), ctx, &next_args));

// For each [constexpr] arg, mark the associated Param as constexpr.
AstEnv constexpr_env;
size_t argc = node->config()->args().size();
size_t paramc = proc->config().params().size();
if (argc != paramc) {
return TypeInferenceErrorStatus(
node->span(), nullptr,
absl::StrFormat("spawn had wrong argument count; want: %d got: %d",
paramc, argc));
}
for (int i = 0; i < node->config()->args().size(); i++) {
XLS_ASSIGN_OR_RETURN(InterpValue value,
ConstexprEvaluator::EvaluateToValue(
ctx->import_data(), ctx->type_info(),
ctx->warnings(), ctx->GetCurrentParametricEnv(),
node->config()->args()[i], nullptr));
constexpr_env.Add(proc->config().params()[i], value);
}

// TODO(rspringer): 2022-05-26: We can't currently lazily evaluate `next` args
// in the BytecodeEmitter, since that'd lead to a circular dependency between
// it and the ConstexprEvaluator, so we have to do it eagerly here.
// Un-wind that, if possible.
XLS_RETURN_IF_ERROR(ConstexprEvaluator::Evaluate(
ctx->import_data(), ctx->type_info(), ctx->warnings(),
ctx->GetCurrentParametricEnv(),
down_cast<Invocation*>(node->next()->args()[0]),
/*type=*/nullptr));

XLS_RETURN_IF_ERROR(DeduceInstantiation(ctx, node->config(), config_args,
resolve_config, constexpr_env)
.status());

XLS_ASSIGN_OR_RETURN(TypeInfo * config_ti,
ctx->type_info()->GetInvocationTypeInfoOrError(
node->config(), caller_parametric_env));

{
std::unique_ptr<DeduceCtx> proc_level_ctx =
ctx->MakeCtx(config_ti, proc->owner());
proc_level_ctx->fn_stack().push_back(FnStackEntry::MakeTop(proc->owner()));
XLS_RETURN_IF_ERROR(TypecheckProcConstAsserts(*proc, proc_level_ctx.get()));
}

// Now we need to get the [constexpr] Proc member values so we can set them
// when typechecking the `next` function. Those values are the elements in the
// `config` function's terminating XlsTuple.
// 1. Get the last statement in the `config` function.
absl::Nullable<const XlsTuple*> config_tuple = proc->GetConfigTuple();

// 2. Extract the value of each element and associate with the corresponding
// Proc member (in decl. order).
constexpr_env.Clear();

if (config_tuple == nullptr) {
if (!proc->members().empty()) {
return TypeInferenceErrorStatus(
proc->config().span(), nullptr,
absl::StrFormat("Proc '%s' has %d members to configure, but no proc "
"configuration tuple was provided",
proc->identifier(), proc->members().size()));
}
} else {
// When a config tuple is present (and returned from the `config()`) we
// constexpr evaluate all of its elements and place those in the constexpr
// env.
XLS_RET_CHECK_EQ(config_tuple->members().size(), proc->members().size());
for (int i = 0; i < config_tuple->members().size(); i++) {
XLS_ASSIGN_OR_RETURN(InterpValue value,
ConstexprEvaluator::EvaluateToValue(
ctx->import_data(), config_ti, ctx->warnings(),
ctx->GetCurrentParametricEnv(),
config_tuple->members()[i], nullptr));
constexpr_env.Add(proc->members()[i], value);
}
}

// With all the proc members placed in the constexpr env, now we can deduce
// the instantiation of `next()`, which accesses these members.
XLS_RETURN_IF_ERROR(DeduceInstantiation(ctx, node->next(), next_args,
resolve_next, constexpr_env)
.status());

return Type::MakeUnit();
}

absl::StatusOr<std::unique_ptr<Type>> DeduceNameRef(const NameRef* node,
DeduceCtx* ctx) {
AstNode* name_def = ToAstNode(node->name_def());
Expand Down
Loading

0 comments on commit af3a0e3

Please sign in to comment.