Make codegen implement shra as a function.

This allows us to bind the intermediate result `$signed(to_shift) >>> amount` to a reg with a defined width, which avoids the `$signed()` having a self-determined width, which is a lint finding.

Fixes #1642.

PiperOrigin-RevId: 682526771

xls/codegen/BUILD

@@ -281,6 +281,8 @@ cc_library(
     deps = [
         "//xls/codegen/vast",
         "//xls/ir:source_location",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/strings:str_format",
     ],
 )

xls/codegen/lint_annotate.h

@@ -19,7 +19,8 @@
 #include <utility>
 #include <vector>
 
-#include "absl/strings/str_format.h"
+#include "absl/base/optimization.h"
+#include "absl/log/log.h"
 #include "xls/codegen/vast/vast.h"
 
 namespace xls::verilog {
@@ -59,7 +60,8 @@ inline std::string LintToString(Lint flag) {
     case Lint::kMultiply:
       return "MULTIPLY";
   }
-  return absl::StrFormat("<invalid Lint(%d)>", static_cast<int>(flag));
+  LOG(FATAL) << "Unknown lint flag: " << static_cast<int>(flag);
+  ABSL_UNREACHABLE();
 }
 
 }  // namespace xls::verilog

xls/codegen/module_builder.cc

@@ -1437,6 +1437,7 @@ bool ModuleBuilder::MustEmitAsFunction(Node* node) {
     case Op::kBitSliceUpdate:
     case Op::kSDiv:
     case Op::kUDiv:
+    case Op::kShra:
       return true;
     case Op::kPrioritySel:
       return node->As<PrioritySelect>()->cases().size() > 1;
@@ -1486,6 +1487,11 @@ absl::StatusOr<std::string> ModuleBuilder::VerilogFunctionName(Node* node) {
       CHECK_EQ(node->BitCountOrDie(), node->operand(1)->BitCountOrDie());
       return absl::StrFormat("%s_%db", OpToString(node->op()),
                              node->BitCountOrDie());
+    case Op::kShra:
+      CHECK_EQ(node->BitCountOrDie(), node->operand(0)->BitCountOrDie());
+      return absl::StrFormat("%s_%db_by_%db", OpToString(node->op()),
+                             node->BitCountOrDie(),
+                             node->operand(1)->BitCountOrDie());
     default:
       LOG(FATAL) << "Cannot emit node as function: " << node->ToString();
   }
@@ -1905,6 +1911,68 @@ VerilogFunction* DefineSDivFunction(Node* node, std::string_view function_name,
   return func;
 }
 
+VerilogFunction* DefineShraFunction(Node* node, std::string_view function_name,
+                                    ModuleSection* section) {
+  CHECK_EQ(node->op(), Op::kShra);
+  VerilogFile* file = section->file();
+
+  // Disable lint for the signed_result reg, which is declared signed to catch
+  // the output of `$signed(to_shift) >> amount`.
+  ScopedLintDisable lint_disable(section, {Lint::kSignedType});
+
+  VerilogFunction* func = section->Add<VerilogFunction>(
+      node->loc(), function_name,
+      file->BitVectorType(node->BitCountOrDie(), node->loc()));
+  CHECK_EQ(node->operand_count(), 2);
+
+  IndexableExpression* to_shift = func->AddArgument(
+      "to_shift",
+      file->BitVectorType(node->operand(0)->BitCountOrDie(), node->loc()),
+      node->loc());
+  IndexableExpression* shift_amount = func->AddArgument(
+      "shift_amount",
+      file->BitVectorType(node->operand(1)->BitCountOrDie(), node->loc()),
+      node->loc());
+
+  LogicRef* signed_result =
+      func->AddRegDef(node->loc(), "signed_result",
+                      file->BitVectorType(node->BitCountOrDie(), node->loc(),
+                                          /*is_signed=*/true),
+                      /*init=*/nullptr);
+
+  // To perform an arithmetic shift right the left operand must be cast to a
+  // signed value, ie:
+  //
+  //   signed_result := $signed(x) >>> y
+  //
+  // We bind the intermediate result to a variable to prevent the signed
+  // result from having a self-determined width. Self-determined widths of
+  // signed values is normally suspicious because it can result in accidentally
+  // dropping sign-extended bits. In our case, it should be safe because we'll
+  // immediately cast to unsigned, but linters can still flag this. When we bind
+  // the signed result to a variable, the width is no longer self-determined.
+  //
+  // Ultimately, we cast the expression with $unsigned and return the value as
+  // an unsigned type to prevent usages from leaking the signed property.
+  //
+  //   $unsigned(signed_result)
+  //
+  // Without the unsigned, the '>>>' expression would be treated as a signed
+  // value potentially affecting the evaluation of 'op'.  This unsigned cast
+  // is also necessary for correctness of the shift evaluation when the shift
+  // appears in a ternary expression because of Verilog type rules.
+  func->AddStatement<BlockingAssignment>(
+      node->loc(), signed_result,
+      file->Shra(file->Make<SignedCast>(node->loc(), to_shift), shift_amount,
+                 node->loc()));
+
+  func->AddStatement<BlockingAssignment>(
+      node->loc(), func->return_value_ref(),
+      file->Make<UnsignedCast>(node->loc(), signed_result));
+
+  return func;
+}
+
 }  // namespace
 
 absl::StatusOr<VerilogFunction*> ModuleBuilder::DefineFunction(Node* node) {
@@ -1957,6 +2025,9 @@ absl::StatusOr<VerilogFunction*> ModuleBuilder::DefineFunction(Node* node) {
     case Op::kSDiv:
       func = DefineSDivFunction(node, function_name, functions_section_);
       break;
+    case Op::kShra:
+      func = DefineShraFunction(node, function_name, functions_section_);
+      break;
     default:
       LOG(FATAL) << "Cannot define node as function: " << node->ToString();
   }

xls/codegen/module_builder_test.cc

@@ -766,6 +766,46 @@ TEST_P(ModuleBuilderTest, DifferentArrayPrioritySelects) {
                                  file.Emit());
 }
 
