feat(ir): add DataType.get_byte() to eliminate magic dtype byte-size literals

docs/en/dev/passes/37-synthesize_allreduce_signals.md

@@ -42,7 +42,7 @@ For every host-orchestration function:
 
 ```python
 __allreduce_signal_world_size_0 = pld.system.world_size()
-__allreduce_signal_buf_0: pl.Ptr = pld.tensor.alloc_window_buffer(__allreduce_signal_world_size_0 * 4)
+__allreduce_signal_buf_0: pl.Ptr = pld.tensor.alloc_window_buffer(__allreduce_signal_world_size_0 * pl.INT32.get_byte())
 __allreduce_signal_0 = pld.tensor.window(
     __allreduce_signal_buf_0,
     [__allreduce_signal_world_size_0, 1],

docs/en/user/01-language_guide.md

@@ -436,7 +436,7 @@ def host_orch(
     inputs: pl.Tensor[[2, 1, 256], pl.FP32],
     outputs: pl.Out[pl.Tensor[[2, 1, 256], pl.FP32]],
 ):
-    data_buf = pld.alloc_window_buffer(256 * 4)
+    data_buf = pld.alloc_window_buffer(256 * pl.FP32.get_byte())
     for r in pl.range(pld.world_size()):
         data = pld.window(data_buf, [1, 256], dtype=pl.FP32)
         chip_orch(inputs[r], outputs[r], data, (r + 1) % pld.world_size(),

docs/zh-cn/dev/passes/37-synthesize_allreduce_signals.md

@@ -38,7 +38,7 @@ allocation 处理，并放入 allreduce data buffer 所属的通信域。
 
 ```python
 __allreduce_signal_world_size_0 = pld.system.world_size()
