Merge pull request #1096 from cdleary:array-size-convert

PiperOrigin-RevId: 555014997
google · Aug 9, 2023 · 622b2c7 · 622b2c7
2 parents 916686b + a0de9e1
commit 622b2c7
Show file tree

Hide file tree

Showing 7 changed files with 43 additions and 3 deletions.
diff --git a/dependency_support/z3/build_defs.bzl b/dependency_support/z3/build_defs.bzl
@@ -83,7 +83,7 @@ def gen_srcs():
         srcs = [
             "LICENSE.txt",
             "scripts/mk_make.py",
-        ] + native.glob(["src/**", "examples/**"]),
+        ] + native.glob(["src/**", "examples/**"], exclude = GEN_HDRS + GEN_SRCS),
         tools = ["scripts/mk_make.py"],
         outs = MK_MAKE_SRCS + MK_MAKE_HDRS,
         # We can't use $(location) here, since the bundled script internally

diff --git a/dependency_support/z3/bundled.BUILD.bazel b/dependency_support/z3/bundled.BUILD.bazel
@@ -15,7 +15,7 @@
 # Description:
 #   Z3 is a theorem prover from Microsoft Research.
 
-load("@com_google_xls//dependency_support/z3:build_defs.bzl", "gen_srcs", "GEN_SRCS", "GEN_HDRS")
+load("@com_google_xls//dependency_support/z3:build_defs.bzl", "GEN_HDRS", "GEN_SRCS", "gen_srcs")
 
 licenses(["notice"])
 
@@ -121,7 +121,7 @@ cc_library(
         "src/test/fuzzing/*.cpp",
         "src/util/*.cpp",
         "src/util/lp/*.cpp",
-    ]),
+    ], exclude = GEN_SRCS),
     hdrs = GEN_HDRS + glob([
         "src/ackermannization/*.h",
         "src/ackermannization/*.hpp",

diff --git a/xls/dslx/ir_convert/function_converter.cc b/xls/dslx/ir_convert/function_converter.cc
@@ -1760,6 +1760,7 @@ absl::Status FunctionConverter::HandleInvocation(const Invocation* node) {
                       decltype(&FunctionConverter::HandleBuiltinClz)>
       map = {
           {"array_rev", &FunctionConverter::HandleBuiltinArrayRev},
+          {"array_size", &FunctionConverter::HandleBuiltinArraySize},
           {"clz", &FunctionConverter::HandleBuiltinClz},
           {"ctz", &FunctionConverter::HandleBuiltinCtz},
           {"gate!", &FunctionConverter::HandleBuiltinGate},
@@ -2739,6 +2740,15 @@ absl::Status FunctionConverter::HandleBuiltinAndReduce(const Invocation* node) {
   return absl::OkStatus();
 }
 
+absl::Status FunctionConverter::HandleBuiltinArraySize(const Invocation* node) {
+  XLS_RET_CHECK_EQ(node->args().size(), 1);
+  // All array sizes are constexpr since they're based on known types.
+  XLS_ASSIGN_OR_RETURN(InterpValue iv, current_type_info_->GetConstExpr(node));
+  XLS_ASSIGN_OR_RETURN(Value v, InterpValueToValue(iv));
+  DefConst(node, v);
+  return absl::OkStatus();
+}
+
 absl::Status FunctionConverter::HandleBuiltinArrayRev(const Invocation* node) {
   XLS_RET_CHECK_EQ(node->args().size(), 1);
   XLS_ASSIGN_OR_RETURN(BValue arg, Use(node->args()[0]));

diff --git a/xls/dslx/ir_convert/function_converter.h b/xls/dslx/ir_convert/function_converter.h
@@ -361,6 +361,7 @@ class FunctionConverter {
   // Builtin invocation handlers.
   absl::Status HandleBuiltinAndReduce(const Invocation* node);
   absl::Status HandleBuiltinArrayRev(const Invocation* node);
+  absl::Status HandleBuiltinArraySize(const Invocation* node);
   absl::Status HandleBuiltinArraySlice(const Invocation* node);
   absl::Status HandleBuiltinBitSlice(const Invocation* node);
   absl::Status HandleBuiltinBitSliceUpdate(const Invocation* node);

diff --git a/xls/dslx/ir_convert/ir_converter_test.cc b/xls/dslx/ir_convert/ir_converter_test.cc
@@ -1921,6 +1921,20 @@ fn main(x: u8, y: s8) -> u32 {
   ExpectIr(converted, TestName());
 }
 
+TEST(IrConverterTest, ArraySizeBuiltin) {
+  constexpr std::string_view program =
+      R"(
+fn main(x: u8[42]) -> u32 {
+  array_size(x)
+}
+)";
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertModuleForTest(program, ConvertOptions{.emit_positions = false}));
+  ExpectIr(converted, TestName());
+}
+
 }  // namespace
 }  // namespace xls::dslx
 

diff --git a/xls/dslx/ir_convert/testdata/ir_converter_test_ArraySizeBuiltin.ir b/xls/dslx/ir_convert/testdata/ir_converter_test_ArraySizeBuiltin.ir
@@ -0,0 +1,7 @@
+package test_module
+
+file_number 0 "test_module.x"
+
+fn __test_module__main(x: bits[8][42]) -> bits[32] {
+  ret literal.2: bits[32] = literal(value=42, id=2)
+}
diff --git a/xls/dslx/type_system/typecheck.cc b/xls/dslx/type_system/typecheck.cc
@@ -356,6 +356,14 @@ absl::StatusOr<TypeAndParametricEnv> CheckParametricBuiltinInvocation(
   if (callee_nameref->identifier() == "widening_cast") {
     XLS_RETURN_IF_ERROR(CheckIsAcceptableWideningCast(ctx, invocation));
   }
+  // array_size is always a constexpr result since it just needs the type
+  // information
+  if (callee_nameref->identifier() == "array_size") {
+    auto* array_type = down_cast<const ArrayType*>(fn_type->params()[0].get());
+    XLS_ASSIGN_OR_RETURN(int64_t array_size, array_type->size().GetAsInt64());
+    ctx->type_info()->NoteConstExpr(invocation,
+                                    InterpValue::MakeU32(array_size));
+  }
 
   // fsignature returns a tab w/a fn type, not the fn return type (which is
   // what we actually want). We hack up `tab` to make this consistent with