+TEST_P(ModuleBuilderTest, ShraAsFunction) {
+  VerilogFile file = NewVerilogFile();
+  Package package(TestBaseName());
+  FunctionBuilder fb(TestBaseName(), &package);
+  Type* u32 = package.GetBitsType(32);
+  BValue x_param = fb.Param("x", u32);
+  BValue y_param = fb.Param("y", u32);
+  BValue x_shra_y = fb.Shra(x_param, y_param);
+  XLS_ASSERT_OK(fb.Build());
+
+  ModuleBuilder mb(TestBaseName(), &file, codegen_options());
+  XLS_ASSERT_OK_AND_ASSIGN(LogicRef * x, mb.AddInputPort("x", u32));
+  XLS_ASSERT_OK_AND_ASSIGN(LogicRef * y, mb.AddInputPort("y", u32));
+  XLS_ASSERT_OK(
+      mb.EmitAsAssignment("x_smul_y", x_shra_y.node(), {x, y}).status());
+
+  ExpectVerilogEqualToGoldenFile(GoldenFilePath(kTestName, kTestdataPath),
+                                 file.Emit());
+}
+
+TEST_P(ModuleBuilderTest, ShraAsFunctionSingleBit) {
+  VerilogFile file = NewVerilogFile();
+  Package package(TestBaseName());
+  FunctionBuilder fb(TestBaseName(), &package);
+  Type* u1 = package.GetBitsType(1);
+  BValue x_param = fb.Param("x", u1);
+  BValue y_param = fb.Param("y", u1);
+  BValue x_shra_y = fb.Shra(x_param, y_param);
+  XLS_ASSERT_OK(fb.Build());
+
+  ModuleBuilder mb(TestBaseName(), &file, codegen_options());
+  XLS_ASSERT_OK_AND_ASSIGN(LogicRef * x, mb.AddInputPort("x", u1));
+  XLS_ASSERT_OK_AND_ASSIGN(LogicRef * y, mb.AddInputPort("y", u1));
+  XLS_ASSERT_OK(
+      mb.EmitAsAssignment("x_smul_y", x_shra_y.node(), {x, y}).status());
+
+  ExpectVerilogEqualToGoldenFile(GoldenFilePath(kTestName, kTestdataPath),
+                                 file.Emit());
+}
+
 INSTANTIATE_TEST_SUITE_P(ModuleBuilderTestInstantiation, ModuleBuilderTest,
                          testing::ValuesIn(kDefaultSimulationTargets),
                          ParameterizedTestName<ModuleBuilderTest>);

xls/codegen/node_expressions.cc