-__allreduce_signal_buf_0: pl.Ptr = pld.tensor.alloc_window_buffer(__allreduce_signal_world_size_0 * 4)
+__allreduce_signal_buf_0: pl.Ptr = pld.tensor.alloc_window_buffer(__allreduce_signal_world_size_0 * pl.INT32.get_byte())
 __allreduce_signal_0 = pld.tensor.window(
     __allreduce_signal_buf_0,
     [__allreduce_signal_world_size_0, 1],

docs/zh-cn/user/01-language_guide.md

@@ -420,7 +420,7 @@ def host_orch(
     inputs: pl.Tensor[[2, 1, 256], pl.FP32],
     outputs: pl.Out[pl.Tensor[[2, 1, 256], pl.FP32]],
 ):
-    data_buf = pld.alloc_window_buffer(256 * 4)
+    data_buf = pld.alloc_window_buffer(256 * pl.FP32.get_byte())
     for r in pl.range(pld.world_size()):
         data = pld.window(data_buf, [1, 256], dtype=pl.FP32)
         chip_orch(inputs[r], outputs[r], data, (r + 1) % pld.world_size(),

include/pypto/core/dtype.h

@@ -182,6 +182,16 @@ class DataType {
     }
   }
 
+  /**
+   * @brief Get the size in bytes of this data type
+   *
+   * Returns ceil(GetBit() / 8). For sub-byte types (INT4, UINT4, FP4, HF4, BOOL)
+   * this returns 1 since window buffers operate at byte granularity.
+   *
+   * @return Size in bytes
+   */
+  [[nodiscard]] size_t GetByte() const { return (GetBit() + 7) / 8; }
+
   /**
    * @brief Get a human-readable string name for this data type
    *

python/bindings/modules/core.cpp

@@ -59,6 +59,9 @@ void BindCore(nb::module_& m) {
            "Get the size in bits of this data type. Returns the actual bit size for sub-byte types (e.g., 4 "
            "bits "
            "for INT4).")
+      .def("get_byte", &DataType::GetByte,
+           "Get the size in bytes of this data type (ceil(get_bit() / 8)). Returns 1 for sub-byte types "
+           "(e.g., INT4, BOOL).")
       .def("to_string", &DataType::ToString, "Get a human-readable string name for this data type.")
       .def("to_c_type_string", &DataType::ToCTypeString,
            "Get C style type string for code generation (e.g., 'float', 'half', 'int32_t').")

python/pypto/language/distributed/op/system_ops.py

@@ -28,7 +28,7 @@
 Typical use sites for ``world_size``:
 
 * loop bounds: ``for r in pl.range(pld.world_size()): ...``
-* allocation sizes (in bytes): ``pld.tensor.alloc_window_buffer(pld.world_size() * 4)``
+* allocation sizes (in bytes): ``pld.tensor.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())``
 * per-rank tensor shapes:
   ``pld.tensor.window(buf, [pld.world_size()], dtype=pl.INT32)``
 """

python/pypto/language/parser/ast_parser.py

@@ -5286,6 +5286,14 @@ def parse_op_call(self, call: ast.Call) -> ir.Expr:
         if len(attrs) >= 2 and attrs[0] == "pl" and attrs[1] == "const":
             return self._parse_typed_constant(call)
 
+        # pl.<DataType>.get_byte() — fold into compile-time int constant.
+        # e.g. ``pl.FP32.get_byte()`` → 4, ``pl.INT32.get_byte()`` → 4.
+        # The call has no side effects (GetByte is a pure accessor) so folding
+        # at parse time is safe and lets it work directly inside traced bodies
+        # without a separate import of the _window_byte_constants module.
+        if len(attrs) == 3 and attrs[0] == "pl" and attrs[-1] == "get_byte":
+            return self._parse_dtype_get_byte(attrs[1], call)
+
         # pl.{operation} (2-segment, unified dispatch or promoted ops)
         if len(attrs) >= 2 and attrs[0] == "pl" and attrs[1] not in ("tensor", "tile", "system", "array"):
             op_name = attrs[1]
@@ -7280,6 +7288,50 @@ def _parse_typed_constant(self, call: ast.Call) -> ir.Expr:
         else:
             return ir.ConstInt(value, dtype, span)
 
+    def _parse_dtype_get_byte(self, dtype_name: str, call: ast.Call) -> ir.Expr:
+        """Parse ``pl.<DataType>.get_byte()`` as a compile-time int constant.
+
+        ``DataType.get_byte()`` is a pure accessor returning ``ceil(bit_width / 8)``
+        with no side effects, so folding at parse time is safe. This lets traced
+        ``@pl.function`` / ``@pl.jit.host`` bodies use the canonical API directly:
+        ``pld.alloc_window_buffer(256 * pl.FP32.get_byte())`` works without an
+        external import.
+
+        Args:
+            dtype_name: The DataType attribute name (e.g. "FP32", "INT32", "BF16").
+            call: The Call AST node (for span tracking).
+
+        Returns:
+            ``ir.ConstInt(nbytes, INDEX)`` with the byte-size of the data type.
+        """
+        span = self.span_tracker.get_span(call)
+
+        # Lazily build the mapping from dtype attribute name → get_byte() value.
+        # Module-level caching via a class attribute avoids re-scanning on every
+        # call site within a single parser instance (a function body may contain
+        # many alloc_window_buffer calls).
+        try:
+            cache = ASTParser._dtype_byte_cache  # type: ignore[attr-defined]
+        except AttributeError:
+            import pypto.language as _pl  # noqa: PLC0415 (lazy import of DataType constants)
+
+            cache: dict[str, int] = {}
+            for attr_name in dir(_pl):
+                val = getattr(_pl, attr_name)
+                if isinstance(val, DataType):
+                    cache[attr_name] = val.get_byte()
+            ASTParser._dtype_byte_cache = cache  # type: ignore[attr-defined]
+
+        if dtype_name not in cache:
+            known = ", ".join(sorted(cache.keys()))
+            raise ParserTypeError(
+                f"Unknown DataType '{dtype_name}' in pl.{dtype_name}.get_byte()",
+                span=span,
+                hint=f"Valid DataType names: {known}",
+            )
+
+        return ir.ConstInt(cache[dtype_name], DataType.INDEX, span)
+
     def parse_attribute(self, attr: ast.Attribute) -> ir.Expr:
         """Parse attribute access.
 

python/pypto/pypto_core/__init__.pyi

@@ -66,6 +66,15 @@ class DataType:
             The size in bits of the data type
         """
 
+    def get_byte(self) -> int:
+        """
+        Get the size in bytes of this data type (ceil(get_bit() / 8)).
+        Returns 1 for sub-byte types (e.g., INT4, BOOL).
+
+        Returns:
+            The size in bytes of the data type
+        """
+
     def __hash__(self) -> int: ...
     def to_string(self) -> str:
         """

tests/st/distributed/collectives/test_l3_all_to_all.py

@@ -126,8 +126,8 @@ def host_orch(
         inputs: pl.Tensor[[NR, NR, SIZE], pl.FP32],
         outputs: pl.Out[pl.Tensor[[NR, NR, SIZE], pl.FP32]],
     ) -> pl.Tensor[[NR, NR, SIZE], pl.FP32]:
-        data_buf = pld.alloc_window_buffer(NR * SIZE * 4)
-        signal_buf = pld.alloc_window_buffer(NR * 4)
+        data_buf = pld.alloc_window_buffer(NR * SIZE * pl.FP32.get_byte())
+        signal_buf = pld.alloc_window_buffer(NR * pl.INT32.get_byte())
 
         for r in pl.range(pld.world_size()):
             data = pld.window(data_buf, [NR, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_allgather.py

@@ -127,8 +127,8 @@ def host_orch(
         outputs: pl.Out[pl.Tensor[[NR, 1, NR * SIZE], pl.FP32]],
     ) -> pl.Tensor[[NR, 1, NR * SIZE], pl.FP32]:
         """Launch one chip orchestration per rank with shared window buffers."""
-        scratch_buf = pld.alloc_window_buffer(SIZE * 4)  # 1xSIZE x FP32 (4 bytes)
-        signal_buf = pld.alloc_window_buffer(pld.world_size() * 4)  # NR x 1 x INT32
+        scratch_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+        signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())
 
         for r in pl.range(pld.world_size()):
             scratch = pld.window(scratch_buf, [1, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_allreduce.py

@@ -134,8 +134,8 @@ def host_orch(
         outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]],
     ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]:
         """Launch one chip orchestration per rank with shared window buffers."""
-        data_buf = pld.alloc_window_buffer(SIZE * 4)  # 1xSIZE x FP32 (4 bytes)
-        signal_buf = pld.alloc_window_buffer(pld.world_size() * 4)  # NR x 1 x INT32
+        data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+        signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())
 
         for r in pl.range(pld.world_size()):
             data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_allreduce_ring.py

@@ -201,8 +201,8 @@ def host_orch(
             outputs: pl.Out[pl.Tensor[[n_ranks, 1, SIZE], pl.FP32]],
         ) -> pl.Tensor[[n_ranks, 1, SIZE], pl.FP32]:
             """Launch one chip orchestration per rank with shared window buffers."""
-            scratch_buf = pld.alloc_window_buffer(SIZE * 4)  # NR × chunk × FP32
-            signal_buf = pld.alloc_window_buffer(total_rounds * n_ranks * 4)  # rounds × NR × INT32
+            scratch_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(total_rounds * n_ranks * pl.INT32.get_byte())
 
             chunk_elems = SIZE // n_ranks
             for r in pl.range(n_ranks):

tests/st/distributed/collectives/test_l3_broadcast.py

@@ -127,8 +127,8 @@ def host_orch(
         outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]],
     ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]:
         """Launch one chip orchestration per rank with shared window buffers."""
-        scratch_buf = pld.alloc_window_buffer(SIZE * 4)  # 1xSIZE x FP32 (4 bytes)
-        signal_buf = pld.alloc_window_buffer(pld.world_size() * 4)  # NR x 1 x INT32
+        scratch_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+        signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())
 
         for r in pl.range(pld.world_size()):
             scratch = pld.window(scratch_buf, [1, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_reduce_scatter.py

@@ -136,8 +136,8 @@ def host_orch(
         outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]],
     ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]:
         """Launch one chip orchestration per rank with shared window buffers."""
-        scratch_buf = pld.alloc_window_buffer(pld.world_size() * SIZE * 4)  # NR x SIZE FP32
-        signal_buf = pld.alloc_window_buffer(pld.world_size() * 4)  # NR x 1 x INT32
+        scratch_buf = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte())
+        signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())
 
         for r in pl.range(pld.world_size()):
             scratch = pld.window(scratch_buf, [1, pld.world_size() * SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_tensor_all_to_all_intrinsic.py

@@ -90,8 +90,8 @@ def host_orch(
             inputs: pl.Tensor[[nr, nr, SIZE], pl.FP32],
             outputs: pl.Out[pl.Tensor[[nr, nr, SIZE], pl.FP32]],
         ) -> pl.Tensor[[nr, nr, SIZE], pl.FP32]:
-            data_buf = pld.alloc_window_buffer(nr * SIZE * 4)
-            signal_buf = pld.alloc_window_buffer(nr * 4)
+            data_buf = pld.alloc_window_buffer(nr * SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(nr * pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 data = pld.window(data_buf, [nr, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_tensor_allgather_intrinsic.py

@@ -90,8 +90,8 @@ def host_orch(
             inputs: pl.Tensor[[nr, 1, SIZE], pl.FP32],
             outputs: pl.Out[pl.Tensor[[nr, 1, nr * SIZE], pl.FP32]],
         ) -> pl.Tensor[[nr, 1, nr * SIZE], pl.FP32]:
-            data_buf = pld.alloc_window_buffer(nr * SIZE * 4)
-            signal_buf = pld.alloc_window_buffer(nr * 4)
+            data_buf = pld.alloc_window_buffer(nr * SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(nr * pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 data = pld.window(data_buf, [nr, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_tensor_allreduce_intrinsic.py

@@ -110,8 +110,8 @@ def host_orch(
             does not synthesise it; this matches the established pattern in
             ``test_l3_allreduce.py``.
             """
-            data_buf = pld.alloc_window_buffer(SIZE * 4)  # 1xSIZE x FP32 (4 bytes)
-            signal_buf = pld.alloc_window_buffer(pld.world_size() * 4)  # nr x 1 x INT32
+            data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.py

@@ -116,8 +116,8 @@ def host_orch(
             ``alloc_window_buffer`` zero-initialises so per-round
             ``AtomicAdd(0→1)`` / ``WaitGe(1)`` works without explicit reset.
             """
-            data_buf = pld.alloc_window_buffer(SIZE * 4)  # 1 x SIZE x FP32
-            signal_buf = pld.alloc_window_buffer(total_rounds * nr * 4)  # rounds × NR × INT32
+            data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(total_rounds * nr * pl.INT32.get_byte())
 
             for r in pl.range(nr):
                 data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_tensor_barrier_intrinsic.py

@@ -94,8 +94,8 @@ def host_orch(
             inputs: pl.Tensor[[nr, 1, SIZE], pl.FP32],
             outputs: pl.Out[pl.Tensor[[nr, 1, SIZE], pl.FP32]],
         ) -> pl.Tensor[[nr, 1, SIZE], pl.FP32]:
-            data_buf = pld.alloc_window_buffer(SIZE * 4)  # 1xSIZE x FP32
-            signal_buf = pld.alloc_window_buffer(nr * 4)  # nr x 1 x INT32
+            data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(nr * pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_tensor_broadcast_intrinsic.py

@@ -95,8 +95,8 @@ def host_orch(
             inputs: pl.Tensor[[nr, 1, SIZE], pl.FP32],
             outputs: pl.Out[pl.Tensor[[nr, 1, SIZE], pl.FP32]],
         ) -> pl.Tensor[[nr, 1, SIZE], pl.FP32]:
-            data_buf = pld.alloc_window_buffer(SIZE * 4)
-            signal_buf = pld.alloc_window_buffer(nr * 4)  # nr x 1 x INT32
+            data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(nr * pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32)

tests/st/distributed/collectives/test_l3_tensor_reduce_scatter_intrinsic.py

@@ -93,8 +93,8 @@ def host_orch(
             inputs: pl.Tensor[[nr, 1, nr * SIZE], pl.FP32],
             outputs: pl.Out[pl.Tensor[[nr, 1, SIZE], pl.FP32]],
         ) -> pl.Tensor[[nr, 1, SIZE], pl.FP32]:
-            data_buf = pld.alloc_window_buffer(nr * SIZE * 4)
-            signal_buf = pld.alloc_window_buffer(nr * 4)
+            data_buf = pld.alloc_window_buffer(nr * SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(nr * pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 data = pld.window(data_buf, [nr, SIZE], dtype=pl.FP32)

tests/st/distributed/test_l3_ep_dispatch_combine.py

@@ -469,14 +469,14 @@ def host_orch(
             # Window allocations — one buffer per cross-rank slot. Barrier
             # signals are sized [N_RANKS, 1] to host per-src cells, matching
             # count_done_sig[N] / data_done_sig[N] / combine_done_sig[N].
-            pub_counts_buf = pld.alloc_window_buffer(N_RANKS * N_RANKS * L * 4)  # INT32
-            count_done_buf = pld.alloc_window_buffer(N_RANKS * 4)
-            recv_x_buf = pld.alloc_window_buffer(L * R * D * 2)  # BF16
-            recv_w_buf = pld.alloc_window_buffer(L * R * W_PAD * 4)  # FP32
-            recv_idx_buf = pld.alloc_window_buffer(L * R * IDX_PAD * 4)  # INT32
-            data_done_buf = pld.alloc_window_buffer(N_RANKS * 4)
-            routed_y_buf_buf = pld.alloc_window_buffer(N_ROUTES * D * 2)  # BF16
-            combine_done_buf = pld.alloc_window_buffer(N_RANKS * 4)
+            pub_counts_buf = pld.alloc_window_buffer(N_RANKS * N_RANKS * L * pl.INT32.get_byte())
+            count_done_buf = pld.alloc_window_buffer(N_RANKS * pl.INT32.get_byte())
+            recv_x_buf = pld.alloc_window_buffer(L * R * D * pl.BF16.get_byte())
+            recv_w_buf = pld.alloc_window_buffer(L * R * W_PAD * pl.FP32.get_byte())
+            recv_idx_buf = pld.alloc_window_buffer(L * R * IDX_PAD * pl.INT32.get_byte())
+            data_done_buf = pld.alloc_window_buffer(N_RANKS * pl.INT32.get_byte())
+            routed_y_buf_buf = pld.alloc_window_buffer(N_ROUTES * D * pl.BF16.get_byte())
+            combine_done_buf = pld.alloc_window_buffer(N_RANKS * pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 pub_counts = pld.window(pub_counts_buf, [N_RANKS * N_RANKS, L], dtype=pl.INT32)

tests/st/distributed/test_l3_gemm.py

@@ -108,7 +108,7 @@ def host_orch(
             b: pl.Tensor[[K, N], pl.FP32],
             c: pl.Out[pl.Tensor[[2, M0, N], pl.FP32]],
         ) -> pl.Tensor[[2, M0, N], pl.FP32]:
-            scratch_buf = pld.alloc_window_buffer(COMM_ANCHOR_COLS * 4)  # 1x8 x INT32 (32 B)
+            scratch_buf = pld.alloc_window_buffer(COMM_ANCHOR_COLS * pl.INT32.get_byte())
             for r in pl.range(pld.world_size()):
                 scratch = pld.window(scratch_buf, [1, COMM_ANCHOR_COLS], dtype=pl.INT32)
                 self.chip_orch(a[r], b, c[r], scratch, device=r)

tests/st/distributed/test_l3_get.py

@@ -107,9 +107,9 @@ def host_orch(
             inputs: pl.Tensor[[2, 1, SIZE], pl.FP32],
             outputs: pl.Out[pl.Tensor[[2, 1, SIZE], pl.FP32]],
         ) -> pl.Tensor[[2, 1, SIZE], pl.FP32]:
-            src_buf = pld.alloc_window_buffer(SIZE * 4)  # 1xSIZE x FP32
-            dst_buf = pld.alloc_window_buffer(SIZE * 4)
-            signal_buf = pld.alloc_window_buffer(4)  # 1x1 x INT32
+            src_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+            dst_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(pl.INT32.get_byte())  # 1x1 x INT32
 
             for r in pl.range(pld.world_size()):
                 src = pld.window(src_buf, [1, SIZE], dtype=pl.FP32)
@@ -172,9 +172,9 @@ def host_orch(
             inputs: pl.Tensor[[2, 2, SIZE], pl.FP32],
             outputs: pl.Out[pl.Tensor[[2, 1, SIZE], pl.FP32]],
         ) -> pl.Tensor[[2, 1, SIZE], pl.FP32]:
-            src_buf = pld.alloc_window_buffer(2 * SIZE * 4)
-            dst_buf = pld.alloc_window_buffer(2 * SIZE * 4)
-            signal_buf = pld.alloc_window_buffer(4)
+            src_buf = pld.alloc_window_buffer(2 * SIZE * pl.FP32.get_byte())
+            dst_buf = pld.alloc_window_buffer(2 * SIZE * pl.FP32.get_byte())
+            signal_buf = pld.alloc_window_buffer(pl.INT32.get_byte())
 
             for r in pl.range(pld.world_size()):
                 src = pld.window(src_buf, [2, SIZE], dtype=pl.FP32)

tests/st/distributed/test_l3_host_tensor_allgather.py

@@ -85,8 +85,8 @@ def host_orch(
         inputs: pl.Tensor[[NR, 1, SIZE], pl.FP32],
         outputs: pl.Out[pl.Tensor[[NR, 1, NR, SIZE], pl.FP32]],
     ) -> pl.Tensor[[NR, 1, NR, SIZE], pl.FP32]:
-        data_buf = pld.alloc_window_buffer(pld.world_size() * SIZE * 4)
-        signal_buf = pld.alloc_window_buffer(pld.world_size() * 4)
+        data_buf = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte())
+        signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())
 
         for r in pl.range(pld.world_size()):
             data = pld.window(data_buf, [pld.world_size(), SIZE], dtype=pl.FP32)

tests/st/distributed/test_l3_host_tensor_allreduce.py

@@ -77,8 +77,8 @@ def host_orch(
         inputs: pl.Tensor[[NR, 1, SIZE], pl.FP32],
         outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]],
     ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]:
-        data_buf = pld.alloc_window_buffer(SIZE * 4)
-        signal_buf = pld.alloc_window_buffer(pld.world_size() * 4)
+        data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte())
+        signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte())
 
         for r in pl.range(pld.world_size()):
             data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32)