Skip to content

Commit

Permalink
Fix ConcretizeStructAnnotation to deal with deps of default parametri…
Browse files Browse the repository at this point in the history
…c exprs.

This partially fixes #1523, and the IR converter test added here is a derivative of the code in that issue, which is completely fixed by this change.

PiperOrigin-RevId: 671166194
  • Loading branch information
richmckeever authored and copybara-github committed Sep 5, 2024
1 parent 72048d6 commit afc720d
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 6 deletions.
22 changes: 22 additions & 0 deletions xls/dslx/ir_convert/ir_converter_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -877,6 +877,28 @@ fn f(outer_thing_1: u32, outer_thing_2: u32) -> u32 {
ExpectIr(converted, TestName());
}

TEST(IrConverterTest, ParametricDefaultInStruct) {
const char* kProgram = R"(
struct Foo <X: u32, Y: u32 = {X + u32:1}> {
a: uN[X],
b: uN[Y],
}
fn make_zero_foo<X: u32>() -> Foo<X> {
zero!<Foo<X>>()
}
fn test() -> Foo<u32:5> {
make_zero_foo<u32:5>()
}
)";

XLS_ASSERT_OK_AND_ASSIGN(
std::string converted,
ConvertModuleForTest(kProgram, ConvertOptions{.emit_positions = false}));
ExpectIr(converted, TestName());
}

TEST(IrConverterTest, UnrollForSimple) {
const char* kProgram = R"(
fn test() -> u32 {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
package test_module

file_number 0 "test_module.x"

fn __test_module__make_zero_foo__5() -> (bits[5], bits[6]) {
X: bits[32] = literal(value=5, id=1)
ret literal.2: (bits[5], bits[6]) = literal(value=(0, 0), id=2)
}

fn __test_module__test() -> (bits[5], bits[6]) {
ret invoke.3: (bits[5], bits[6]) = invoke(to_apply=__test_module__make_zero_foo__5, id=3)
}
35 changes: 29 additions & 6 deletions xls/dslx/type_system/deduce.cc
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,29 @@ absl::StatusOr<InterpValue> EvaluateConstexprValue(DeduceCtx* ctx,
ctx->GetCurrentParametricEnv(), node, type);
}

// Evaluates the given `dim` to the extent possible with the given `env`.
// Examples:
// - If `dim` is already a constant, the result is `dim`.
// - If `dim` is `X + 1`, and `env` says `X` == `5`, then the result is `6`.
// - If `dim` is `X + 1`, and `env` says `X` == `Y` and does not have `Y`,
// then the result is `Y + 1`.
// - If `dim` is `X + 1`, and `env` says `X` == `Y` and `Y` == `5`, then the
// result is `6`.
TypeDim EvaluateTypeDim(TypeDim dim, const ParametricExpression::Env& env) {
absl::flat_hash_set<std::string> prev_free_variables;
while (std::holds_alternative<TypeDim::OwnedParametric>(dim.value())) {
auto& parametric = std::get<TypeDim::OwnedParametric>(dim.value());
absl::flat_hash_set<std::string> next_free_variables =
parametric->GetFreeVariables();
if (prev_free_variables == next_free_variables) {
break;
}
prev_free_variables = std::move(next_free_variables);
dim = TypeDim(parametric->Evaluate(env));
}
return dim;
}

// Attempts to convert an expression from the full DSL AST into the
// ParametricExpression sub-AST (a limited form that we can embed into a
// TypeDim for later instantiation).
Expand Down Expand Up @@ -1683,6 +1706,10 @@ static absl::StatusOr<std::unique_ptr<Type>> ConcretizeStructAnnotation(
defined_parametric->identifier(),
struct_def->identifier()));
}
XLS_ASSIGN_OR_RETURN(std::unique_ptr<ParametricExpression> parametric_expr,
ExprToParametric(defined_parametric->expr(), ctx));
parametric_env.emplace(defined_parametric->identifier(),
TypeDim(std::move(parametric_expr)));
}

ParametricExpression::Env env;
Expand All @@ -1698,16 +1725,12 @@ static absl::StatusOr<std::unique_ptr<Type>> ConcretizeStructAnnotation(
XLS_ASSIGN_OR_RETURN(
std::unique_ptr<Type> mapped_type,
base_type.MapSize([&](const TypeDim& dim) -> absl::StatusOr<TypeDim> {
if (std::holds_alternative<TypeDim::OwnedParametric>(dim.value())) {
auto& parametric = std::get<TypeDim::OwnedParametric>(dim.value());
return TypeDim(parametric->Evaluate(env));
}
return dim;
return EvaluateTypeDim(dim, env);
}));

// Attach the nominal parametrics to the type, so that we will remember the
// fact that we have instantiated e.g. Foo<M:u32, N:u32> as Foo<5, 6>.
return mapped_type->AddNominalTypeDims(std::move(parametric_env));
return mapped_type->AddNominalTypeDims(parametric_env);
}

absl::StatusOr<std::unique_ptr<Type>> DeduceTypeRefTypeAnnotation(
Expand Down

0 comments on commit afc720d

Please sign in to comment.