Merge pull request #20062 from Micky774:strip_ir

PiperOrigin-RevId: 621930992

Author: jax authors

jax/_src/interpreters/mlir.py

@@ -75,6 +75,13 @@
          "Supports the special value 'sponge' to pick the path from the "
          "environment variable TEST_UNDECLARED_OUTPUTS_DIR.")
 
+_JAX_INCLUDE_DEBUG_INFO_IN_DUMPS = config.DEFINE_string(
+  'jax_include_debug_info_in_dumps',
+  os.getenv('JAX_INCLUDE_DEBUG_INFO_IN_DUMPS', "True"),
+  help="Determine whether or not to keep debug symbols and location information "
+       "when dumping IR code. By default, debug information will be preserved in "
+       "the IR dump. To avoid exposing source code and potentially sensitive "
+       "information, set to false")
 lowerable_effects: effects_lib.EffectTypeSet = effects_lib.lowerable_effects
 
 
@@ -474,9 +481,12 @@ def dump_module_message(module: ir.Module, stage_name: str) -> str:
 def _make_string_safe_for_filename(s: str) -> str:
   return re.sub(r'[^\w.)( -]', '', s)
 
-def module_to_string(module: ir.Module) -> str:
+def module_to_string(module: ir.Module, enable_debug_info=None) -> str:
   output = io.StringIO()
-  module.operation.print(file=output, enable_debug_info=True)
+  if enable_debug_info is None:
+    enable_debug_flag = str.lower(_JAX_INCLUDE_DEBUG_INFO_IN_DUMPS.value)
+    enable_debug_info = enable_debug_flag not in ('false', '0')
+  module.operation.print(file=output, enable_debug_info=enable_debug_info)
   return output.getvalue()
 
 def module_to_bytecode(module: ir.Module) -> bytes:

tests/api_test.py

@@ -1227,6 +1227,14 @@ def f(x, y, *args, **kwargs):
     self.assertIn("kwargs['z']", hlo_str)
     self.assertIn("kwargs['w']", hlo_str)
 
+    hlo_str = mlir.module_to_string(
+      lowered.compiler_ir('stablehlo'),
+      enable_debug_info=False,
+    )
+    for s in ("\"x\"", "y['hi']", "args[0]", "args[1]", "kwargs['z']", "kwargs['w']"):
+      self.assertNotIn(s, hlo_str)
+
+
   @parameterized.parameters([0, 2, [(0, 2)]])
   def test_jit_lower_arg_info_static_argnums(self, static_argnums):
     def f(x, y, *args, **kwargs):
@@ -1242,6 +1250,10 @@ def f(x, y, *args, **kwargs):
     self.assertIn("kwargs['z']", hlo_str)
     self.assertIn("kwargs['w']", hlo_str)
 
+    hlo_str = mlir.module_to_string(ir, enable_debug_info=False)
+    for s in ("\"x\"", "y['hi']", "args[0]", "args[1]", "kwargs['z']", "kwargs['w']"):
+      self.assertNotIn(s, hlo_str)
+
   @parameterized.parameters(['a', 'b', [('a', 'b')]])
   def test_jit_lower_arg_info_static_argnames(self, static_argnames):
     def f(x, y, *args, **kwargs):
@@ -1259,6 +1271,13 @@ def f(x, y, *args, **kwargs):
     self.assertNotIn("kwargs['a']", hlo_str)
     self.assertNotIn("kwargs['b']", hlo_str)
 
+    hlo_str = mlir.module_to_string(ir, enable_debug_info=False)
+    for s in (
+      "\"x\"", "y['hi']", "args[0]", "args[1]", "kwargs['z']",
+      "kwargs['w']", "kwargs['a']", "kwargs['b']"
+    ):
+      self.assertNotIn(s, hlo_str)
+
   def test_jit_lower_result_info(self):
     def f(x, y, z):
       return {'a': x, 'b': [y]}