Merge pull request #21800 from superbobry:typing

PiperOrigin-RevId: 642224964

Author: jax authors

jax/_src/pallas/triton/lowering.py

@@ -253,9 +253,7 @@ def lower_jaxpr_to_triton_module(
     in_shapes,
     grid_mapping: GridMapping,
     name: str,
-    cuda_options: Any,
 ) -> LoweringResult:
-  # TODO(slebedev): Use cuda_options= during lowering.
   jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), instantiate=True)
   with _new_ir_context(), ir.Location.unknown():
     module = ir.Module.create()

jax/_src/pallas/triton/pallas_call_registration.py

@@ -14,9 +14,6 @@
 
 """Module registering a lowering rule for pallas_call on GPU."""
 
-# TODO(sharadmv): Enable type checking.
-# mypy: ignore-errors
-
 from __future__ import annotations
 
 import io
@@ -36,77 +33,13 @@ def normalize_grid(grid: pallas_core.StaticGrid) -> tuple[int, int, int]:
     grid = (grid,)
   elif len(grid) > 3:
     raise ValueError("`grid` should have three or fewer dimensions.")
-  return tuple(grid) + (1,) * (3 - len(grid))
+  return tuple(grid) + (1,) * (3 - len(grid))  # type: ignore
 
 
 def avals_to_layouts(avals):
   return [list(reversed(range(aval.ndim))) for aval in avals]
 
 
-def _pallas_call_ttir_lowering(
-    ctx: mlir.LoweringRuleContext,
-    *in_nodes,
-    jaxpr: jax_core.Jaxpr,
-    name: str,
-    in_shapes: tuple[jax.ShapeDtypeStruct, ...],
-    out_shapes: tuple[jax.ShapeDtypeStruct, ...],
-    debug: bool,
-    input_output_aliases: tuple[tuple[int, int], ...],
-    grid_mapping: pallas_core.GridMapping,
-    triton_params: dict[str, Any] | None = None,
-    num_warps: int,
-    num_stages: int,
-):
-  # TODO(sharadmv): Handle multiple devices with different capabilities.
-  d, *_ = jax.local_devices(backend="gpu")
-  cuda_options = dict(
-      compute_capability=d.compute_capability,
-      num_warps=num_warps,
-      num_stages=num_stages,
-      debug=debug,
-  )
-
-  lowering_result = lowering.lower_jaxpr_to_triton_module(
-      jaxpr, (*in_shapes, *out_shapes), grid_mapping, name, cuda_options
-  )
-  module_op = lowering_result.module.operation
-  if debug:
-    print(module_op.get_asm(enable_debug_info=True, pretty_debug_info=True))
-
-  grid_x, grid_y, grid_z = normalize_grid(lowering_result.grid)
-  out_types = [
-      ir.RankedTensorType.get(shape.shape, mlir.dtype_to_ir_type(shape.dtype))
-      for shape in out_shapes
-  ]
-  buf = io.BytesIO()
-  module_op.write_bytecode(buf)
-  backend_config = dict(
-      name=ir.StringAttr.get(name),
-      ir=ir.StringAttr.get(buf.getvalue()),
-      num_stages=mlir.i32_attr(num_stages),
-      num_warps=mlir.i32_attr(num_warps),
-      grid_x=mlir.i32_attr(grid_x),
-      grid_y=mlir.i32_attr(grid_y),
-      grid_z=mlir.i32_attr(grid_z),
-      debug=ir.BoolAttr.get(debug),
-  )
-  if "serialized_metadata" in (triton_params or {}):
-    # This field is unstable and may be removed in the future.
-    backend_config["serialized_metadata"] = ir.StringAttr.get(
-        triton_params["serialized_metadata"]
-    )
-  return mlir.custom_call(
-      call_target_name="__gpu$xla.gpu.triton",
-      result_types=out_types,
-      operands=in_nodes,
-      backend_config=backend_config,
-      api_version=4,
-      operand_layouts=avals_to_layouts(ctx.avals_in),
-      result_layouts=avals_to_layouts(ctx.avals_out),
-      operand_output_aliases=dict(input_output_aliases),
-  ).results
-
-
 def pallas_call_lowering(
     ctx: mlir.LoweringRuleContext,
     *in_nodes,
@@ -154,17 +87,42 @@ def pallas_call_lowering(
     print(jaxpr)
     print(grid_mapping)
 
-  return _pallas_call_ttir_lowering(
-        ctx,
-        *in_nodes,
-        jaxpr=jaxpr,
-        name=name,
-        in_shapes=in_shapes,
-        out_shapes=out_shapes,
-        debug=debug,
-        input_output_aliases=input_output_aliases,
-        grid_mapping=grid_mapping,
-        triton_params=triton_params,
-        num_warps=num_warps,
-        num_stages=num_stages,
+  lowering_result = lowering.lower_jaxpr_to_triton_module(
+      jaxpr, (*in_shapes, *out_shapes), grid_mapping, name,
+  )
+  module_op = lowering_result.module.operation
+  if debug:
+    print(module_op.get_asm(enable_debug_info=True, pretty_debug_info=True))
+
+  grid_x, grid_y, grid_z = normalize_grid(lowering_result.grid)
+  out_types = [
+      ir.RankedTensorType.get(shape.shape, mlir.dtype_to_ir_type(shape.dtype))
+      for shape in out_shapes
+  ]
+  buf = io.BytesIO()
+  module_op.write_bytecode(buf)
+  backend_config = dict(
+      name=ir.StringAttr.get(name),
+      ir=ir.StringAttr.get(buf.getvalue()),  # type: ignore
+      num_stages=mlir.i32_attr(num_stages),
+      num_warps=mlir.i32_attr(num_warps),
+      grid_x=mlir.i32_attr(grid_x),
+      grid_y=mlir.i32_attr(grid_y),
+      grid_z=mlir.i32_attr(grid_z),
+      debug=ir.BoolAttr.get(debug),
+  )
+  if "serialized_metadata" in (triton_params or {}):
+    # This field is unstable and may be removed in the future.
+    backend_config["serialized_metadata"] = ir.StringAttr.get(
+        triton_params["serialized_metadata"]
     )
+  return mlir.custom_call(
+      call_target_name="__gpu$xla.gpu.triton",
+      result_types=out_types,
+      operands=in_nodes,
+      backend_config=backend_config,
+      api_version=4,
+      operand_layouts=avals_to_layouts(ctx.avals_in),
+      result_layouts=avals_to_layouts(ctx.avals_out),
+      operand_output_aliases=dict(input_output_aliases),
+  ).results