@@ -310,25 +310,9 @@ absl::StatusOr<Expression*> EmitShift(Node* shift, Expression* operand,
                                       Expression* shift_amount,
                                       VerilogFile* file) {
   Expression* shifted_operand;
-  if (shift->op() == Op::kShra) {
-    // To perform an arithmetic shift right the left operand must be cast to a
-    // signed value, ie:
-    //
-    //   $signed(x) >>> y
-    //
-    // Also, wrap the expression in $unsigned to prevent the signed property
-    // from leaking out into the rest of the expression.
-    //
-    //   $unsigned($signed(x) >>> y) op ...
-    //
-    // Without the unsigned the '>>>' expression would be treated as a signed
-    // value potentially affecting the evaluation of 'op'.  This unsigned cast
-    // is also necessary for correctness of the shift evaluation when the shift
-    // appears in a ternary expression because of Verilog type rules.
-    shifted_operand = file->Make<UnsignedCast>(
-        shift->loc(), file->Shra(file->Make<SignedCast>(shift->loc(), operand),
-                                 shift_amount, shift->loc()));
-  } else if (shift->op() == Op::kShrl) {
+  XLS_RET_CHECK_NE(shift->op(), Op::kShra)
+      << absl::StreamFormat("Shra is handled by emitting a function.");
+  if (shift->op() == Op::kShrl) {
     shifted_operand = file->Shrl(operand, shift_amount, shift->loc());
   } else {
     CHECK_EQ(shift->op(), Op::kShll);

xls/codegen/testdata/module_builder_test_ShraAsFunction.svtxt

@@ -0,0 +1,16 @@
+module ShraAsFunction(
+  input wire [31:0] x,
+  input wire [31:0] y
+);
+  // lint_off SIGNED_TYPE
+  function automatic [31:0] shra_32b_by_32b (input reg [31:0] to_shift, input reg [31:0] shift_amount);
+    reg signed [31:0] signed_result;
+    begin
+      signed_result = $signed(to_shift) >>> shift_amount;
+      shra_32b_by_32b = $unsigned(signed_result);
+    end
+  endfunction
+  // lint_on SIGNED_TYPE
+  wire [31:0] x_smul_y;
+  assign x_smul_y = shra_32b_by_32b(x, y);
+endmodule

xls/codegen/testdata/module_builder_test_ShraAsFunction.vtxt

@@ -0,0 +1,16 @@
+module ShraAsFunction(
+  input wire [31:0] x,
+  input wire [31:0] y
+);
+  // lint_off SIGNED_TYPE
+  function automatic [31:0] shra_32b_by_32b (input reg [31:0] to_shift, input reg [31:0] shift_amount);
+    reg signed [31:0] signed_result;
+    begin
+      signed_result = $signed(to_shift) >>> shift_amount;
+      shra_32b_by_32b = $unsigned(signed_result);
+    end
+  endfunction
+  // lint_on SIGNED_TYPE
+  wire [31:0] x_smul_y;
+  assign x_smul_y = shra_32b_by_32b(x, y);
+endmodule

xls/codegen/testdata/module_builder_test_ShraAsFunctionSingleBit.svtxt

@@ -0,0 +1,16 @@
+module ShraAsFunctionSingleBit(
+  input wire x,
+  input wire y
+);
+  // lint_off SIGNED_TYPE
+  function automatic logic shra_1b_by_1b (input reg to_shift, input reg shift_amount);
+    reg signed signed_result;
+    begin
+      signed_result = $signed(to_shift) >>> shift_amount;
+      shra_1b_by_1b = $unsigned(signed_result);
+    end
+  endfunction
+  // lint_on SIGNED_TYPE
+  wire x_smul_y;
+  assign x_smul_y = shra_1b_by_1b(x, y);
+endmodule

xls/codegen/testdata/module_builder_test_ShraAsFunctionSingleBit.vtxt

@@ -0,0 +1,16 @@
+module ShraAsFunctionSingleBit(
+  input wire x,
+  input wire y
+);
+  // lint_off SIGNED_TYPE
+  function automatic shra_1b_by_1b (input reg to_shift, input reg shift_amount);
+    reg signed signed_result;
+    begin
+      signed_result = $signed(to_shift) >>> shift_amount;
+      shra_1b_by_1b = $unsigned(signed_result);
+    end
+  endfunction
+  // lint_on SIGNED_TYPE
+  wire x_smul_y;
+  assign x_smul_y = shra_1b_by_1b(x, y);
+endmodule