From 8896320a7c3e287ba7c2b095bf181bdf58713d1a Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 17:03:51 -0800 Subject: [PATCH 01/10] cake_kda_det_train_kcp: the affine merge selects a register-blocked build per launch shape The merge program ships one record per register-blocked build (column block, thread tile, packed FMA pairs on sm_100a / sm_103a); every build gives the same bits as the previous one. A new select rule kind, "chains", picks the record from the launch's (ncols, transpose_m, n_chains) by regions, and merge_record() names it. The kernel backend's positional merge ABI follows the selected record's column block by default. --- .../cake_kda_det_train_kcp/README.md | 15 +- .../cake_kda_det_train_kcp/cake_kcp.py | 47 ++++- .../cake_kda_det_train_kcp/cake_kcp_jit.py | 161 +++++++++++++++--- .../test_cake_kda_det_train_kcp_build.py | 89 +++++++--- 4 files changed, 254 insertions(+), 58 deletions(-) diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/README.md b/flashinfer/experimental/cake_kda_det_train_kcp/README.md index 6c9d8e9b5f6..6eeb09f431e 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/README.md +++ b/flashinfer/experimental/cake_kda_det_train_kcp/README.md @@ -472,14 +472,15 @@ families of a phase in one launch; off by default). The files of this package: -* `cake_kcp_jit.py` -- the `MODULES` registry of the eleven shipped programs (the +* `cake_kcp_jit.py` -- the `MODULES` registry of the twelve shipped programs (the fragment walks of the forward and backward state recurrences in their `offset` and `tape` forms, the BF16x3 transfer walk (two CTAs per segment and head, each over one 64-column half of the state), the affine - merge at `cols_per_cta` 8 and 16, and the two chunk stages of the + merge in one record per register-blocked build, and the two chunk stages of the context-parallel route, `kcp_fwd_o` (the output from the K1 tapes) and - `kcp_dav` (`dA_qk` and the `dv1` K2 reads); selection by `always` or by - `columns` on the merge's `ncols`, grids by `kcp_segments_heads`, + `kcp_dav` (`dA_qk` and the `dv1` K2 reads); selection by `always`, by + `walk_form`, or by `chains` on the merge's `ncols`, `transpose_m` and chain + count, grids by `kcp_segments_heads`, `kcp_chains`, `kcp_tile_items` or `kcp_pairs`), the logical vocabulary of the KCP arguments, `run_program(name, args, device_index=)`, the lowering of the host backend's walk and merge requests to the registry's launches, @@ -499,18 +500,18 @@ The files of this package: tests; without the file those tests skip. The package reads no environment variable. * `cake_build.py` + `build.json` -- the loader shared by every Cake-exported - package (byte-identical copies) and the manifest of the eleven shipped + package (byte-identical copies) and the manifest of the twelve shipped programs (schema `cake-build/1`, one entry per `_`); `aot/` is never tracked (`.gitignore`). * `csrc/cake_ffi_checks.h` -- the TVM-FFI helper header the bindings include. -* `csrc/cake_kda_det_train__kcp_{kernel,binding}.cu` of the fifteen +* `csrc/cake_kda_det_train__kcp_{kernel,binding}.cu` of the sixteen context-parallel-only programs and `csrc/cake_kda_det_train__{kernel,binding}.cu` of the two chunk stages (`kcp_fwd_o`, `kcp_dav`), rendered like the single-rank ones (digests in `build.json` and `build.test.json`). ## Status -The registry holds eleven shipped programs (`cake_kcp_jit.MODULES`, in +The registry holds twelve shipped programs (`cake_kcp_jit.MODULES`, in `build.json`) and six test-only programs (`cake_kcp_jit.TEST_MODULES`, in `build.test.json`) for `sm_100a`, `sm_103a` and `sm_107a`. diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py index e8c9e98d12d..c0763f96407 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py @@ -842,13 +842,20 @@ def __init__( self, programs: Any | None = None, *, - cols_per_cta: int = 8, + cols_per_cta: int | None = None, path: str = "fi-rendered", ) -> None: self._programs = programs - if cols_per_cta not in (8, 16, 32) or D % cols_per_cta: - raise ValueError(f"cols_per_cta must be 8, 16 or 32, got {cols_per_cta}") - self.cols_per_cta = int(cols_per_cta) + if cols_per_cta is not None and ( + cols_per_cta not in (8, 16, 32, 64) or D % cols_per_cta + ): + raise ValueError( + f"cols_per_cta must be None (the registry's merge record per launch) or 8, 16, 32 or 64, got {cols_per_cta}" + ) + # The K4 column block of the positional merge ABI: None = the column block of the registry record the + # launch's (ncols, n_chains, transpose_m) selects (cake_kcp_jit.merge_record), the same record the + # registry's run_program launches. + self.cols_per_cta = None if cols_per_cta is None else int(cols_per_cta) self.path = path self._dev: dict[tuple[Any, ...], torch.Tensor] = {} self._checked_chains: set[int] = set() @@ -1367,7 +1374,14 @@ def _merge_args( call.merge_log.append( (str(args.get("family")), int(ch.ncols), int(ch.n_chains), cols) ) # D1 evidence: None = registry record - cols = self.cols_per_cta if cols is None else cols + if cols is None: + cols = ( + self.cols_per_cta + if self.cols_per_cta is not None + else _registry_merge_cols( + ch.ncols, ch.n_chains, int(args["transpose_m"]) + ) + ) if ch.ncols % cols: raise ValueError( f"ncols {ch.ncols} is not a multiple of cols_per_cta {cols}" @@ -1401,6 +1415,14 @@ def _merge_args( return abi, ch.n_chains * (ch.ncols // cols) +def _registry_merge_cols(ncols: int, n_chains: int, transpose_m: int) -> int: + """``cols_per_cta`` of the K4 record the registry selects for the launch (:func:`.cake_kcp_jit.merge_record`).""" + from . import cake_kcp_jit # noqa: PLC0415 - H3's registry; optional at import time + + record = cake_kcp_jit.merge_record(ncols, n_chains, transpose_m) + return int(cake_kcp_jit.MODULES[record]["grid_rule"]["cols_per_cta"]) + + def merge_cols_per_cta(ncols: int, n_chains: int, rule: str) -> int | None: """D1 (the column-rule decision): the K4 ``cols_per_cta`` a launch asks for explicitly, ``None`` for the registry's choice of record (``fixed``: the shipped rule, c8 state / c16 map). ``by_chains`` keeps the 8-column @@ -2385,9 +2407,20 @@ def _launch( elif transfer_walk != TRANSFER_WALKS[0]: record += f"_{transfer_walk}" elif program == "affine_merge": - record += ( - "_c8" if int(abi[ABI_PARAMS[program].index("ncols")]) == D else "_c16" + # the registry's record of the launch; the positional ABI carries its column block + params = ABI_PARAMS[program] + ncols = int(abi[params.index("ncols")]) + cols = int(abi[params.index("cols_per_cta")]) + record = jit.merge_record( + ncols, + int(grid) // (ncols // cols), + int(abi[params.index("transpose_m")]), ) + if int(jit.MODULES[record]["grid_rule"]["cols_per_cta"]) != cols: + raise ValueError( + f"the merge ABI carries cols_per_cta {cols}; the registry's record {record} is built for " + f"{jit.MODULES[record]['grid_rule']['cols_per_cta']}" + ) values = dict(zip(ABI_PARAMS[program], abi, strict=True)) jit.run_module( record, values, grid=(int(grid), 1, 1), device_index=device_index diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index 0f21037ecb4..6995e9e2d9f 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -53,9 +53,11 @@ FP32-class offset; * ``kcp_affine_merge`` -- the exact FP32 affine chains over tile pools (``grid_rule`` ``kcp_chains``: ``n_chains * (ncols / cols_per_cta)`` CTAs); - two variants by the chain's column count (``select_rule`` ``columns``: - ``cols_per_cta`` 8 for 128-column state chains, 16 for 256-column map - chains, the faster build of each on GB300); + one record per register-blocked build (column block ``cols_per_cta``, + thread tile, packed FMA pairs on sm_100a / sm_103a), every build giving + the same bits (each output is one FP32 FMA chain in a fixed order); the + ``select_rule`` ``chains`` picks the build per launch from its column + count, operator layout and chain count (:func:`merge_record`); * ``kcp_fwd_o`` -- the chunk output stage of the forward at the rank's shape: ``out`` from the ``h`` / ``v_new`` tapes the forward walk writes and the intra stage's ``qg`` / ``Aqk`` rows, one CTA per (chunk, head) tile of @@ -701,7 +703,7 @@ "select_rule": {"kind": "always"}, "variant": None, }, - "kcp_affine_merge_c8": { + "kcp_affine_merge_c64": { "arches": ["sm_100a", "sm_103a", "sm_107a"], "arg_plan": [ ["buffer", "m"], @@ -723,18 +725,18 @@ ["grid", "grid_z"], ], "ffi_entry": "run", - # One CTA per chain and block of ``cols_per_cta`` columns; the 8-column - # build is the fastest of the three on 128-column state chains (GB300). - "grid_rule": {"cols_per_cta": 8, "kind": "kcp_chains"}, - "kernel_symbol": "kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp", + # The 8 x 8 thread tile over 64 columns (packed FMA pairs on sm_100a / sm_103a): the widest + # launches (provisional rule). + "grid_rule": {"cols_per_cta": 64, "kind": "kcp_chains"}, + "kernel_symbol": "kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp", "launch": { "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 37120, + "dynamic_smem_bytes": 0, "use_pdl": False, }, - "library": "cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp", + "library": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp", "logical": { "c": "c", "chain_start": "chain_start", @@ -752,8 +754,61 @@ "y_final": "y_final", }, "program": "kcp_affine_merge", - "select_rule": {"kind": "columns", "ncols": 128}, - "variant": "c8", + "select_rule": {"kind": "chains", "regions": [[128, None, 48, None]]}, + "variant": "c64", + }, + "kcp_affine_merge_c32": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "m"], + ["buffer", "c"], + ["buffer", "x0"], + ["buffer", "y"], + ["buffer", "chain_start"], + ["buffer", "elem_m"], + ["buffer", "elem_c"], + ["buffer", "elem_y"], + ["buffer", "x0_index"], + ["buffer", "y_final"], + ["parameter", "ncols"], + ["parameter", "transpose_m"], + ["parameter", "write_states"], + ["parameter", "cols_per_cta"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # The 4 x 8 thread tile over 32 columns (packed FMA pairs on sm_100a / sm_103a) (provisional rule). + "grid_rule": {"cols_per_cta": 32, "kind": "kcp_chains"}, + "kernel_symbol": "kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 0, + "use_pdl": False, + }, + "library": "cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp", + "logical": { + "c": "c", + "chain_start": "chain_start", + "cols_per_cta": "cols_per_cta", + "elem_c": "elem_c", + "elem_m": "elem_m", + "elem_y": "elem_y", + "m": "m", + "ncols": "ncols", + "transpose_m": "transpose_m", + "write_states": "write_states", + "x0": "x0", + "x0_index": "x0_index", + "y": "y", + "y_final": "y_final", + }, + "program": "kcp_affine_merge", + "select_rule": {"kind": "chains", "regions": [[128, None, 17, 47]]}, + "variant": "c32", }, "kcp_affine_merge_c16": { "arches": ["sm_100a", "sm_103a", "sm_107a"], @@ -777,18 +832,17 @@ ["grid", "grid_z"], ], "ffi_entry": "run", - # The 16-column build is the fastest of the three on 256-column map - # chains (``X = [B | M]``; GB300). + # The 4 x 4 thread tile over 16 columns: few chains and the map chains (provisional rule). "grid_rule": {"cols_per_cta": 16, "kind": "kcp_chains"}, - "kernel_symbol": "kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp", + "kernel_symbol": "kernel_cake_kda_det_train_15e81414b27b4b831808_kcp", "launch": { "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 41216, + "dynamic_smem_bytes": 0, "use_pdl": False, }, - "library": "cake_kda_det_train_89f5ff59f03eb170e3bf_kcp", + "library": "cake_kda_det_train_15e81414b27b4b831808_kcp", "logical": { "c": "c", "chain_start": "chain_start", @@ -806,7 +860,10 @@ "y_final": "y_final", }, "program": "kcp_affine_merge", - "select_rule": {"kind": "columns", "ncols": 256}, + "select_rule": { + "kind": "chains", + "regions": [[128, None, 1, 16], [256, None, 1, None]], + }, "variant": "c16", }, "kcp_fwd_o": { @@ -1285,8 +1342,27 @@ def records_of_program(program: str) -> tuple[str, ...]: return tuple(name for name, rec in MODULES.items() if rec["program"] == program) +def _in_region( + region: Sequence[Any], ncols: int, n_chains: int, transpose_m: int +) -> bool: + """Whether a ``chains`` region ``[ncols, transpose_m, min_chains, max_chains]`` admits the launch + (``transpose_m`` / ``max_chains`` ``None``: any).""" + width, layout, low, high = region + return ( + ncols == int(width) + and (layout is None or transpose_m == int(layout)) + and n_chains >= int(low) + and (high is None or n_chains <= int(high)) + ) + + def _admits( - rule: Mapping[str, Any], *, ncols: int | None, form: str | None = None + rule: Mapping[str, Any], + *, + ncols: int | None, + form: str | None = None, + n_chains: int | None = None, + transpose_m: int | None = None, ) -> bool: kind = rule["kind"] if kind == "always": @@ -1295,17 +1371,32 @@ def _admits( return ncols is not None and int(ncols) == int(rule["ncols"]) if kind == "walk_form": return form is not None and str(form) == str(rule["form"]) + if kind == "chains": + if ncols is None or n_chains is None or transpose_m is None: + return False + return any( + _in_region(region, int(ncols), int(n_chains), int(transpose_m)) + for region in rule["regions"] + ) raise ValueError(f"unknown Cake KDA KCP selection rule: {kind!r}") def select_record( - program: str, *, ncols: int | None = None, form: str | None = None + program: str, + *, + ncols: int | None = None, + form: str | None = None, + n_chains: int | None = None, + transpose_m: int | None = None, ) -> str: """The one record of ``program`` the launch's scalars select. ``always`` is the rule of a single-record program; ``columns`` (``{"kind": - "columns", "ncols": N}``) picks the merge variant whose column count is - the chain's ``ncols`` (128 for state chains, 256 for map chains); + "columns", "ncols": N}``) picks the variant whose column count is the + chain's ``ncols``; ``chains`` (``{"kind": "chains", "regions": [[ncols, + transpose_m, min_chains, max_chains], ...]}``) picks the merge build whose + regions hold the launch's ``(ncols, transpose_m, n_chains)`` + (``transpose_m`` / ``max_chains`` ``None``: any; :func:`merge_record`); ``walk_form`` (``{"kind": "walk_form", "form": F}``) picks the K1 / K2 walk record of form ``F`` -- ``"full"`` (one CTA per segment and head) or ``"split"`` (the two-CTA record over 64 value columns each), the @@ -1317,15 +1408,34 @@ def select_record( admitted = [ name for name in records_of_program(program) - if _admits(MODULES[name]["select_rule"], ncols=ncols, form=form) + if _admits( + MODULES[name]["select_rule"], + ncols=ncols, + form=form, + n_chains=n_chains, + transpose_m=transpose_m, + ) ] if len(admitted) != 1: raise ValueError( - f"program {program!r}: {len(admitted)} records admit ncols={ncols} form={form} ({admitted}); expected one" + f"program {program!r}: {len(admitted)} records admit ncols={ncols} form={form} n_chains={n_chains} " + f"transpose_m={transpose_m} ({admitted}); expected one" ) return admitted[0] +def merge_record(ncols: int, n_chains: int, transpose_m: int) -> str: + """The ``kcp_affine_merge`` record of a launch of ``n_chains`` chains over ``ncols`` columns with operator + layout ``transpose_m``: the build the records' ``chains`` regions name (every build gives the same bits, the FMA + order of every output being fixed; the rule picks the fastest measured build per launch shape).""" + return select_record( + "kcp_affine_merge", + ncols=int(ncols), + n_chains=int(n_chains), + transpose_m=int(transpose_m), + ) + + def walk_form( program: str, chains: Sequence[int], @@ -1994,7 +2104,7 @@ def _merge(args: Mapping[str, Any], call: Any, device_index: int) -> None: return device = str(call.device) ncols = int(chains.ncols) - record = select_record("kcp_affine_merge", ncols=ncols) + record = merge_record(ncols, n_chains, int(args["transpose_m"])) dummy = _dummy_tile(device) values: dict[str, Any] = { "m": _flat_pool(args["m"], "kcp_affine_merge.m"), @@ -2084,6 +2194,7 @@ def run_program(name: str, args: Mapping[str, Any], *, device_index: int) -> Non "library_name", "load", "load_test", + "merge_record", "records_of_program", "run_module", "run_program", diff --git a/tests/experimental/test_cake_kda_det_train_kcp_build.py b/tests/experimental/test_cake_kda_det_train_kcp_build.py index 5563204cf6e..1be6495c836 100644 --- a/tests/experimental/test_cake_kda_det_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_det_train_kcp_build.py @@ -74,6 +74,35 @@ ["buffer", "do_tma"], ["buffer", "vnew_tma"], ["buffer", "aqk_tma"], ["buffer", "dAqk"], ["buffer", "dv1"], ["parameter", "num_heads"], ["parameter", "scale"], *GRID, ] # fmt: skip +# The K4 merge records (the register-blocked v2 builds, one per CONFIGS entry the select rule names): record -> +# (build keywords, threads per CTA = 128 * cols_per_cta / the thread tile's outputs). +MERGE_V2_MODULE = "tart.kda_chunk_kcp_affine_merge.sm100_sm103_sm107.kda_chunk_det_affine_merge_fp32_v2_kcp" +MERGE_RECORDS = { + "kcp_affine_merge_c64": ( + {"chunk_k": 32, "cols_per_cta": 64, "pairs": True, "ring": 4, "tile": "8x8"}, + 128, + ), + "kcp_affine_merge_c32": ( + {"chunk_k": 32, "cols_per_cta": 32, "pairs": True, "ring": 4, "tile": "4x8"}, + 128, + ), + "kcp_affine_merge_c16": ( + {"chunk_k": 16, "cols_per_cta": 16, "pairs": False, "ring": 4, "tile": "4x4"}, + 128, + ), +} +# The K4 select rule at sample launches: (ncols, transpose_m, n_chains) -> record. +MERGE_SELECT = { + (128, 0, 64): "kcp_affine_merge_c64", + (128, 1, 64): "kcp_affine_merge_c64", + (128, 0, 48): "kcp_affine_merge_c64", + (128, 1, 47): "kcp_affine_merge_c32", + (128, 0, 17): "kcp_affine_merge_c32", + (128, 0, 16): "kcp_affine_merge_c16", + (128, 1, 1): "kcp_affine_merge_c16", + (256, 0, 24): "kcp_affine_merge_c16", + (256, 1, 96): "kcp_affine_merge_c16", +} PLANS = { "kcp_fwdh_offset": FWDH_PLAN, "kcp_fwdh_offset_vh64": FWDH_PLAN, @@ -82,8 +111,7 @@ "kcp_dhu_offset_vh64": DHU_PLAN, "kcp_dhu_tape": DHU_PLAN, "kcp_transfer_walk": TRANSFER_PLAN, - "kcp_affine_merge_c8": MERGE_PLAN, - "kcp_affine_merge_c16": MERGE_PLAN, + **dict.fromkeys(MERGE_RECORDS, MERGE_PLAN), "kcp_fwd_o": FWD_O_PLAN, "kcp_dav": DAV_PLAN, "kcp_transfer_walk_three_mma": TRANSFER_PLAN, @@ -127,14 +155,7 @@ f"tart.kda_chunk_kcp_transfer_walk.{_W}.kda_chunk_kcp_transfer_walk_bf16x3_v1", {"vh": 64}, ), - "kcp_affine_merge_c8": ( - f"tart.kda_chunk_kcp_affine_merge.{_W}.kda_chunk_kcp_affine_merge_fp32_v1", - {"cols_per_cta": 8}, - ), - "kcp_affine_merge_c16": ( - f"tart.kda_chunk_kcp_affine_merge.{_W}.kda_chunk_kcp_affine_merge_fp32_v1", - {"cols_per_cta": 16}, - ), + **{name: (MERGE_V2_MODULE, kwargs) for name, (kwargs, _) in MERGE_RECORDS.items()}, "kcp_fwd_o": (f"tart.kda_chunk_kcp_fwd_o.{_W}.kda_chunk_det_fwd_o_v1_kcp", {}), "kcp_dav": (f"tart.kda_chunk_bwd_dav.{_W}.kda_chunk_train_bwd_dav_v1", {}), "kcp_transfer_walk_three_mma": ( @@ -211,8 +232,15 @@ def test_record_matches_the_frozen_abi(name: str) -> None: ) assert record["launch"]["cluster"] == (1, 1, 1) assert record["launch"]["block"] == ( - (128, 1, 1) if record["program"] == "kcp_affine_merge" else (192, 1, 1) + (MERGE_RECORDS[name][1], 1, 1) + if record["program"] == "kcp_affine_merge" + else (192, 1, 1) ) + if record["program"] == "kcp_affine_merge": + assert record["grid_rule"] == { + "cols_per_cta": MERGE_RECORDS[name][0]["cols_per_cta"], + "kind": "kcp_chains", + } assert isinstance(record["launch"]["dynamic_smem_bytes"], int) expected_name = ( record["program"] @@ -247,13 +275,31 @@ def test_program_names_follow_the_identity_rule() -> None: ) -def test_selection_and_grids() -> None: - assert kcp_jit.select_record("kcp_affine_merge", ncols=128) == "kcp_affine_merge_c8" - assert ( - kcp_jit.select_record("kcp_affine_merge", ncols=256) == "kcp_affine_merge_c16" - ) +def test_merge_selection_covers_every_launch_once() -> None: + for (ncols, transpose_m, n_chains), record in MERGE_SELECT.items(): + assert ( + kcp_jit.select_record( + "kcp_affine_merge", + ncols=ncols, + n_chains=n_chains, + transpose_m=transpose_m, + ) + == record + ), (ncols, transpose_m, n_chains) + assert kcp_jit.merge_record(ncols, n_chains, transpose_m) == record + # every state / map launch shape is admitted by exactly one record (select_record raises otherwise) + for ncols in (128, 256): + for transpose_m in (0, 1): + for n_chains in range(1, 300): + kcp_jit.merge_record(ncols, n_chains, transpose_m) with pytest.raises(ValueError): - kcp_jit.select_record("kcp_affine_merge", ncols=64) + kcp_jit.merge_record(64, 4, 0) + # the chains rule needs the chain count and the layout + with pytest.raises(ValueError): + kcp_jit.select_record("kcp_affine_merge", ncols=128) + + +def test_selection_and_grids() -> None: assert kcp_jit.select_record("kcp_fwdh_offset", form="full") == "kcp_fwdh_offset" assert ( kcp_jit.select_record("kcp_fwdh_offset", form="split") == "kcp_fwdh_offset_vh64" @@ -287,8 +333,13 @@ def test_selection_and_grids() -> None: 1, 1, ) - assert kcp_jit.launch_grid("kcp_affine_merge_c8", n_chains=2, ncols=128) == ( - 32, + assert kcp_jit.launch_grid("kcp_affine_merge_c64", n_chains=2, ncols=128) == ( + 4, + 1, + 1, + ) + assert kcp_jit.launch_grid("kcp_affine_merge_c32", n_chains=3, ncols=256) == ( + 24, 1, 1, ) From 85f72bac6d2961cc4e2d4ea1f2eb8b1528255c91 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 17:15:17 -0800 Subject: [PATCH 02/10] cake_kda_det_train_kcp: render the register-blocked affine merge builds Three merge builds replace the two previous column-block builds: 64 columns with an 8 x 8 thread tile, 32 columns with a 4 x 8 tile (both with packed FMA pairs on sm_100a / sm_103a) and 16 columns with a 4 x 4 tile. build.json gains their entries; the other programs keep their units and entries. --- .../cake_kda_det_train_kcp/build.json | 151 +- .../cake_kda_det_train_kcp/cake_kcp_jit.py | 6 +- ...train_15e81414b27b4b831808_kcp_binding.cu} | 25 +- ...t_train_15e81414b27b4b831808_kcp_kernel.cu | 900 +++++ ..._train_7ddde81b137fb82d27d5_kcp_binding.cu | 156 + ...t_train_7ddde81b137fb82d27d5_kcp_kernel.cu | 3185 +++++++++++++++++ ...t_train_89f5ff59f03eb170e3bf_kcp_kernel.cu | 1312 ------- ...train_c2ad3b63ad48e23ccbdb_kcp_binding.cu} | 29 +- ...t_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu | 1854 ++++++++++ ...t_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu | 1032 ------ 10 files changed, 6241 insertions(+), 2409 deletions(-) rename flashinfer/experimental/cake_kda_det_train_kcp/csrc/{cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu => cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu} (86%) create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu delete mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu rename flashinfer/experimental/cake_kda_det_train_kcp/csrc/{cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu => cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu} (85%) create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu delete mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/build.json b/flashinfer/experimental/cake_kda_det_train_kcp/build.json index 011c2262054..4e2d1bbae06 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/build.json +++ b/flashinfer/experimental/cake_kda_det_train_kcp/build.json @@ -55,11 +55,74 @@ } }, "generator": { - "cake_revision": "2de30eb1576dc83255aacf9d6c6e0294b9e54c31", + "cake_revision": "07bc0a47f570fcbafc700e3082201db8813c3365", "cuda": "13.4.92", - "rendered_at": "2026-10-31T16:27:07Z" + "rendered_at": "2026-11-03T23:32:53Z" }, "libraries": [ + { + "arch": "sm_100a", + "closure_sha256": "1ebfa9e2976961a1ca7742ef0e8e9b8429e5ed298f2b8de61e555c40ac0945f2", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_kda_det_train_15e81414b27b4b831808_kcp_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_15e81414b27b4b831808_kcp_sm_100a.so", + "sha256": { + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu": "a65f44614f617eabf7a6ac5635a00ffb5547d177460e3f66a4a014600d263288", + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu": "29d89e83f021d69047583d2008293e6321209ec0b6d37b86e3663d6f8bbafa6f" + }, + "sources": [ + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu", + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "3e88ddd0f81033a7cca8ea00e0634de10f88d88a41b0744f9d4b43b676a6b02b", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_kda_det_train_15e81414b27b4b831808_kcp_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_15e81414b27b4b831808_kcp_sm_103a.so", + "sha256": { + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu": "a65f44614f617eabf7a6ac5635a00ffb5547d177460e3f66a4a014600d263288", + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu": "29d89e83f021d69047583d2008293e6321209ec0b6d37b86e3663d6f8bbafa6f" + }, + "sources": [ + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu", + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "208de082eb390338fcae73c52912bd052041f0c3d74cd310cae5eda5e624e2ed", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_kda_det_train_15e81414b27b4b831808_kcp_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_15e81414b27b4b831808_kcp_sm_107a.so", + "sha256": { + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu": "a65f44614f617eabf7a6ac5635a00ffb5547d177460e3f66a4a014600d263288", + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu": "29d89e83f021d69047583d2008293e6321209ec0b6d37b86e3663d6f8bbafa6f" + }, + "sources": [ + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu", + "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "a4751e59d5ad550996fe0b1aaaa85ecebb03a5e962a4b5b1a14c2571469df361", @@ -377,65 +440,65 @@ }, { "arch": "sm_100a", - "closure_sha256": "b5ad1a828a49cb1f8af65e8ac1d0b84a9cc91f4f82b6f3b67d9210441652c5d9", + "closure_sha256": "c53a166bc23db11af242f70da51f2b6e7b80d0c54b7f785485bba05d16bc81d8", "entry": "run", "gencode": "-gencode=arch=compute_100a,code=sm_100a", "min_cuda": "12.8", - "name": "cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_sm_100a", + "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_100a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_sm_100a.so", + "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_100a.so", "sha256": { - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu": "5b7ce22f51d32d264506e618a6715c182a9584543c0e689c14e413e135fe3748", - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu": "9973de6e1c2b9ebd2a1d53df082f44c3ad5b819558356106d02c107123a372b2" + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu": "7d580a68978de3029cdcb0b44c43361f37acdc101cd7a074907dc8d5604f0eba", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "7e30683f1adde5744209052239eb51e83b0e2dd6537acd74624d4f0e33f3a5a5" }, "sources": [ - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu", - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu" + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu" ] }, { "arch": "sm_103a", - "closure_sha256": "ae83da3bd0d50a8f6e603bdfc3b0b7fed5253ab9b4a0b6b1bc069c24dcd95b9c", + "closure_sha256": "26e4143bc8f2d06f256f5af75a6fad89f8da44971d505e0143b90f7a1f458116", "entry": "run", "gencode": "-gencode=arch=compute_103a,code=sm_103a", "min_cuda": "12.9", - "name": "cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_sm_103a", + "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_103a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_sm_103a.so", + "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_103a.so", "sha256": { - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu": "5b7ce22f51d32d264506e618a6715c182a9584543c0e689c14e413e135fe3748", - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu": "9973de6e1c2b9ebd2a1d53df082f44c3ad5b819558356106d02c107123a372b2" + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu": "7d580a68978de3029cdcb0b44c43361f37acdc101cd7a074907dc8d5604f0eba", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "7e30683f1adde5744209052239eb51e83b0e2dd6537acd74624d4f0e33f3a5a5" }, "sources": [ - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu", - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu" + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu" ] }, { "arch": "sm_107a", - "closure_sha256": "7d205a74023b9ea2de9039d82782c1bbcd3575857017f75e4cf743131d30af0d", + "closure_sha256": "3f0f09169653ee87e8757c7f767e100e63d18d6e644e558434136ebefed78f9f", "entry": "run", "gencode": "-gencode=arch=compute_107a,code=sm_107a", "min_cuda": "13.4", - "name": "cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_sm_107a", + "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_107a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_sm_107a.so", + "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_107a.so", "sha256": { - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu": "5b7ce22f51d32d264506e618a6715c182a9584543c0e689c14e413e135fe3748", - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu": "9973de6e1c2b9ebd2a1d53df082f44c3ad5b819558356106d02c107123a372b2" + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu": "7d580a68978de3029cdcb0b44c43361f37acdc101cd7a074907dc8d5604f0eba", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "7e30683f1adde5744209052239eb51e83b0e2dd6537acd74624d4f0e33f3a5a5" }, "sources": [ - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu", - "csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu" + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu" ] }, { @@ -566,65 +629,65 @@ }, { "arch": "sm_100a", - "closure_sha256": "2926ce0521fa108e8538ad8e80f7ffee9fde3f2fd20a6ed47dda31ff53c0f824", + "closure_sha256": "1a0c814bd56af8f1e65850e271552044644e0eaef1924c566bce388d2c78f3d6", "entry": "run", "gencode": "-gencode=arch=compute_100a,code=sm_100a", "min_cuda": "12.8", - "name": "cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_sm_100a", + "name": "cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_sm_100a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_sm_100a.so", + "path": "aot/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_sm_100a.so", "sha256": { - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu": "d40485190b8ef6223c56a5351a5387e367ddc23b8f40ebd4ff53aed970e81cf8", - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu": "4b5886fd742b6b1ff11a9a3c7addb368f11bfef85b93c1454b7e1b200302d9c4" + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu": "41a9136194cb90aaf48e56f92016dd33b48a04b6e3856091f7c0b1784978cef2", + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu": "dfb7c0637ebedd91e4d7148e0948bf45412341bcab1e8eac3d993585cc3640a1" }, "sources": [ - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu", - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu" + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu", + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu" ] }, { "arch": "sm_103a", - "closure_sha256": "92472bc6a32f795923b9ccaee651d6c6d463264cce70155ac128fbd0974da720", + "closure_sha256": "f0090b6101df8c4f31f611413822400a7101d1891504bc477edf2676d2b6c627", "entry": "run", "gencode": "-gencode=arch=compute_103a,code=sm_103a", "min_cuda": "12.9", - "name": "cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_sm_103a", + "name": "cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_sm_103a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_sm_103a.so", + "path": "aot/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_sm_103a.so", "sha256": { - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu": "d40485190b8ef6223c56a5351a5387e367ddc23b8f40ebd4ff53aed970e81cf8", - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu": "4b5886fd742b6b1ff11a9a3c7addb368f11bfef85b93c1454b7e1b200302d9c4" + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu": "41a9136194cb90aaf48e56f92016dd33b48a04b6e3856091f7c0b1784978cef2", + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu": "dfb7c0637ebedd91e4d7148e0948bf45412341bcab1e8eac3d993585cc3640a1" }, "sources": [ - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu", - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu" + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu", + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu" ] }, { "arch": "sm_107a", - "closure_sha256": "b08f636354223519a5908f228bc0b036ca9e79c6895580a256295612d47f965e", + "closure_sha256": "db12938a23ea2a2dbe0c45b2713937681f8b1e0540e820e3c7fb45f955ebe98f", "entry": "run", "gencode": "-gencode=arch=compute_107a,code=sm_107a", "min_cuda": "13.4", - "name": "cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_sm_107a", + "name": "cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_sm_107a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_sm_107a.so", + "path": "aot/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_sm_107a.so", "sha256": { - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu": "d40485190b8ef6223c56a5351a5387e367ddc23b8f40ebd4ff53aed970e81cf8", - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu": "4b5886fd742b6b1ff11a9a3c7addb368f11bfef85b93c1454b7e1b200302d9c4" + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu": "41a9136194cb90aaf48e56f92016dd33b48a04b6e3856091f7c0b1784978cef2", + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu": "dfb7c0637ebedd91e4d7148e0948bf45412341bcab1e8eac3d993585cc3640a1" }, "sources": [ - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu", - "csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu" + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu", + "csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu" ] }, { diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index 6995e9e2d9f..2519bccb64d 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -733,7 +733,7 @@ "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 143360, "use_pdl": False, }, "library": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp", @@ -786,7 +786,7 @@ "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 110592, "use_pdl": False, }, "library": "cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp", @@ -839,7 +839,7 @@ "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 61440, "use_pdl": False, }, "library": "cake_kda_det_train_15e81414b27b4b831808_kcp", diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu similarity index 86% rename from flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu rename to flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu index 88c89a01d56..c065412316b 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_binding.cu +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_15e81414b27b4b831808_kcp, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); +extern "C" __global__ void kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); -namespace cake_binding_kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp { +namespace cake_binding_kernel_cake_kda_det_train_15e81414b27b4b831808_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp"; +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_15e81414b27b4b831808_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -117,6 +117,15 @@ void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, Tensor TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_15e81414b27b4b831808_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 61440), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } void* v_m = DataPtr(a_m); void* v_c = DataPtr(a_c); void* v_x0 = DataPtr(a_x0); @@ -135,13 +144,13 @@ void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, Tensor cudaLaunchConfig_t config{}; config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); config.blockDim = dim3(128, 1, 1); - config.dynamicSmemBytes = 41216; + config.dynamicSmemBytes = 61440; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_15e81414b27b4b831808_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp +} // namespace cake_binding_kernel_cake_kda_det_train_15e81414b27b4b831808_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_15e81414b27b4b831808_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu new file mode 100644 index 00000000000..72146109f05 --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu @@ -0,0 +1,900 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LAUNCH_MIN_BLOCKS 3 + +#include + +extern "C" { + +__global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(128) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Prologue + float* xs = reinterpret_cast(smem_raw + 0); + const int xs_addr = smem + 0; + float* ms = reinterpret_cast(smem_raw + 20480); + const int ms_addr = smem + 20480; + + // === Task calls (declaration order) === + int nblk = ncols / 16; + int chain = blockIdx.x / nblk; + int col0 = blockIdx.x % nblk * 16; + int col_tile = col0 / 128; + int v0 = col0 % 128; + int warp_1 = warp; + int g = lane / 4; + int cg = lane % 4; + int e_begin = chain_start[chain]; + int e_end = chain_start[chain + 1]; + unsigned int ring_w = ms_addr + (unsigned int)(warp_1 * 10240); + float acc[16]; + float cn[16]; + float mv[16]; + float xv[4]; + if (transpose_m == 0) { + int row0 = warp_1 * 32 + g * 4; + long long g_off = (long long)col_tile * 16384 + (long long)(row0 * 128 + v0 + cg * 4); + unsigned int x_thr = xs_addr + (unsigned int)((row0 * 20 + cg * 4) * 4); + unsigned int x_rd = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd = ring_w + (unsigned int)(g * 16); + int cp_i = lane / 8; + int cp_j = lane % 8; + int cp_src = cp_i * 128 + warp_1 * 32 + 4 * cp_j; + unsigned int cp_dst = ring_w + (unsigned int)((cp_i * 32 + 4 * cp_j) * 4); + long long m_cur = 0; + if (e_begin < e_end) { + m_cur = elem_m[e_begin]; + long long src0 = m_cur + (long long)cp_src; + unsigned int dst0 = cp_dst + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0), "l"(m + src0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 512), "l"(m + (src0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1024), "l"(m + (src0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1536), "l"(m + (src0 + 1536))); + asm volatile("cp.async.commit_group;"); + long long src0_0 = m_cur + (long long)(cp_src + 2048); + unsigned int dst0_0 = cp_dst + 2560; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0), "l"(m + src0_0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 512), "l"(m + (src0_0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1024), "l"(m + (src0_0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1536), "l"(m + (src0_0 + 1536))); + asm volatile("cp.async.commit_group;"); + long long src0_1 = m_cur + (long long)(cp_src + 4096); + unsigned int dst0_1 = cp_dst + 5120; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1), "l"(m + src0_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 512), "l"(m + (src0_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1024), "l"(m + (src0_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1536), "l"(m + (src0_1 + 1536))); + asm volatile("cp.async.commit_group;"); + long long c_first = -1; + if (col_tile == 0) { + c_first = elem_c[e_begin]; + } + if (c_first >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 128); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 256); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 384); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + long long x0_off = x0_index[chain]; + if (x0_off >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 128); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 256); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 384); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + } else { + int diag = row0 - v0 - cg * 4; + if (col_tile == 0) { + diag = -4096; + } + acc[0] = ((diag == 0) ? 1.0f : 0.0f); + acc[1] = ((diag == 1) ? 1.0f : 0.0f); + acc[2] = ((diag == 2) ? 1.0f : 0.0f); + acc[3] = ((diag == 3) ? 1.0f : 0.0f); + acc[4] = ((diag + 1 == 0) ? 1.0f : 0.0f); + acc[5] = ((diag + 1 == 1) ? 1.0f : 0.0f); + acc[6] = ((diag + 1 == 2) ? 1.0f : 0.0f); + acc[7] = ((diag + 1 == 3) ? 1.0f : 0.0f); + acc[8] = ((diag + 2 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag + 2 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag + 2 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag + 2 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag + 3 == 0) ? 1.0f : 0.0f); + acc[13] = ((diag + 3 == 1) ? 1.0f : 0.0f); + acc[14] = ((diag + 3 == 2) ? 1.0f : 0.0f); + acc[15] = ((diag + 3 == 3) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 80), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 160), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 240), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first = elem_y[e_begin]; + if (y_first >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first + g_off + 384) = _v4; + } + } + } + } else { + long long y_last = y_final[chain]; + if (y_last >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last + g_off + 384) = _v4; + } + } + } + __syncthreads(); + unsigned int buf = 0; + #pragma unroll 1 + for (int e = e_begin; e < e_end; e++) { + int has_next = 0; + long long m_next = 0; + long long c_next = -1; + long long y_pub = -1; + if (e + 1 < e_end) { + has_next = 1; + m_next = elem_m[e + 1]; + if (col_tile == 0) { + c_next = elem_c[e + 1]; + } + if (write_states != 0) { + y_pub = elem_y[e + 1]; + } + } else { + y_pub = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + unsigned int x_rd_e = x_rd + buf * 10240; + #pragma unroll 1 + for (int ch = 0; ch < 8; ch++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc = ch + 3; + if (nc < 8) { + long long src0_2 = m_cur + (long long)(cp_src + nc * 2048); + unsigned int dst0_2 = cp_dst + (unsigned int)(nc % 4 * 2560); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2), "l"(m + src0_2)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 512), "l"(m + (src0_2 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1024), "l"(m + (src0_2 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1536), "l"(m + (src0_2 + 1536))); + } else if (has_next != 0) { + long long src0_3 = m_next + (long long)(cp_src + (nc - 8) * 2048); + unsigned int dst0_3 = cp_dst + (unsigned int)(nc % 4 * 2560); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3), "l"(m + src0_3)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 512), "l"(m + (src0_3 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1024), "l"(m + (src0_3 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1536), "l"(m + (src0_3 + 1536))); + } + asm volatile("cp.async.commit_group;"); + if (ch == 1) { + if (c_next >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 128); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 256); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 384); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + #pragma unroll + for (int kk = 0; kk < 16; kk++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd + (unsigned int)(ch % 4 * 2560) + (unsigned int)(kk * 128))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e + (unsigned int)(ch * 1280) + (unsigned int)(kk * 80))); + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[1], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[1], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[1], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[1], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[2], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[2], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[2], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[2], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[3], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[3], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[3], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[3], xv[3], acc[15]); + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (10240 - buf * 10240)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (10240 - buf * 10240) + 80), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (10240 - buf * 10240) + 160), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (10240 - buf * 10240) + 240), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (y_pub >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub + g_off + 384) = _v4; + } + } + __syncthreads(); + buf = 1 - buf; + m_cur = m_next; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } else { + int row0_1 = warp_1 * 32 + g; + long long g_off_1 = (long long)col_tile * 16384 + (long long)(row0_1 * 128 + v0 + cg * 4); + unsigned int x_thr_1 = xs_addr + (unsigned int)((row0_1 * 20 + cg * 4) * 4); + unsigned int x_rd_1 = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd_1 = ring_w + (unsigned int)(g * 80); + int cp_i_1 = lane / 4; + int cp_j_1 = lane % 4; + int cp_src_1 = cp_i_1 * 128 + warp_1 * 4096 + 4 * cp_j_1; + unsigned int cp_dst_1 = ring_w + (unsigned int)((cp_i_1 * 20 + 4 * cp_j_1) * 4); + long long m_cur_1 = 0; + if (e_begin < e_end) { + m_cur_1 = elem_m[e_begin]; + long long src0_4 = m_cur_1 + (long long)cp_src_1; + unsigned int dst0_4 = cp_dst_1 + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4), "l"(m + src0_4)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 640), "l"(m + (src0_4 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1280), "l"(m + (src0_4 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1920), "l"(m + (src0_4 + 3072))); + asm volatile("cp.async.commit_group;"); + long long src0_2_1 = m_cur_1 + (long long)(cp_src_1 + 16); + unsigned int dst0_2_1 = cp_dst_1 + 2560; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1), "l"(m + src0_2_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 640), "l"(m + (src0_2_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1280), "l"(m + (src0_2_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1920), "l"(m + (src0_2_1 + 3072))); + asm volatile("cp.async.commit_group;"); + long long src0_3_1 = m_cur_1 + (long long)(cp_src_1 + 32); + unsigned int dst0_3_1 = cp_dst_1 + 5120; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1), "l"(m + src0_3_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 640), "l"(m + (src0_3_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1280), "l"(m + (src0_3_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1920), "l"(m + (src0_3_1 + 3072))); + asm volatile("cp.async.commit_group;"); + long long c_first_1 = -1; + if (col_tile == 0) { + c_first_1 = elem_c[e_begin]; + } + if (c_first_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 1024); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2048); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 3072); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + long long x0_off_1 = x0_index[chain]; + if (x0_off_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 1024); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2048); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 3072); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + } else { + int diag_1 = row0_1 - v0 - cg * 4; + if (col_tile == 0) { + diag_1 = -4096; + } + acc[0] = ((diag_1 == 0) ? 1.0f : 0.0f); + acc[1] = ((diag_1 == 1) ? 1.0f : 0.0f); + acc[2] = ((diag_1 == 2) ? 1.0f : 0.0f); + acc[3] = ((diag_1 == 3) ? 1.0f : 0.0f); + acc[4] = ((diag_1 + 8 == 0) ? 1.0f : 0.0f); + acc[5] = ((diag_1 + 8 == 1) ? 1.0f : 0.0f); + acc[6] = ((diag_1 + 8 == 2) ? 1.0f : 0.0f); + acc[7] = ((diag_1 + 8 == 3) ? 1.0f : 0.0f); + acc[8] = ((diag_1 + 16 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag_1 + 16 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag_1 + 16 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag_1 + 16 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag_1 + 24 == 0) ? 1.0f : 0.0f); + acc[13] = ((diag_1 + 24 == 1) ? 1.0f : 0.0f); + acc[14] = ((diag_1 + 24 == 2) ? 1.0f : 0.0f); + acc[15] = ((diag_1 + 24 == 3) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 640), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 1280), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 1920), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first_1 = elem_y[e_begin]; + if (y_first_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 3072) = _v4; + } + } + } + } else { + long long y_last_1 = y_final[chain]; + if (y_last_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 3072) = _v4; + } + } + } + __syncthreads(); + unsigned int buf_1 = 0; + #pragma unroll 1 + for (int e_1 = e_begin; e_1 < e_end; e_1++) { + int has_next_1 = 0; + long long m_next_1 = 0; + long long c_next_1 = -1; + long long y_pub_1 = -1; + if (e_1 + 1 < e_end) { + has_next_1 = 1; + m_next_1 = elem_m[e_1 + 1]; + if (col_tile == 0) { + c_next_1 = elem_c[e_1 + 1]; + } + if (write_states != 0) { + y_pub_1 = elem_y[e_1 + 1]; + } + } else { + y_pub_1 = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + unsigned int x_rd_e_1 = x_rd_1 + buf_1 * 10240; + #pragma unroll 1 + for (int ch_1 = 0; ch_1 < 8; ch_1++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc_1 = ch_1 + 3; + if (nc_1 < 8) { + long long src0_5 = m_cur_1 + (long long)(cp_src_1 + nc_1 * 16); + unsigned int dst0_5 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 2560); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5), "l"(m + src0_5)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 640), "l"(m + (src0_5 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1280), "l"(m + (src0_5 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1920), "l"(m + (src0_5 + 3072))); + } else if (has_next_1 != 0) { + long long src0_6 = m_next_1 + (long long)(cp_src_1 + (nc_1 - 8) * 16); + unsigned int dst0_6 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 2560); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6), "l"(m + src0_6)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 640), "l"(m + (src0_6 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1280), "l"(m + (src0_6 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1920), "l"(m + (src0_6 + 3072))); + } + asm volatile("cp.async.commit_group;"); + if (ch_1 == 1) { + if (c_next_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 1024); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2048); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 3072); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + #pragma unroll + for (int q = 0; q < 4; q++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 2560) + (unsigned int)(q * 16))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[4])), "=r"(*reinterpret_cast(&mv[(4) + 1])), "=r"(*reinterpret_cast(&mv[(4) + 2])), "=r"(*reinterpret_cast(&mv[(4) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 2560) + (unsigned int)(q * 16) + 640)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[8])), "=r"(*reinterpret_cast(&mv[(8) + 1])), "=r"(*reinterpret_cast(&mv[(8) + 2])), "=r"(*reinterpret_cast(&mv[(8) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 2560) + (unsigned int)(q * 16) + 1280)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[12])), "=r"(*reinterpret_cast(&mv[(12) + 1])), "=r"(*reinterpret_cast(&mv[(12) + 2])), "=r"(*reinterpret_cast(&mv[(12) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 2560) + (unsigned int)(q * 16) + 1920)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320))); + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[4], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[4], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[4], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[4], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[8], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[8], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[8], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[8], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[12], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[12], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[12], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[12], xv[3], acc[15]); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320) + 80)); + acc[0] = __fmaf_rn(mv[1], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[1], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[1], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[1], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[5], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[5], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[5], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[5], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[9], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[9], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[9], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[9], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[13], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[13], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[13], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[13], xv[3], acc[15]); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320) + 160)); + acc[0] = __fmaf_rn(mv[2], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[2], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[2], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[2], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[6], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[6], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[6], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[6], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[10], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[10], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[10], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[10], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[14], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[14], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[14], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[14], xv[3], acc[15]); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320) + 240)); + acc[0] = __fmaf_rn(mv[3], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[3], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[3], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[3], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[7], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[7], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[7], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[7], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[11], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[11], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[11], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[11], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[15], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[15], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[15], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[15], xv[3], acc[15]); + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (10240 - buf_1 * 10240)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (10240 - buf_1 * 10240) + 640), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (10240 - buf_1 * 10240) + 1280), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (10240 - buf_1 * 10240) + 1920), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (y_pub_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 3072) = _v4; + } + } + __syncthreads(); + buf_1 = 1 - buf_1; + m_cur_1 = m_next_1; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } +} + +} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu new file mode 100644 index 00000000000..9bbb0e25218 --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu @@ -0,0 +1,156 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp, for the device source that defines it. Do not edit. +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); + +namespace cake_binding_kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) { + TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b + << " has no int64 value in a host check"; +} + +int64_t CMod(int64_t a, int64_t b) { + if (b == 0) NoInt64(a, "%", b); + return b == -1 ? 0 : a % b; +} + +void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, TensorView a_chain_start, TensorView a_elem_m, TensorView a_elem_c, TensorView a_elem_y, TensorView a_x0_index, TensorView a_y_final, int64_t a_ncols, int64_t a_transpose_m, int64_t a_write_states, int64_t a_cols_per_cta, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_m.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_m, "m"); + check_dtype(a_m, DLDataType{kDLFloat, 32, 1}, "m"); + check_cuda_tensor(a_c, "c"); + check_dtype(a_c, DLDataType{kDLFloat, 32, 1}, "c"); + check_cuda_tensor(a_x0, "x0"); + check_dtype(a_x0, DLDataType{kDLFloat, 32, 1}, "x0"); + check_cuda_tensor(a_y, "y"); + check_dtype(a_y, DLDataType{kDLFloat, 32, 1}, "y"); + check_cuda_tensor(a_chain_start, "chain_start"); + check_dtype(a_chain_start, DLDataType{kDLInt, 32, 1}, "chain_start"); + check_cuda_tensor(a_elem_m, "elem_m"); + check_dtype(a_elem_m, DLDataType{kDLInt, 64, 1}, "elem_m"); + check_cuda_tensor(a_elem_c, "elem_c"); + check_dtype(a_elem_c, DLDataType{kDLInt, 64, 1}, "elem_c"); + check_cuda_tensor(a_elem_y, "elem_y"); + check_dtype(a_elem_y, DLDataType{kDLInt, 64, 1}, "elem_y"); + check_cuda_tensor(a_x0_index, "x0_index"); + check_dtype(a_x0_index, DLDataType{kDLInt, 64, 1}, "x0_index"); + check_cuda_tensor(a_y_final, "y_final"); + check_dtype(a_y_final, DLDataType{kDLInt, 64, 1}, "y_final"); + TVM_FFI_CHECK(a_ncols >= -2147483648LL && a_ncols <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'ncols' = " << a_ncols + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_transpose_m >= -2147483648LL && a_transpose_m <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'transpose_m' = " << a_transpose_m + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_write_states >= -2147483648LL && a_write_states <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'write_states' = " << a_write_states + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_cols_per_cta >= -2147483648LL && a_cols_per_cta <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'cols_per_cta' = " << a_cols_per_cta + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_m, "m"); + check_contiguous(a_c, "c"); + check_same_device(a_c, a_m, "c", "m"); + check_contiguous(a_x0, "x0"); + check_same_device(a_x0, a_m, "x0", "m"); + check_contiguous(a_y, "y"); + check_same_device(a_y, a_m, "y", "m"); + check_contiguous(a_chain_start, "chain_start"); + check_same_device(a_chain_start, a_m, "chain_start", "m"); + check_contiguous(a_elem_m, "elem_m"); + check_same_device(a_elem_m, a_m, "elem_m", "m"); + check_contiguous(a_elem_c, "elem_c"); + check_same_device(a_elem_c, a_m, "elem_c", "m"); + check_contiguous(a_elem_y, "elem_y"); + check_same_device(a_elem_y, a_m, "elem_y", "m"); + check_contiguous(a_x0_index, "x0_index"); + check_same_device(a_x0_index, a_m, "x0_index", "m"); + check_contiguous(a_y_final, "y_final"); + check_same_device(a_y_final, a_m, "y_final", "m"); + TVM_FFI_CHECK((a_cols_per_cta == 64LL), ValueError) + << KERNEL_NAME << ": requires (cols_per_cta == 64)"; + TVM_FFI_CHECK((CMod(a_ncols, 128LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(ncols, 128) == 0)"; + TVM_FFI_CHECK((a_ncols >= 128LL), ValueError) + << KERNEL_NAME << ": requires (ncols >= 128)"; + TVM_FFI_CHECK((a_ncols <= 256LL), ValueError) + << KERNEL_NAME << ": requires (ncols <= 256)"; + TVM_FFI_CHECK((a_transpose_m >= 0LL), ValueError) + << KERNEL_NAME << ": requires (transpose_m >= 0)"; + TVM_FFI_CHECK((a_transpose_m <= 1LL), ValueError) + << KERNEL_NAME << ": requires (transpose_m <= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 143360), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + void* v_m = DataPtr(a_m); + void* v_c = DataPtr(a_c); + void* v_x0 = DataPtr(a_x0); + void* v_y = DataPtr(a_y); + void* v_chain_start = DataPtr(a_chain_start); + void* v_elem_m = DataPtr(a_elem_m); + void* v_elem_c = DataPtr(a_elem_c); + void* v_elem_y = DataPtr(a_elem_y); + void* v_x0_index = DataPtr(a_x0_index); + void* v_y_final = DataPtr(a_y_final); + int32_t v_ncols = static_cast(a_ncols); + int32_t v_transpose_m = static_cast(a_transpose_m); + int32_t v_write_states = static_cast(a_write_states); + int32_t v_cols_per_cta = static_cast(a_cols_per_cta); + void* kernel_args[] = {&v_m, &v_c, &v_x0, &v_y, &v_chain_start, &v_elem_m, &v_elem_c, &v_elem_y, &v_x0_index, &v_y_final, &v_ncols, &v_transpose_m, &v_write_states, &v_cols_per_cta}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(128, 1, 1); + config.dynamicSmemBytes = 143360; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu new file mode 100644 index 00000000000..42c6117094a --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu @@ -0,0 +1,3185 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LAUNCH_MIN_BLOCKS 1 + +#include +#if __CUDA_ARCH__ == 1070 +#else + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} +#endif + +extern "C" { + +__global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(128) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Prologue + float* xs = reinterpret_cast(smem_raw + 0); + const int xs_addr = smem + 0; + float* ms = reinterpret_cast(smem_raw + 69632); + const int ms_addr = smem + 69632; + + // === Task calls (declaration order) === + int nblk = ncols / 64; + int chain = blockIdx.x / nblk; + int col0 = blockIdx.x % nblk * 64; + int col_tile = col0 / 128; + int v0 = col0 % 128; + int warp_1 = warp; + int g = lane / 8; + int cg = lane % 8; + int e_begin = chain_start[chain]; + int e_end = chain_start[chain + 1]; + unsigned int ring_w = ms_addr + (unsigned int)(warp_1 * 18432); + float acc[64]; + float cn[64]; + float mv[32]; + float xv[8]; + if (transpose_m == 0) { + int row0 = warp_1 * 32 + g * 4; + long long g_off = (long long)col_tile * 16384 + (long long)(row0 * 128 + v0 + cg * 4); + unsigned int x_thr = xs_addr + (unsigned int)((row0 * 68 + cg * 4) * 4); + unsigned int x_rd = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd = ring_w + (unsigned int)(g * 16); + int cp_i = lane / 8; + int cp_j = lane % 8; + int cp_src = cp_i * 128 + warp_1 * 32 + 4 * cp_j; + unsigned int cp_dst = ring_w + (unsigned int)((cp_i * 32 + 4 * cp_j) * 4); + long long m_cur = 0; + if (e_begin < e_end) { + m_cur = elem_m[e_begin]; + long long src0 = m_cur + (long long)cp_src; + unsigned int dst0 = cp_dst + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0), "l"(m + src0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 512), "l"(m + (src0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1024), "l"(m + (src0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1536), "l"(m + (src0 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 2048), "l"(m + (src0 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 2560), "l"(m + (src0 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 3072), "l"(m + (src0 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 3584), "l"(m + (src0 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_0 = m_cur + (long long)(cp_src + 4096); + unsigned int dst0_0 = cp_dst + 4608; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0), "l"(m + src0_0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 512), "l"(m + (src0_0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1024), "l"(m + (src0_0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1536), "l"(m + (src0_0 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 2048), "l"(m + (src0_0 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 2560), "l"(m + (src0_0 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 3072), "l"(m + (src0_0 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 3584), "l"(m + (src0_0 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_1 = m_cur + (long long)(cp_src + 8192); + unsigned int dst0_1 = cp_dst + 9216; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1), "l"(m + src0_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 512), "l"(m + (src0_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1024), "l"(m + (src0_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1536), "l"(m + (src0_1 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 2048), "l"(m + (src0_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 2560), "l"(m + (src0_1 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 3072), "l"(m + (src0_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 3584), "l"(m + (src0_1 + 3584))); + asm volatile("cp.async.commit_group;"); + long long c_first = -1; + if (col_tile == 0) { + c_first = elem_c[e_begin]; + } + if (c_first >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 32); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 128); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 160); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 256); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 288); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 384); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 416); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2048); + cn[32 + 0] = _v4.x; + cn[32 + 1] = _v4.y; + cn[32 + 2] = _v4.z; + cn[32 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2080); + cn[36 + 0] = _v4.x; + cn[36 + 1] = _v4.y; + cn[36 + 2] = _v4.z; + cn[36 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2176); + cn[40 + 0] = _v4.x; + cn[40 + 1] = _v4.y; + cn[40 + 2] = _v4.z; + cn[40 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2208); + cn[44 + 0] = _v4.x; + cn[44 + 1] = _v4.y; + cn[44 + 2] = _v4.z; + cn[44 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2304); + cn[48 + 0] = _v4.x; + cn[48 + 1] = _v4.y; + cn[48 + 2] = _v4.z; + cn[48 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2336); + cn[52 + 0] = _v4.x; + cn[52 + 1] = _v4.y; + cn[52 + 2] = _v4.z; + cn[52 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2432); + cn[56 + 0] = _v4.x; + cn[56 + 1] = _v4.y; + cn[56 + 2] = _v4.z; + cn[56 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 2464); + cn[60 + 0] = _v4.x; + cn[60 + 1] = _v4.y; + cn[60 + 2] = _v4.z; + cn[60 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + cn[32] = 0.0f; + cn[33] = 0.0f; + cn[34] = 0.0f; + cn[35] = 0.0f; + cn[36] = 0.0f; + cn[37] = 0.0f; + cn[38] = 0.0f; + cn[39] = 0.0f; + cn[40] = 0.0f; + cn[41] = 0.0f; + cn[42] = 0.0f; + cn[43] = 0.0f; + cn[44] = 0.0f; + cn[45] = 0.0f; + cn[46] = 0.0f; + cn[47] = 0.0f; + cn[48] = 0.0f; + cn[49] = 0.0f; + cn[50] = 0.0f; + cn[51] = 0.0f; + cn[52] = 0.0f; + cn[53] = 0.0f; + cn[54] = 0.0f; + cn[55] = 0.0f; + cn[56] = 0.0f; + cn[57] = 0.0f; + cn[58] = 0.0f; + cn[59] = 0.0f; + cn[60] = 0.0f; + cn[61] = 0.0f; + cn[62] = 0.0f; + cn[63] = 0.0f; + } + } + long long x0_off = x0_index[chain]; + if (x0_off >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 32); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 128); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 160); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 256); + acc[16 + 0] = _v4.x; + acc[16 + 1] = _v4.y; + acc[16 + 2] = _v4.z; + acc[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 288); + acc[20 + 0] = _v4.x; + acc[20 + 1] = _v4.y; + acc[20 + 2] = _v4.z; + acc[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 384); + acc[24 + 0] = _v4.x; + acc[24 + 1] = _v4.y; + acc[24 + 2] = _v4.z; + acc[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 416); + acc[28 + 0] = _v4.x; + acc[28 + 1] = _v4.y; + acc[28 + 2] = _v4.z; + acc[28 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2048); + acc[32 + 0] = _v4.x; + acc[32 + 1] = _v4.y; + acc[32 + 2] = _v4.z; + acc[32 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2080); + acc[36 + 0] = _v4.x; + acc[36 + 1] = _v4.y; + acc[36 + 2] = _v4.z; + acc[36 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2176); + acc[40 + 0] = _v4.x; + acc[40 + 1] = _v4.y; + acc[40 + 2] = _v4.z; + acc[40 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2208); + acc[44 + 0] = _v4.x; + acc[44 + 1] = _v4.y; + acc[44 + 2] = _v4.z; + acc[44 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2304); + acc[48 + 0] = _v4.x; + acc[48 + 1] = _v4.y; + acc[48 + 2] = _v4.z; + acc[48 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2336); + acc[52 + 0] = _v4.x; + acc[52 + 1] = _v4.y; + acc[52 + 2] = _v4.z; + acc[52 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2432); + acc[56 + 0] = _v4.x; + acc[56 + 1] = _v4.y; + acc[56 + 2] = _v4.z; + acc[56 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 2464); + acc[60 + 0] = _v4.x; + acc[60 + 1] = _v4.y; + acc[60 + 2] = _v4.z; + acc[60 + 3] = _v4.w; + } + } else { + int diag = row0 - v0 - cg * 4; + if (col_tile == 0) { + diag = -4096; + } + acc[0] = ((diag == 0) ? 1.0f : 0.0f); + acc[1] = ((diag == 1) ? 1.0f : 0.0f); + acc[2] = ((diag == 2) ? 1.0f : 0.0f); + acc[3] = ((diag == 3) ? 1.0f : 0.0f); + acc[4] = ((diag == 32) ? 1.0f : 0.0f); + acc[5] = ((diag == 33) ? 1.0f : 0.0f); + acc[6] = ((diag == 34) ? 1.0f : 0.0f); + acc[7] = ((diag == 35) ? 1.0f : 0.0f); + acc[8] = ((diag + 1 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag + 1 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag + 1 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag + 1 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag + 1 == 32) ? 1.0f : 0.0f); + acc[13] = ((diag + 1 == 33) ? 1.0f : 0.0f); + acc[14] = ((diag + 1 == 34) ? 1.0f : 0.0f); + acc[15] = ((diag + 1 == 35) ? 1.0f : 0.0f); + acc[16] = ((diag + 2 == 0) ? 1.0f : 0.0f); + acc[17] = ((diag + 2 == 1) ? 1.0f : 0.0f); + acc[18] = ((diag + 2 == 2) ? 1.0f : 0.0f); + acc[19] = ((diag + 2 == 3) ? 1.0f : 0.0f); + acc[20] = ((diag + 2 == 32) ? 1.0f : 0.0f); + acc[21] = ((diag + 2 == 33) ? 1.0f : 0.0f); + acc[22] = ((diag + 2 == 34) ? 1.0f : 0.0f); + acc[23] = ((diag + 2 == 35) ? 1.0f : 0.0f); + acc[24] = ((diag + 3 == 0) ? 1.0f : 0.0f); + acc[25] = ((diag + 3 == 1) ? 1.0f : 0.0f); + acc[26] = ((diag + 3 == 2) ? 1.0f : 0.0f); + acc[27] = ((diag + 3 == 3) ? 1.0f : 0.0f); + acc[28] = ((diag + 3 == 32) ? 1.0f : 0.0f); + acc[29] = ((diag + 3 == 33) ? 1.0f : 0.0f); + acc[30] = ((diag + 3 == 34) ? 1.0f : 0.0f); + acc[31] = ((diag + 3 == 35) ? 1.0f : 0.0f); + acc[32] = ((diag + 16 == 0) ? 1.0f : 0.0f); + acc[33] = ((diag + 16 == 1) ? 1.0f : 0.0f); + acc[34] = ((diag + 16 == 2) ? 1.0f : 0.0f); + acc[35] = ((diag + 16 == 3) ? 1.0f : 0.0f); + acc[36] = ((diag + 16 == 32) ? 1.0f : 0.0f); + acc[37] = ((diag + 16 == 33) ? 1.0f : 0.0f); + acc[38] = ((diag + 16 == 34) ? 1.0f : 0.0f); + acc[39] = ((diag + 16 == 35) ? 1.0f : 0.0f); + acc[40] = ((diag + 17 == 0) ? 1.0f : 0.0f); + acc[41] = ((diag + 17 == 1) ? 1.0f : 0.0f); + acc[42] = ((diag + 17 == 2) ? 1.0f : 0.0f); + acc[43] = ((diag + 17 == 3) ? 1.0f : 0.0f); + acc[44] = ((diag + 17 == 32) ? 1.0f : 0.0f); + acc[45] = ((diag + 17 == 33) ? 1.0f : 0.0f); + acc[46] = ((diag + 17 == 34) ? 1.0f : 0.0f); + acc[47] = ((diag + 17 == 35) ? 1.0f : 0.0f); + acc[48] = ((diag + 18 == 0) ? 1.0f : 0.0f); + acc[49] = ((diag + 18 == 1) ? 1.0f : 0.0f); + acc[50] = ((diag + 18 == 2) ? 1.0f : 0.0f); + acc[51] = ((diag + 18 == 3) ? 1.0f : 0.0f); + acc[52] = ((diag + 18 == 32) ? 1.0f : 0.0f); + acc[53] = ((diag + 18 == 33) ? 1.0f : 0.0f); + acc[54] = ((diag + 18 == 34) ? 1.0f : 0.0f); + acc[55] = ((diag + 18 == 35) ? 1.0f : 0.0f); + acc[56] = ((diag + 19 == 0) ? 1.0f : 0.0f); + acc[57] = ((diag + 19 == 1) ? 1.0f : 0.0f); + acc[58] = ((diag + 19 == 2) ? 1.0f : 0.0f); + acc[59] = ((diag + 19 == 3) ? 1.0f : 0.0f); + acc[60] = ((diag + 19 == 32) ? 1.0f : 0.0f); + acc[61] = ((diag + 19 == 33) ? 1.0f : 0.0f); + acc[62] = ((diag + 19 == 34) ? 1.0f : 0.0f); + acc[63] = ((diag + 19 == 35) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 128), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 272), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 400), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 544), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 672), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 816), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 944), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 4352), "r"(__float_as_uint(acc[32])), "r"(__float_as_uint(acc[33])), "r"(__float_as_uint(acc[34])), "r"(__float_as_uint(acc[35])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 4480), "r"(__float_as_uint(acc[36])), "r"(__float_as_uint(acc[37])), "r"(__float_as_uint(acc[38])), "r"(__float_as_uint(acc[39])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 4624), "r"(__float_as_uint(acc[40])), "r"(__float_as_uint(acc[41])), "r"(__float_as_uint(acc[42])), "r"(__float_as_uint(acc[43])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 4752), "r"(__float_as_uint(acc[44])), "r"(__float_as_uint(acc[45])), "r"(__float_as_uint(acc[46])), "r"(__float_as_uint(acc[47])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 4896), "r"(__float_as_uint(acc[48])), "r"(__float_as_uint(acc[49])), "r"(__float_as_uint(acc[50])), "r"(__float_as_uint(acc[51])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 5024), "r"(__float_as_uint(acc[52])), "r"(__float_as_uint(acc[53])), "r"(__float_as_uint(acc[54])), "r"(__float_as_uint(acc[55])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 5168), "r"(__float_as_uint(acc[56])), "r"(__float_as_uint(acc[57])), "r"(__float_as_uint(acc[58])), "r"(__float_as_uint(acc[59])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 5296), "r"(__float_as_uint(acc[60])), "r"(__float_as_uint(acc[61])), "r"(__float_as_uint(acc[62])), "r"(__float_as_uint(acc[63])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first = elem_y[e_begin]; + if (y_first >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first + g_off + 32) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first + g_off + 160) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_first + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_first + g_off + 288) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_first + g_off + 384) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_first + g_off + 416) = _v4; + } + { + float4 _v4 = make_float4(acc[32 + 0], acc[32 + 1], acc[32 + 2], acc[32 + 3]); + *reinterpret_cast(y + y_first + g_off + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[36 + 0], acc[36 + 1], acc[36 + 2], acc[36 + 3]); + *reinterpret_cast(y + y_first + g_off + 2080) = _v4; + } + { + float4 _v4 = make_float4(acc[40 + 0], acc[40 + 1], acc[40 + 2], acc[40 + 3]); + *reinterpret_cast(y + y_first + g_off + 2176) = _v4; + } + { + float4 _v4 = make_float4(acc[44 + 0], acc[44 + 1], acc[44 + 2], acc[44 + 3]); + *reinterpret_cast(y + y_first + g_off + 2208) = _v4; + } + { + float4 _v4 = make_float4(acc[48 + 0], acc[48 + 1], acc[48 + 2], acc[48 + 3]); + *reinterpret_cast(y + y_first + g_off + 2304) = _v4; + } + { + float4 _v4 = make_float4(acc[52 + 0], acc[52 + 1], acc[52 + 2], acc[52 + 3]); + *reinterpret_cast(y + y_first + g_off + 2336) = _v4; + } + { + float4 _v4 = make_float4(acc[56 + 0], acc[56 + 1], acc[56 + 2], acc[56 + 3]); + *reinterpret_cast(y + y_first + g_off + 2432) = _v4; + } + { + float4 _v4 = make_float4(acc[60 + 0], acc[60 + 1], acc[60 + 2], acc[60 + 3]); + *reinterpret_cast(y + y_first + g_off + 2464) = _v4; + } + } + } + } else { + long long y_last = y_final[chain]; + if (y_last >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last + g_off + 32) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last + g_off + 160) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_last + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_last + g_off + 288) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_last + g_off + 384) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_last + g_off + 416) = _v4; + } + { + float4 _v4 = make_float4(acc[32 + 0], acc[32 + 1], acc[32 + 2], acc[32 + 3]); + *reinterpret_cast(y + y_last + g_off + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[36 + 0], acc[36 + 1], acc[36 + 2], acc[36 + 3]); + *reinterpret_cast(y + y_last + g_off + 2080) = _v4; + } + { + float4 _v4 = make_float4(acc[40 + 0], acc[40 + 1], acc[40 + 2], acc[40 + 3]); + *reinterpret_cast(y + y_last + g_off + 2176) = _v4; + } + { + float4 _v4 = make_float4(acc[44 + 0], acc[44 + 1], acc[44 + 2], acc[44 + 3]); + *reinterpret_cast(y + y_last + g_off + 2208) = _v4; + } + { + float4 _v4 = make_float4(acc[48 + 0], acc[48 + 1], acc[48 + 2], acc[48 + 3]); + *reinterpret_cast(y + y_last + g_off + 2304) = _v4; + } + { + float4 _v4 = make_float4(acc[52 + 0], acc[52 + 1], acc[52 + 2], acc[52 + 3]); + *reinterpret_cast(y + y_last + g_off + 2336) = _v4; + } + { + float4 _v4 = make_float4(acc[56 + 0], acc[56 + 1], acc[56 + 2], acc[56 + 3]); + *reinterpret_cast(y + y_last + g_off + 2432) = _v4; + } + { + float4 _v4 = make_float4(acc[60 + 0], acc[60 + 1], acc[60 + 2], acc[60 + 3]); + *reinterpret_cast(y + y_last + g_off + 2464) = _v4; + } + } + } + __syncthreads(); + unsigned int buf = 0; + #pragma unroll 1 + for (int e = e_begin; e < e_end; e++) { + int has_next = 0; + long long m_next = 0; + long long c_next = -1; + long long y_pub = -1; + if (e + 1 < e_end) { + has_next = 1; + m_next = elem_m[e + 1]; + if (col_tile == 0) { + c_next = elem_c[e + 1]; + } + if (write_states != 0) { + y_pub = elem_y[e + 1]; + } + } else { + y_pub = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + acc[16] = cn[16]; + acc[17] = cn[17]; + acc[18] = cn[18]; + acc[19] = cn[19]; + acc[20] = cn[20]; + acc[21] = cn[21]; + acc[22] = cn[22]; + acc[23] = cn[23]; + acc[24] = cn[24]; + acc[25] = cn[25]; + acc[26] = cn[26]; + acc[27] = cn[27]; + acc[28] = cn[28]; + acc[29] = cn[29]; + acc[30] = cn[30]; + acc[31] = cn[31]; + acc[32] = cn[32]; + acc[33] = cn[33]; + acc[34] = cn[34]; + acc[35] = cn[35]; + acc[36] = cn[36]; + acc[37] = cn[37]; + acc[38] = cn[38]; + acc[39] = cn[39]; + acc[40] = cn[40]; + acc[41] = cn[41]; + acc[42] = cn[42]; + acc[43] = cn[43]; + acc[44] = cn[44]; + acc[45] = cn[45]; + acc[46] = cn[46]; + acc[47] = cn[47]; + acc[48] = cn[48]; + acc[49] = cn[49]; + acc[50] = cn[50]; + acc[51] = cn[51]; + acc[52] = cn[52]; + acc[53] = cn[53]; + acc[54] = cn[54]; + acc[55] = cn[55]; + acc[56] = cn[56]; + acc[57] = cn[57]; + acc[58] = cn[58]; + acc[59] = cn[59]; + acc[60] = cn[60]; + acc[61] = cn[61]; + acc[62] = cn[62]; + acc[63] = cn[63]; + unsigned int x_rd_e = x_rd + buf * 34816; + #pragma unroll 1 + for (int ch = 0; ch < 4; ch++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc = ch + 3; + if (nc < 4) { + long long src0_2 = m_cur + (long long)(cp_src + nc * 4096); + unsigned int dst0_2 = cp_dst + (unsigned int)(nc % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2), "l"(m + src0_2)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 512), "l"(m + (src0_2 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1024), "l"(m + (src0_2 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1536), "l"(m + (src0_2 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 2048), "l"(m + (src0_2 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 2560), "l"(m + (src0_2 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 3072), "l"(m + (src0_2 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 3584), "l"(m + (src0_2 + 3584))); + } else if (has_next != 0) { + long long src0_3 = m_next + (long long)(cp_src + (nc - 4) * 4096); + unsigned int dst0_3 = cp_dst + (unsigned int)(nc % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3), "l"(m + src0_3)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 512), "l"(m + (src0_3 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1024), "l"(m + (src0_3 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1536), "l"(m + (src0_3 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 2048), "l"(m + (src0_3 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 2560), "l"(m + (src0_3 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 3072), "l"(m + (src0_3 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 3584), "l"(m + (src0_3 + 3584))); + } + asm volatile("cp.async.commit_group;"); + if (ch == 1) { + if (c_next >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 32); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 128); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 160); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 256); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 288); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 384); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 416); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2048); + cn[32 + 0] = _v4.x; + cn[32 + 1] = _v4.y; + cn[32 + 2] = _v4.z; + cn[32 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2080); + cn[36 + 0] = _v4.x; + cn[36 + 1] = _v4.y; + cn[36 + 2] = _v4.z; + cn[36 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2176); + cn[40 + 0] = _v4.x; + cn[40 + 1] = _v4.y; + cn[40 + 2] = _v4.z; + cn[40 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2208); + cn[44 + 0] = _v4.x; + cn[44 + 1] = _v4.y; + cn[44 + 2] = _v4.z; + cn[44 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2304); + cn[48 + 0] = _v4.x; + cn[48 + 1] = _v4.y; + cn[48 + 2] = _v4.z; + cn[48 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2336); + cn[52 + 0] = _v4.x; + cn[52 + 1] = _v4.y; + cn[52 + 2] = _v4.z; + cn[52 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2432); + cn[56 + 0] = _v4.x; + cn[56 + 1] = _v4.y; + cn[56 + 2] = _v4.z; + cn[56 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 2464); + cn[60 + 0] = _v4.x; + cn[60 + 1] = _v4.y; + cn[60 + 2] = _v4.z; + cn[60 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + cn[32] = 0.0f; + cn[33] = 0.0f; + cn[34] = 0.0f; + cn[35] = 0.0f; + cn[36] = 0.0f; + cn[37] = 0.0f; + cn[38] = 0.0f; + cn[39] = 0.0f; + cn[40] = 0.0f; + cn[41] = 0.0f; + cn[42] = 0.0f; + cn[43] = 0.0f; + cn[44] = 0.0f; + cn[45] = 0.0f; + cn[46] = 0.0f; + cn[47] = 0.0f; + cn[48] = 0.0f; + cn[49] = 0.0f; + cn[50] = 0.0f; + cn[51] = 0.0f; + cn[52] = 0.0f; + cn[53] = 0.0f; + cn[54] = 0.0f; + cn[55] = 0.0f; + cn[56] = 0.0f; + cn[57] = 0.0f; + cn[58] = 0.0f; + cn[59] = 0.0f; + cn[60] = 0.0f; + cn[61] = 0.0f; + cn[62] = 0.0f; + cn[63] = 0.0f; + } + } + #pragma unroll + for (int kk = 0; kk < 32; kk++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd + (unsigned int)(ch % 4 * 4608) + (unsigned int)(kk * 128))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[4])), "=r"(*reinterpret_cast(&mv[(4) + 1])), "=r"(*reinterpret_cast(&mv[(4) + 2])), "=r"(*reinterpret_cast(&mv[(4) + 3])) + : "r"(m_rd + (unsigned int)(ch % 4 * 4608) + (unsigned int)(kk * 128) + 64)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e + (unsigned int)(ch * 8704) + (unsigned int)(kk * 272))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e + (unsigned int)(ch * 8704) + (unsigned int)(kk * 272) + 128)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[0], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[0], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[0], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[0], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[1], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[1], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[1], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[1], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[1], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[1], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[1], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[1], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[2], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[2], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[2], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[2], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[2], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[2], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[2], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[2], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[3], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[3], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[3], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[3], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[3], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[3], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[3], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[3], xv[7], acc[31]); + acc[32] = __fmaf_rn(mv[4], xv[0], acc[32]); + acc[33] = __fmaf_rn(mv[4], xv[1], acc[33]); + acc[34] = __fmaf_rn(mv[4], xv[2], acc[34]); + acc[35] = __fmaf_rn(mv[4], xv[3], acc[35]); + acc[36] = __fmaf_rn(mv[4], xv[4], acc[36]); + acc[37] = __fmaf_rn(mv[4], xv[5], acc[37]); + acc[38] = __fmaf_rn(mv[4], xv[6], acc[38]); + acc[39] = __fmaf_rn(mv[4], xv[7], acc[39]); + acc[40] = __fmaf_rn(mv[5], xv[0], acc[40]); + acc[41] = __fmaf_rn(mv[5], xv[1], acc[41]); + acc[42] = __fmaf_rn(mv[5], xv[2], acc[42]); + acc[43] = __fmaf_rn(mv[5], xv[3], acc[43]); + acc[44] = __fmaf_rn(mv[5], xv[4], acc[44]); + acc[45] = __fmaf_rn(mv[5], xv[5], acc[45]); + acc[46] = __fmaf_rn(mv[5], xv[6], acc[46]); + acc[47] = __fmaf_rn(mv[5], xv[7], acc[47]); + acc[48] = __fmaf_rn(mv[6], xv[0], acc[48]); + acc[49] = __fmaf_rn(mv[6], xv[1], acc[49]); + acc[50] = __fmaf_rn(mv[6], xv[2], acc[50]); + acc[51] = __fmaf_rn(mv[6], xv[3], acc[51]); + acc[52] = __fmaf_rn(mv[6], xv[4], acc[52]); + acc[53] = __fmaf_rn(mv[6], xv[5], acc[53]); + acc[54] = __fmaf_rn(mv[6], xv[6], acc[54]); + acc[55] = __fmaf_rn(mv[6], xv[7], acc[55]); + acc[56] = __fmaf_rn(mv[7], xv[0], acc[56]); + acc[57] = __fmaf_rn(mv[7], xv[1], acc[57]); + acc[58] = __fmaf_rn(mv[7], xv[2], acc[58]); + acc[59] = __fmaf_rn(mv[7], xv[3], acc[59]); + acc[60] = __fmaf_rn(mv[7], xv[4], acc[60]); + acc[61] = __fmaf_rn(mv[7], xv[5], acc[61]); + acc[62] = __fmaf_rn(mv[7], xv[6], acc[62]); + acc[63] = __fmaf_rn(mv[7], xv[7], acc[63]); +#else + float2 m2_0 = make_float2(mv[0], mv[0]); + float2 a2_0 = make_float2(acc[0], acc[1]); + a2_0 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[0], xv[1]), a2_0); + acc[0] = a2_0.x; + acc[1] = a2_0.y; + float2 a2_1 = make_float2(acc[2], acc[3]); + a2_1 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[2], xv[3]), a2_1); + acc[2] = a2_1.x; + acc[3] = a2_1.y; + float2 a2_2 = make_float2(acc[4], acc[5]); + a2_2 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[4], xv[5]), a2_2); + acc[4] = a2_2.x; + acc[5] = a2_2.y; + float2 a2_3 = make_float2(acc[6], acc[7]); + a2_3 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[6], xv[7]), a2_3); + acc[6] = a2_3.x; + acc[7] = a2_3.y; + float2 m2_1 = make_float2(mv[1], mv[1]); + float2 a2_4 = make_float2(acc[8], acc[9]); + a2_4 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[0], xv[1]), a2_4); + acc[8] = a2_4.x; + acc[9] = a2_4.y; + float2 a2_5 = make_float2(acc[10], acc[11]); + a2_5 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[2], xv[3]), a2_5); + acc[10] = a2_5.x; + acc[11] = a2_5.y; + float2 a2_6 = make_float2(acc[12], acc[13]); + a2_6 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[4], xv[5]), a2_6); + acc[12] = a2_6.x; + acc[13] = a2_6.y; + float2 a2_7 = make_float2(acc[14], acc[15]); + a2_7 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[6], xv[7]), a2_7); + acc[14] = a2_7.x; + acc[15] = a2_7.y; + float2 m2_2 = make_float2(mv[2], mv[2]); + float2 a2_8 = make_float2(acc[16], acc[17]); + a2_8 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[0], xv[1]), a2_8); + acc[16] = a2_8.x; + acc[17] = a2_8.y; + float2 a2_9 = make_float2(acc[18], acc[19]); + a2_9 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[2], xv[3]), a2_9); + acc[18] = a2_9.x; + acc[19] = a2_9.y; + float2 a2_10 = make_float2(acc[20], acc[21]); + a2_10 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[4], xv[5]), a2_10); + acc[20] = a2_10.x; + acc[21] = a2_10.y; + float2 a2_11 = make_float2(acc[22], acc[23]); + a2_11 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[6], xv[7]), a2_11); + acc[22] = a2_11.x; + acc[23] = a2_11.y; + float2 m2_3 = make_float2(mv[3], mv[3]); + float2 a2_12 = make_float2(acc[24], acc[25]); + a2_12 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[0], xv[1]), a2_12); + acc[24] = a2_12.x; + acc[25] = a2_12.y; + float2 a2_13 = make_float2(acc[26], acc[27]); + a2_13 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[2], xv[3]), a2_13); + acc[26] = a2_13.x; + acc[27] = a2_13.y; + float2 a2_14 = make_float2(acc[28], acc[29]); + a2_14 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[4], xv[5]), a2_14); + acc[28] = a2_14.x; + acc[29] = a2_14.y; + float2 a2_15 = make_float2(acc[30], acc[31]); + a2_15 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[6], xv[7]), a2_15); + acc[30] = a2_15.x; + acc[31] = a2_15.y; + float2 m2_4 = make_float2(mv[4], mv[4]); + float2 a2_16 = make_float2(acc[32], acc[33]); + a2_16 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[0], xv[1]), a2_16); + acc[32] = a2_16.x; + acc[33] = a2_16.y; + float2 a2_17 = make_float2(acc[34], acc[35]); + a2_17 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[2], xv[3]), a2_17); + acc[34] = a2_17.x; + acc[35] = a2_17.y; + float2 a2_18 = make_float2(acc[36], acc[37]); + a2_18 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[4], xv[5]), a2_18); + acc[36] = a2_18.x; + acc[37] = a2_18.y; + float2 a2_19 = make_float2(acc[38], acc[39]); + a2_19 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[6], xv[7]), a2_19); + acc[38] = a2_19.x; + acc[39] = a2_19.y; + float2 m2_5 = make_float2(mv[5], mv[5]); + float2 a2_20 = make_float2(acc[40], acc[41]); + a2_20 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[0], xv[1]), a2_20); + acc[40] = a2_20.x; + acc[41] = a2_20.y; + float2 a2_21 = make_float2(acc[42], acc[43]); + a2_21 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[2], xv[3]), a2_21); + acc[42] = a2_21.x; + acc[43] = a2_21.y; + float2 a2_22 = make_float2(acc[44], acc[45]); + a2_22 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[4], xv[5]), a2_22); + acc[44] = a2_22.x; + acc[45] = a2_22.y; + float2 a2_23 = make_float2(acc[46], acc[47]); + a2_23 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[6], xv[7]), a2_23); + acc[46] = a2_23.x; + acc[47] = a2_23.y; + float2 m2_6 = make_float2(mv[6], mv[6]); + float2 a2_24 = make_float2(acc[48], acc[49]); + a2_24 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[0], xv[1]), a2_24); + acc[48] = a2_24.x; + acc[49] = a2_24.y; + float2 a2_25 = make_float2(acc[50], acc[51]); + a2_25 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[2], xv[3]), a2_25); + acc[50] = a2_25.x; + acc[51] = a2_25.y; + float2 a2_26 = make_float2(acc[52], acc[53]); + a2_26 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[4], xv[5]), a2_26); + acc[52] = a2_26.x; + acc[53] = a2_26.y; + float2 a2_27 = make_float2(acc[54], acc[55]); + a2_27 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[6], xv[7]), a2_27); + acc[54] = a2_27.x; + acc[55] = a2_27.y; + float2 m2_7 = make_float2(mv[7], mv[7]); + float2 a2_28 = make_float2(acc[56], acc[57]); + a2_28 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[0], xv[1]), a2_28); + acc[56] = a2_28.x; + acc[57] = a2_28.y; + float2 a2_29 = make_float2(acc[58], acc[59]); + a2_29 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[2], xv[3]), a2_29); + acc[58] = a2_29.x; + acc[59] = a2_29.y; + float2 a2_30 = make_float2(acc[60], acc[61]); + a2_30 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[4], xv[5]), a2_30); + acc[60] = a2_30.x; + acc[61] = a2_30.y; + float2 a2_31 = make_float2(acc[62], acc[63]); + a2_31 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[6], xv[7]), a2_31); + acc[62] = a2_31.x; + acc[63] = a2_31.y; +#endif + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 128), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 272), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 400), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 544), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 672), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 816), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 944), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 4352), "r"(__float_as_uint(acc[32])), "r"(__float_as_uint(acc[33])), "r"(__float_as_uint(acc[34])), "r"(__float_as_uint(acc[35])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 4480), "r"(__float_as_uint(acc[36])), "r"(__float_as_uint(acc[37])), "r"(__float_as_uint(acc[38])), "r"(__float_as_uint(acc[39])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 4624), "r"(__float_as_uint(acc[40])), "r"(__float_as_uint(acc[41])), "r"(__float_as_uint(acc[42])), "r"(__float_as_uint(acc[43])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 4752), "r"(__float_as_uint(acc[44])), "r"(__float_as_uint(acc[45])), "r"(__float_as_uint(acc[46])), "r"(__float_as_uint(acc[47])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 4896), "r"(__float_as_uint(acc[48])), "r"(__float_as_uint(acc[49])), "r"(__float_as_uint(acc[50])), "r"(__float_as_uint(acc[51])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 5024), "r"(__float_as_uint(acc[52])), "r"(__float_as_uint(acc[53])), "r"(__float_as_uint(acc[54])), "r"(__float_as_uint(acc[55])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 5168), "r"(__float_as_uint(acc[56])), "r"(__float_as_uint(acc[57])), "r"(__float_as_uint(acc[58])), "r"(__float_as_uint(acc[59])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816) + 5296), "r"(__float_as_uint(acc[60])), "r"(__float_as_uint(acc[61])), "r"(__float_as_uint(acc[62])), "r"(__float_as_uint(acc[63])) : "memory"); + if (y_pub >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub + g_off + 32) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub + g_off + 160) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_pub + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_pub + g_off + 288) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_pub + g_off + 384) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_pub + g_off + 416) = _v4; + } + { + float4 _v4 = make_float4(acc[32 + 0], acc[32 + 1], acc[32 + 2], acc[32 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[36 + 0], acc[36 + 1], acc[36 + 2], acc[36 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2080) = _v4; + } + { + float4 _v4 = make_float4(acc[40 + 0], acc[40 + 1], acc[40 + 2], acc[40 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2176) = _v4; + } + { + float4 _v4 = make_float4(acc[44 + 0], acc[44 + 1], acc[44 + 2], acc[44 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2208) = _v4; + } + { + float4 _v4 = make_float4(acc[48 + 0], acc[48 + 1], acc[48 + 2], acc[48 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2304) = _v4; + } + { + float4 _v4 = make_float4(acc[52 + 0], acc[52 + 1], acc[52 + 2], acc[52 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2336) = _v4; + } + { + float4 _v4 = make_float4(acc[56 + 0], acc[56 + 1], acc[56 + 2], acc[56 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2432) = _v4; + } + { + float4 _v4 = make_float4(acc[60 + 0], acc[60 + 1], acc[60 + 2], acc[60 + 3]); + *reinterpret_cast(y + y_pub + g_off + 2464) = _v4; + } + } + __syncthreads(); + buf = 1 - buf; + m_cur = m_next; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } else { + int row0_1 = warp_1 * 32 + g; + long long g_off_1 = (long long)col_tile * 16384 + (long long)(row0_1 * 128 + v0 + cg * 4); + unsigned int x_thr_1 = xs_addr + (unsigned int)((row0_1 * 68 + cg * 4) * 4); + unsigned int x_rd_1 = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd_1 = ring_w + (unsigned int)(g * 144); + int cp_i_1 = lane / 8; + int cp_j_1 = lane % 8; + int cp_src_1 = cp_i_1 * 128 + warp_1 * 4096 + 4 * cp_j_1; + unsigned int cp_dst_1 = ring_w + (unsigned int)((cp_i_1 * 36 + 4 * cp_j_1) * 4); + long long m_cur_1 = 0; + if (e_begin < e_end) { + m_cur_1 = elem_m[e_begin]; + long long src0_4 = m_cur_1 + (long long)cp_src_1; + unsigned int dst0_4 = cp_dst_1 + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4), "l"(m + src0_4)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 576), "l"(m + (src0_4 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1152), "l"(m + (src0_4 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1728), "l"(m + (src0_4 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 2304), "l"(m + (src0_4 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 2880), "l"(m + (src0_4 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 3456), "l"(m + (src0_4 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 4032), "l"(m + (src0_4 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_2_1 = m_cur_1 + (long long)(cp_src_1 + 32); + unsigned int dst0_2_1 = cp_dst_1 + 4608; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1), "l"(m + src0_2_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 576), "l"(m + (src0_2_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1152), "l"(m + (src0_2_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1728), "l"(m + (src0_2_1 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 2304), "l"(m + (src0_2_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 2880), "l"(m + (src0_2_1 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 3456), "l"(m + (src0_2_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 4032), "l"(m + (src0_2_1 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_3_1 = m_cur_1 + (long long)(cp_src_1 + 64); + unsigned int dst0_3_1 = cp_dst_1 + 9216; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1), "l"(m + src0_3_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 576), "l"(m + (src0_3_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1152), "l"(m + (src0_3_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1728), "l"(m + (src0_3_1 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 2304), "l"(m + (src0_3_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 2880), "l"(m + (src0_3_1 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 3456), "l"(m + (src0_3_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 4032), "l"(m + (src0_3_1 + 3584))); + asm volatile("cp.async.commit_group;"); + long long c_first_1 = -1; + if (col_tile == 0) { + c_first_1 = elem_c[e_begin]; + } + if (c_first_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 32); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 512); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 544); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 1024); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 1056); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 1536); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 1568); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2048); + cn[32 + 0] = _v4.x; + cn[32 + 1] = _v4.y; + cn[32 + 2] = _v4.z; + cn[32 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2080); + cn[36 + 0] = _v4.x; + cn[36 + 1] = _v4.y; + cn[36 + 2] = _v4.z; + cn[36 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2560); + cn[40 + 0] = _v4.x; + cn[40 + 1] = _v4.y; + cn[40 + 2] = _v4.z; + cn[40 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2592); + cn[44 + 0] = _v4.x; + cn[44 + 1] = _v4.y; + cn[44 + 2] = _v4.z; + cn[44 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 3072); + cn[48 + 0] = _v4.x; + cn[48 + 1] = _v4.y; + cn[48 + 2] = _v4.z; + cn[48 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 3104); + cn[52 + 0] = _v4.x; + cn[52 + 1] = _v4.y; + cn[52 + 2] = _v4.z; + cn[52 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 3584); + cn[56 + 0] = _v4.x; + cn[56 + 1] = _v4.y; + cn[56 + 2] = _v4.z; + cn[56 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 3616); + cn[60 + 0] = _v4.x; + cn[60 + 1] = _v4.y; + cn[60 + 2] = _v4.z; + cn[60 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + cn[32] = 0.0f; + cn[33] = 0.0f; + cn[34] = 0.0f; + cn[35] = 0.0f; + cn[36] = 0.0f; + cn[37] = 0.0f; + cn[38] = 0.0f; + cn[39] = 0.0f; + cn[40] = 0.0f; + cn[41] = 0.0f; + cn[42] = 0.0f; + cn[43] = 0.0f; + cn[44] = 0.0f; + cn[45] = 0.0f; + cn[46] = 0.0f; + cn[47] = 0.0f; + cn[48] = 0.0f; + cn[49] = 0.0f; + cn[50] = 0.0f; + cn[51] = 0.0f; + cn[52] = 0.0f; + cn[53] = 0.0f; + cn[54] = 0.0f; + cn[55] = 0.0f; + cn[56] = 0.0f; + cn[57] = 0.0f; + cn[58] = 0.0f; + cn[59] = 0.0f; + cn[60] = 0.0f; + cn[61] = 0.0f; + cn[62] = 0.0f; + cn[63] = 0.0f; + } + } + long long x0_off_1 = x0_index[chain]; + if (x0_off_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 32); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 512); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 544); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 1024); + acc[16 + 0] = _v4.x; + acc[16 + 1] = _v4.y; + acc[16 + 2] = _v4.z; + acc[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 1056); + acc[20 + 0] = _v4.x; + acc[20 + 1] = _v4.y; + acc[20 + 2] = _v4.z; + acc[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 1536); + acc[24 + 0] = _v4.x; + acc[24 + 1] = _v4.y; + acc[24 + 2] = _v4.z; + acc[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 1568); + acc[28 + 0] = _v4.x; + acc[28 + 1] = _v4.y; + acc[28 + 2] = _v4.z; + acc[28 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2048); + acc[32 + 0] = _v4.x; + acc[32 + 1] = _v4.y; + acc[32 + 2] = _v4.z; + acc[32 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2080); + acc[36 + 0] = _v4.x; + acc[36 + 1] = _v4.y; + acc[36 + 2] = _v4.z; + acc[36 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2560); + acc[40 + 0] = _v4.x; + acc[40 + 1] = _v4.y; + acc[40 + 2] = _v4.z; + acc[40 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2592); + acc[44 + 0] = _v4.x; + acc[44 + 1] = _v4.y; + acc[44 + 2] = _v4.z; + acc[44 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 3072); + acc[48 + 0] = _v4.x; + acc[48 + 1] = _v4.y; + acc[48 + 2] = _v4.z; + acc[48 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 3104); + acc[52 + 0] = _v4.x; + acc[52 + 1] = _v4.y; + acc[52 + 2] = _v4.z; + acc[52 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 3584); + acc[56 + 0] = _v4.x; + acc[56 + 1] = _v4.y; + acc[56 + 2] = _v4.z; + acc[56 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 3616); + acc[60 + 0] = _v4.x; + acc[60 + 1] = _v4.y; + acc[60 + 2] = _v4.z; + acc[60 + 3] = _v4.w; + } + } else { + int diag_1 = row0_1 - v0 - cg * 4; + if (col_tile == 0) { + diag_1 = -4096; + } + acc[0] = ((diag_1 == 0) ? 1.0f : 0.0f); + acc[1] = ((diag_1 == 1) ? 1.0f : 0.0f); + acc[2] = ((diag_1 == 2) ? 1.0f : 0.0f); + acc[3] = ((diag_1 == 3) ? 1.0f : 0.0f); + acc[4] = ((diag_1 == 32) ? 1.0f : 0.0f); + acc[5] = ((diag_1 == 33) ? 1.0f : 0.0f); + acc[6] = ((diag_1 == 34) ? 1.0f : 0.0f); + acc[7] = ((diag_1 == 35) ? 1.0f : 0.0f); + acc[8] = ((diag_1 + 4 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag_1 + 4 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag_1 + 4 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag_1 + 4 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag_1 + 4 == 32) ? 1.0f : 0.0f); + acc[13] = ((diag_1 + 4 == 33) ? 1.0f : 0.0f); + acc[14] = ((diag_1 + 4 == 34) ? 1.0f : 0.0f); + acc[15] = ((diag_1 + 4 == 35) ? 1.0f : 0.0f); + acc[16] = ((diag_1 + 8 == 0) ? 1.0f : 0.0f); + acc[17] = ((diag_1 + 8 == 1) ? 1.0f : 0.0f); + acc[18] = ((diag_1 + 8 == 2) ? 1.0f : 0.0f); + acc[19] = ((diag_1 + 8 == 3) ? 1.0f : 0.0f); + acc[20] = ((diag_1 + 8 == 32) ? 1.0f : 0.0f); + acc[21] = ((diag_1 + 8 == 33) ? 1.0f : 0.0f); + acc[22] = ((diag_1 + 8 == 34) ? 1.0f : 0.0f); + acc[23] = ((diag_1 + 8 == 35) ? 1.0f : 0.0f); + acc[24] = ((diag_1 + 12 == 0) ? 1.0f : 0.0f); + acc[25] = ((diag_1 + 12 == 1) ? 1.0f : 0.0f); + acc[26] = ((diag_1 + 12 == 2) ? 1.0f : 0.0f); + acc[27] = ((diag_1 + 12 == 3) ? 1.0f : 0.0f); + acc[28] = ((diag_1 + 12 == 32) ? 1.0f : 0.0f); + acc[29] = ((diag_1 + 12 == 33) ? 1.0f : 0.0f); + acc[30] = ((diag_1 + 12 == 34) ? 1.0f : 0.0f); + acc[31] = ((diag_1 + 12 == 35) ? 1.0f : 0.0f); + acc[32] = ((diag_1 + 16 == 0) ? 1.0f : 0.0f); + acc[33] = ((diag_1 + 16 == 1) ? 1.0f : 0.0f); + acc[34] = ((diag_1 + 16 == 2) ? 1.0f : 0.0f); + acc[35] = ((diag_1 + 16 == 3) ? 1.0f : 0.0f); + acc[36] = ((diag_1 + 16 == 32) ? 1.0f : 0.0f); + acc[37] = ((diag_1 + 16 == 33) ? 1.0f : 0.0f); + acc[38] = ((diag_1 + 16 == 34) ? 1.0f : 0.0f); + acc[39] = ((diag_1 + 16 == 35) ? 1.0f : 0.0f); + acc[40] = ((diag_1 + 20 == 0) ? 1.0f : 0.0f); + acc[41] = ((diag_1 + 20 == 1) ? 1.0f : 0.0f); + acc[42] = ((diag_1 + 20 == 2) ? 1.0f : 0.0f); + acc[43] = ((diag_1 + 20 == 3) ? 1.0f : 0.0f); + acc[44] = ((diag_1 + 20 == 32) ? 1.0f : 0.0f); + acc[45] = ((diag_1 + 20 == 33) ? 1.0f : 0.0f); + acc[46] = ((diag_1 + 20 == 34) ? 1.0f : 0.0f); + acc[47] = ((diag_1 + 20 == 35) ? 1.0f : 0.0f); + acc[48] = ((diag_1 + 24 == 0) ? 1.0f : 0.0f); + acc[49] = ((diag_1 + 24 == 1) ? 1.0f : 0.0f); + acc[50] = ((diag_1 + 24 == 2) ? 1.0f : 0.0f); + acc[51] = ((diag_1 + 24 == 3) ? 1.0f : 0.0f); + acc[52] = ((diag_1 + 24 == 32) ? 1.0f : 0.0f); + acc[53] = ((diag_1 + 24 == 33) ? 1.0f : 0.0f); + acc[54] = ((diag_1 + 24 == 34) ? 1.0f : 0.0f); + acc[55] = ((diag_1 + 24 == 35) ? 1.0f : 0.0f); + acc[56] = ((diag_1 + 28 == 0) ? 1.0f : 0.0f); + acc[57] = ((diag_1 + 28 == 1) ? 1.0f : 0.0f); + acc[58] = ((diag_1 + 28 == 2) ? 1.0f : 0.0f); + acc[59] = ((diag_1 + 28 == 3) ? 1.0f : 0.0f); + acc[60] = ((diag_1 + 28 == 32) ? 1.0f : 0.0f); + acc[61] = ((diag_1 + 28 == 33) ? 1.0f : 0.0f); + acc[62] = ((diag_1 + 28 == 34) ? 1.0f : 0.0f); + acc[63] = ((diag_1 + 28 == 35) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 128), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 1088), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 1216), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 2176), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 2304), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 3264), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 3392), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 4352), "r"(__float_as_uint(acc[32])), "r"(__float_as_uint(acc[33])), "r"(__float_as_uint(acc[34])), "r"(__float_as_uint(acc[35])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 4480), "r"(__float_as_uint(acc[36])), "r"(__float_as_uint(acc[37])), "r"(__float_as_uint(acc[38])), "r"(__float_as_uint(acc[39])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 5440), "r"(__float_as_uint(acc[40])), "r"(__float_as_uint(acc[41])), "r"(__float_as_uint(acc[42])), "r"(__float_as_uint(acc[43])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 5568), "r"(__float_as_uint(acc[44])), "r"(__float_as_uint(acc[45])), "r"(__float_as_uint(acc[46])), "r"(__float_as_uint(acc[47])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 6528), "r"(__float_as_uint(acc[48])), "r"(__float_as_uint(acc[49])), "r"(__float_as_uint(acc[50])), "r"(__float_as_uint(acc[51])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 6656), "r"(__float_as_uint(acc[52])), "r"(__float_as_uint(acc[53])), "r"(__float_as_uint(acc[54])), "r"(__float_as_uint(acc[55])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 7616), "r"(__float_as_uint(acc[56])), "r"(__float_as_uint(acc[57])), "r"(__float_as_uint(acc[58])), "r"(__float_as_uint(acc[59])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 7744), "r"(__float_as_uint(acc[60])), "r"(__float_as_uint(acc[61])), "r"(__float_as_uint(acc[62])), "r"(__float_as_uint(acc[63])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first_1 = elem_y[e_begin]; + if (y_first_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 32) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 512) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 544) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 1056) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 1536) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 1568) = _v4; + } + { + float4 _v4 = make_float4(acc[32 + 0], acc[32 + 1], acc[32 + 2], acc[32 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[36 + 0], acc[36 + 1], acc[36 + 2], acc[36 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2080) = _v4; + } + { + float4 _v4 = make_float4(acc[40 + 0], acc[40 + 1], acc[40 + 2], acc[40 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2560) = _v4; + } + { + float4 _v4 = make_float4(acc[44 + 0], acc[44 + 1], acc[44 + 2], acc[44 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2592) = _v4; + } + { + float4 _v4 = make_float4(acc[48 + 0], acc[48 + 1], acc[48 + 2], acc[48 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 3072) = _v4; + } + { + float4 _v4 = make_float4(acc[52 + 0], acc[52 + 1], acc[52 + 2], acc[52 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 3104) = _v4; + } + { + float4 _v4 = make_float4(acc[56 + 0], acc[56 + 1], acc[56 + 2], acc[56 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 3584) = _v4; + } + { + float4 _v4 = make_float4(acc[60 + 0], acc[60 + 1], acc[60 + 2], acc[60 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 3616) = _v4; + } + } + } + } else { + long long y_last_1 = y_final[chain]; + if (y_last_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 32) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 512) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 544) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 1056) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 1536) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 1568) = _v4; + } + { + float4 _v4 = make_float4(acc[32 + 0], acc[32 + 1], acc[32 + 2], acc[32 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[36 + 0], acc[36 + 1], acc[36 + 2], acc[36 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2080) = _v4; + } + { + float4 _v4 = make_float4(acc[40 + 0], acc[40 + 1], acc[40 + 2], acc[40 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2560) = _v4; + } + { + float4 _v4 = make_float4(acc[44 + 0], acc[44 + 1], acc[44 + 2], acc[44 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2592) = _v4; + } + { + float4 _v4 = make_float4(acc[48 + 0], acc[48 + 1], acc[48 + 2], acc[48 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 3072) = _v4; + } + { + float4 _v4 = make_float4(acc[52 + 0], acc[52 + 1], acc[52 + 2], acc[52 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 3104) = _v4; + } + { + float4 _v4 = make_float4(acc[56 + 0], acc[56 + 1], acc[56 + 2], acc[56 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 3584) = _v4; + } + { + float4 _v4 = make_float4(acc[60 + 0], acc[60 + 1], acc[60 + 2], acc[60 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 3616) = _v4; + } + } + } + __syncthreads(); + unsigned int buf_1 = 0; + #pragma unroll 1 + for (int e_1 = e_begin; e_1 < e_end; e_1++) { + int has_next_1 = 0; + long long m_next_1 = 0; + long long c_next_1 = -1; + long long y_pub_1 = -1; + if (e_1 + 1 < e_end) { + has_next_1 = 1; + m_next_1 = elem_m[e_1 + 1]; + if (col_tile == 0) { + c_next_1 = elem_c[e_1 + 1]; + } + if (write_states != 0) { + y_pub_1 = elem_y[e_1 + 1]; + } + } else { + y_pub_1 = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + acc[16] = cn[16]; + acc[17] = cn[17]; + acc[18] = cn[18]; + acc[19] = cn[19]; + acc[20] = cn[20]; + acc[21] = cn[21]; + acc[22] = cn[22]; + acc[23] = cn[23]; + acc[24] = cn[24]; + acc[25] = cn[25]; + acc[26] = cn[26]; + acc[27] = cn[27]; + acc[28] = cn[28]; + acc[29] = cn[29]; + acc[30] = cn[30]; + acc[31] = cn[31]; + acc[32] = cn[32]; + acc[33] = cn[33]; + acc[34] = cn[34]; + acc[35] = cn[35]; + acc[36] = cn[36]; + acc[37] = cn[37]; + acc[38] = cn[38]; + acc[39] = cn[39]; + acc[40] = cn[40]; + acc[41] = cn[41]; + acc[42] = cn[42]; + acc[43] = cn[43]; + acc[44] = cn[44]; + acc[45] = cn[45]; + acc[46] = cn[46]; + acc[47] = cn[47]; + acc[48] = cn[48]; + acc[49] = cn[49]; + acc[50] = cn[50]; + acc[51] = cn[51]; + acc[52] = cn[52]; + acc[53] = cn[53]; + acc[54] = cn[54]; + acc[55] = cn[55]; + acc[56] = cn[56]; + acc[57] = cn[57]; + acc[58] = cn[58]; + acc[59] = cn[59]; + acc[60] = cn[60]; + acc[61] = cn[61]; + acc[62] = cn[62]; + acc[63] = cn[63]; + unsigned int x_rd_e_1 = x_rd_1 + buf_1 * 34816; + #pragma unroll 1 + for (int ch_1 = 0; ch_1 < 4; ch_1++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc_1 = ch_1 + 3; + if (nc_1 < 4) { + long long src0_5 = m_cur_1 + (long long)(cp_src_1 + nc_1 * 32); + unsigned int dst0_5 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5), "l"(m + src0_5)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 576), "l"(m + (src0_5 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1152), "l"(m + (src0_5 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1728), "l"(m + (src0_5 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 2304), "l"(m + (src0_5 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 2880), "l"(m + (src0_5 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 3456), "l"(m + (src0_5 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 4032), "l"(m + (src0_5 + 3584))); + } else if (has_next_1 != 0) { + long long src0_6 = m_next_1 + (long long)(cp_src_1 + (nc_1 - 4) * 32); + unsigned int dst0_6 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6), "l"(m + src0_6)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 576), "l"(m + (src0_6 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1152), "l"(m + (src0_6 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1728), "l"(m + (src0_6 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 2304), "l"(m + (src0_6 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 2880), "l"(m + (src0_6 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 3456), "l"(m + (src0_6 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 4032), "l"(m + (src0_6 + 3584))); + } + asm volatile("cp.async.commit_group;"); + if (ch_1 == 1) { + if (c_next_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 32); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 512); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 544); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 1024); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 1056); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 1536); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 1568); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2048); + cn[32 + 0] = _v4.x; + cn[32 + 1] = _v4.y; + cn[32 + 2] = _v4.z; + cn[32 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2080); + cn[36 + 0] = _v4.x; + cn[36 + 1] = _v4.y; + cn[36 + 2] = _v4.z; + cn[36 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2560); + cn[40 + 0] = _v4.x; + cn[40 + 1] = _v4.y; + cn[40 + 2] = _v4.z; + cn[40 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2592); + cn[44 + 0] = _v4.x; + cn[44 + 1] = _v4.y; + cn[44 + 2] = _v4.z; + cn[44 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 3072); + cn[48 + 0] = _v4.x; + cn[48 + 1] = _v4.y; + cn[48 + 2] = _v4.z; + cn[48 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 3104); + cn[52 + 0] = _v4.x; + cn[52 + 1] = _v4.y; + cn[52 + 2] = _v4.z; + cn[52 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 3584); + cn[56 + 0] = _v4.x; + cn[56 + 1] = _v4.y; + cn[56 + 2] = _v4.z; + cn[56 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 3616); + cn[60 + 0] = _v4.x; + cn[60 + 1] = _v4.y; + cn[60 + 2] = _v4.z; + cn[60 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + cn[32] = 0.0f; + cn[33] = 0.0f; + cn[34] = 0.0f; + cn[35] = 0.0f; + cn[36] = 0.0f; + cn[37] = 0.0f; + cn[38] = 0.0f; + cn[39] = 0.0f; + cn[40] = 0.0f; + cn[41] = 0.0f; + cn[42] = 0.0f; + cn[43] = 0.0f; + cn[44] = 0.0f; + cn[45] = 0.0f; + cn[46] = 0.0f; + cn[47] = 0.0f; + cn[48] = 0.0f; + cn[49] = 0.0f; + cn[50] = 0.0f; + cn[51] = 0.0f; + cn[52] = 0.0f; + cn[53] = 0.0f; + cn[54] = 0.0f; + cn[55] = 0.0f; + cn[56] = 0.0f; + cn[57] = 0.0f; + cn[58] = 0.0f; + cn[59] = 0.0f; + cn[60] = 0.0f; + cn[61] = 0.0f; + cn[62] = 0.0f; + cn[63] = 0.0f; + } + } + #pragma unroll + for (int q = 0; q < 8; q++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[4])), "=r"(*reinterpret_cast(&mv[(4) + 1])), "=r"(*reinterpret_cast(&mv[(4) + 2])), "=r"(*reinterpret_cast(&mv[(4) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 576)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[8])), "=r"(*reinterpret_cast(&mv[(8) + 1])), "=r"(*reinterpret_cast(&mv[(8) + 2])), "=r"(*reinterpret_cast(&mv[(8) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 1152)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[12])), "=r"(*reinterpret_cast(&mv[(12) + 1])), "=r"(*reinterpret_cast(&mv[(12) + 2])), "=r"(*reinterpret_cast(&mv[(12) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 1728)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[16])), "=r"(*reinterpret_cast(&mv[(16) + 1])), "=r"(*reinterpret_cast(&mv[(16) + 2])), "=r"(*reinterpret_cast(&mv[(16) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 2304)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[20])), "=r"(*reinterpret_cast(&mv[(20) + 1])), "=r"(*reinterpret_cast(&mv[(20) + 2])), "=r"(*reinterpret_cast(&mv[(20) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 2880)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[24])), "=r"(*reinterpret_cast(&mv[(24) + 1])), "=r"(*reinterpret_cast(&mv[(24) + 2])), "=r"(*reinterpret_cast(&mv[(24) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 3456)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[28])), "=r"(*reinterpret_cast(&mv[(28) + 1])), "=r"(*reinterpret_cast(&mv[(28) + 2])), "=r"(*reinterpret_cast(&mv[(28) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 4032)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 128)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[0], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[0], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[0], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[0], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[4], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[4], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[4], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[4], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[4], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[4], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[4], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[4], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[8], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[8], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[8], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[8], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[8], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[8], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[8], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[8], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[12], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[12], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[12], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[12], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[12], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[12], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[12], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[12], xv[7], acc[31]); + acc[32] = __fmaf_rn(mv[16], xv[0], acc[32]); + acc[33] = __fmaf_rn(mv[16], xv[1], acc[33]); + acc[34] = __fmaf_rn(mv[16], xv[2], acc[34]); + acc[35] = __fmaf_rn(mv[16], xv[3], acc[35]); + acc[36] = __fmaf_rn(mv[16], xv[4], acc[36]); + acc[37] = __fmaf_rn(mv[16], xv[5], acc[37]); + acc[38] = __fmaf_rn(mv[16], xv[6], acc[38]); + acc[39] = __fmaf_rn(mv[16], xv[7], acc[39]); + acc[40] = __fmaf_rn(mv[20], xv[0], acc[40]); + acc[41] = __fmaf_rn(mv[20], xv[1], acc[41]); + acc[42] = __fmaf_rn(mv[20], xv[2], acc[42]); + acc[43] = __fmaf_rn(mv[20], xv[3], acc[43]); + acc[44] = __fmaf_rn(mv[20], xv[4], acc[44]); + acc[45] = __fmaf_rn(mv[20], xv[5], acc[45]); + acc[46] = __fmaf_rn(mv[20], xv[6], acc[46]); + acc[47] = __fmaf_rn(mv[20], xv[7], acc[47]); + acc[48] = __fmaf_rn(mv[24], xv[0], acc[48]); + acc[49] = __fmaf_rn(mv[24], xv[1], acc[49]); + acc[50] = __fmaf_rn(mv[24], xv[2], acc[50]); + acc[51] = __fmaf_rn(mv[24], xv[3], acc[51]); + acc[52] = __fmaf_rn(mv[24], xv[4], acc[52]); + acc[53] = __fmaf_rn(mv[24], xv[5], acc[53]); + acc[54] = __fmaf_rn(mv[24], xv[6], acc[54]); + acc[55] = __fmaf_rn(mv[24], xv[7], acc[55]); + acc[56] = __fmaf_rn(mv[28], xv[0], acc[56]); + acc[57] = __fmaf_rn(mv[28], xv[1], acc[57]); + acc[58] = __fmaf_rn(mv[28], xv[2], acc[58]); + acc[59] = __fmaf_rn(mv[28], xv[3], acc[59]); + acc[60] = __fmaf_rn(mv[28], xv[4], acc[60]); + acc[61] = __fmaf_rn(mv[28], xv[5], acc[61]); + acc[62] = __fmaf_rn(mv[28], xv[6], acc[62]); + acc[63] = __fmaf_rn(mv[28], xv[7], acc[63]); +#else + float2 m2_8 = make_float2(mv[0], mv[0]); + float2 a2_32 = make_float2(acc[0], acc[1]); + a2_32 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[0], xv[1]), a2_32); + acc[0] = a2_32.x; + acc[1] = a2_32.y; + float2 a2_33 = make_float2(acc[2], acc[3]); + a2_33 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[2], xv[3]), a2_33); + acc[2] = a2_33.x; + acc[3] = a2_33.y; + float2 a2_34 = make_float2(acc[4], acc[5]); + a2_34 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[4], xv[5]), a2_34); + acc[4] = a2_34.x; + acc[5] = a2_34.y; + float2 a2_35 = make_float2(acc[6], acc[7]); + a2_35 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[6], xv[7]), a2_35); + acc[6] = a2_35.x; + acc[7] = a2_35.y; + float2 m2_9 = make_float2(mv[4], mv[4]); + float2 a2_36 = make_float2(acc[8], acc[9]); + a2_36 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[0], xv[1]), a2_36); + acc[8] = a2_36.x; + acc[9] = a2_36.y; + float2 a2_37 = make_float2(acc[10], acc[11]); + a2_37 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[2], xv[3]), a2_37); + acc[10] = a2_37.x; + acc[11] = a2_37.y; + float2 a2_38 = make_float2(acc[12], acc[13]); + a2_38 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[4], xv[5]), a2_38); + acc[12] = a2_38.x; + acc[13] = a2_38.y; + float2 a2_39 = make_float2(acc[14], acc[15]); + a2_39 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[6], xv[7]), a2_39); + acc[14] = a2_39.x; + acc[15] = a2_39.y; + float2 m2_10 = make_float2(mv[8], mv[8]); + float2 a2_40 = make_float2(acc[16], acc[17]); + a2_40 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[0], xv[1]), a2_40); + acc[16] = a2_40.x; + acc[17] = a2_40.y; + float2 a2_41 = make_float2(acc[18], acc[19]); + a2_41 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[2], xv[3]), a2_41); + acc[18] = a2_41.x; + acc[19] = a2_41.y; + float2 a2_42 = make_float2(acc[20], acc[21]); + a2_42 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[4], xv[5]), a2_42); + acc[20] = a2_42.x; + acc[21] = a2_42.y; + float2 a2_43 = make_float2(acc[22], acc[23]); + a2_43 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[6], xv[7]), a2_43); + acc[22] = a2_43.x; + acc[23] = a2_43.y; + float2 m2_11 = make_float2(mv[12], mv[12]); + float2 a2_44 = make_float2(acc[24], acc[25]); + a2_44 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[0], xv[1]), a2_44); + acc[24] = a2_44.x; + acc[25] = a2_44.y; + float2 a2_45 = make_float2(acc[26], acc[27]); + a2_45 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[2], xv[3]), a2_45); + acc[26] = a2_45.x; + acc[27] = a2_45.y; + float2 a2_46 = make_float2(acc[28], acc[29]); + a2_46 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[4], xv[5]), a2_46); + acc[28] = a2_46.x; + acc[29] = a2_46.y; + float2 a2_47 = make_float2(acc[30], acc[31]); + a2_47 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[6], xv[7]), a2_47); + acc[30] = a2_47.x; + acc[31] = a2_47.y; + float2 m2_12 = make_float2(mv[16], mv[16]); + float2 a2_48 = make_float2(acc[32], acc[33]); + a2_48 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[0], xv[1]), a2_48); + acc[32] = a2_48.x; + acc[33] = a2_48.y; + float2 a2_49 = make_float2(acc[34], acc[35]); + a2_49 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[2], xv[3]), a2_49); + acc[34] = a2_49.x; + acc[35] = a2_49.y; + float2 a2_50 = make_float2(acc[36], acc[37]); + a2_50 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[4], xv[5]), a2_50); + acc[36] = a2_50.x; + acc[37] = a2_50.y; + float2 a2_51 = make_float2(acc[38], acc[39]); + a2_51 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[6], xv[7]), a2_51); + acc[38] = a2_51.x; + acc[39] = a2_51.y; + float2 m2_13 = make_float2(mv[20], mv[20]); + float2 a2_52 = make_float2(acc[40], acc[41]); + a2_52 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[0], xv[1]), a2_52); + acc[40] = a2_52.x; + acc[41] = a2_52.y; + float2 a2_53 = make_float2(acc[42], acc[43]); + a2_53 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[2], xv[3]), a2_53); + acc[42] = a2_53.x; + acc[43] = a2_53.y; + float2 a2_54 = make_float2(acc[44], acc[45]); + a2_54 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[4], xv[5]), a2_54); + acc[44] = a2_54.x; + acc[45] = a2_54.y; + float2 a2_55 = make_float2(acc[46], acc[47]); + a2_55 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[6], xv[7]), a2_55); + acc[46] = a2_55.x; + acc[47] = a2_55.y; + float2 m2_14 = make_float2(mv[24], mv[24]); + float2 a2_56 = make_float2(acc[48], acc[49]); + a2_56 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[0], xv[1]), a2_56); + acc[48] = a2_56.x; + acc[49] = a2_56.y; + float2 a2_57 = make_float2(acc[50], acc[51]); + a2_57 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[2], xv[3]), a2_57); + acc[50] = a2_57.x; + acc[51] = a2_57.y; + float2 a2_58 = make_float2(acc[52], acc[53]); + a2_58 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[4], xv[5]), a2_58); + acc[52] = a2_58.x; + acc[53] = a2_58.y; + float2 a2_59 = make_float2(acc[54], acc[55]); + a2_59 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[6], xv[7]), a2_59); + acc[54] = a2_59.x; + acc[55] = a2_59.y; + float2 m2_15 = make_float2(mv[28], mv[28]); + float2 a2_60 = make_float2(acc[56], acc[57]); + a2_60 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[0], xv[1]), a2_60); + acc[56] = a2_60.x; + acc[57] = a2_60.y; + float2 a2_61 = make_float2(acc[58], acc[59]); + a2_61 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[2], xv[3]), a2_61); + acc[58] = a2_61.x; + acc[59] = a2_61.y; + float2 a2_62 = make_float2(acc[60], acc[61]); + a2_62 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[4], xv[5]), a2_62); + acc[60] = a2_62.x; + acc[61] = a2_62.y; + float2 a2_63 = make_float2(acc[62], acc[63]); + a2_63 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[6], xv[7]), a2_63); + acc[62] = a2_63.x; + acc[63] = a2_63.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 272)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 400)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[1], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[1], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[1], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[1], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[1], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[1], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[1], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[1], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[5], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[5], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[5], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[5], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[5], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[5], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[5], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[5], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[9], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[9], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[9], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[9], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[9], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[9], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[9], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[9], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[13], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[13], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[13], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[13], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[13], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[13], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[13], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[13], xv[7], acc[31]); + acc[32] = __fmaf_rn(mv[17], xv[0], acc[32]); + acc[33] = __fmaf_rn(mv[17], xv[1], acc[33]); + acc[34] = __fmaf_rn(mv[17], xv[2], acc[34]); + acc[35] = __fmaf_rn(mv[17], xv[3], acc[35]); + acc[36] = __fmaf_rn(mv[17], xv[4], acc[36]); + acc[37] = __fmaf_rn(mv[17], xv[5], acc[37]); + acc[38] = __fmaf_rn(mv[17], xv[6], acc[38]); + acc[39] = __fmaf_rn(mv[17], xv[7], acc[39]); + acc[40] = __fmaf_rn(mv[21], xv[0], acc[40]); + acc[41] = __fmaf_rn(mv[21], xv[1], acc[41]); + acc[42] = __fmaf_rn(mv[21], xv[2], acc[42]); + acc[43] = __fmaf_rn(mv[21], xv[3], acc[43]); + acc[44] = __fmaf_rn(mv[21], xv[4], acc[44]); + acc[45] = __fmaf_rn(mv[21], xv[5], acc[45]); + acc[46] = __fmaf_rn(mv[21], xv[6], acc[46]); + acc[47] = __fmaf_rn(mv[21], xv[7], acc[47]); + acc[48] = __fmaf_rn(mv[25], xv[0], acc[48]); + acc[49] = __fmaf_rn(mv[25], xv[1], acc[49]); + acc[50] = __fmaf_rn(mv[25], xv[2], acc[50]); + acc[51] = __fmaf_rn(mv[25], xv[3], acc[51]); + acc[52] = __fmaf_rn(mv[25], xv[4], acc[52]); + acc[53] = __fmaf_rn(mv[25], xv[5], acc[53]); + acc[54] = __fmaf_rn(mv[25], xv[6], acc[54]); + acc[55] = __fmaf_rn(mv[25], xv[7], acc[55]); + acc[56] = __fmaf_rn(mv[29], xv[0], acc[56]); + acc[57] = __fmaf_rn(mv[29], xv[1], acc[57]); + acc[58] = __fmaf_rn(mv[29], xv[2], acc[58]); + acc[59] = __fmaf_rn(mv[29], xv[3], acc[59]); + acc[60] = __fmaf_rn(mv[29], xv[4], acc[60]); + acc[61] = __fmaf_rn(mv[29], xv[5], acc[61]); + acc[62] = __fmaf_rn(mv[29], xv[6], acc[62]); + acc[63] = __fmaf_rn(mv[29], xv[7], acc[63]); +#else + float2 m2_16 = make_float2(mv[1], mv[1]); + float2 a2_64 = make_float2(acc[0], acc[1]); + a2_64 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[0], xv[1]), a2_64); + acc[0] = a2_64.x; + acc[1] = a2_64.y; + float2 a2_65 = make_float2(acc[2], acc[3]); + a2_65 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[2], xv[3]), a2_65); + acc[2] = a2_65.x; + acc[3] = a2_65.y; + float2 a2_66 = make_float2(acc[4], acc[5]); + a2_66 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[4], xv[5]), a2_66); + acc[4] = a2_66.x; + acc[5] = a2_66.y; + float2 a2_67 = make_float2(acc[6], acc[7]); + a2_67 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[6], xv[7]), a2_67); + acc[6] = a2_67.x; + acc[7] = a2_67.y; + float2 m2_17 = make_float2(mv[5], mv[5]); + float2 a2_68 = make_float2(acc[8], acc[9]); + a2_68 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[0], xv[1]), a2_68); + acc[8] = a2_68.x; + acc[9] = a2_68.y; + float2 a2_69 = make_float2(acc[10], acc[11]); + a2_69 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[2], xv[3]), a2_69); + acc[10] = a2_69.x; + acc[11] = a2_69.y; + float2 a2_70 = make_float2(acc[12], acc[13]); + a2_70 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[4], xv[5]), a2_70); + acc[12] = a2_70.x; + acc[13] = a2_70.y; + float2 a2_71 = make_float2(acc[14], acc[15]); + a2_71 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[6], xv[7]), a2_71); + acc[14] = a2_71.x; + acc[15] = a2_71.y; + float2 m2_18 = make_float2(mv[9], mv[9]); + float2 a2_72 = make_float2(acc[16], acc[17]); + a2_72 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[0], xv[1]), a2_72); + acc[16] = a2_72.x; + acc[17] = a2_72.y; + float2 a2_73 = make_float2(acc[18], acc[19]); + a2_73 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[2], xv[3]), a2_73); + acc[18] = a2_73.x; + acc[19] = a2_73.y; + float2 a2_74 = make_float2(acc[20], acc[21]); + a2_74 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[4], xv[5]), a2_74); + acc[20] = a2_74.x; + acc[21] = a2_74.y; + float2 a2_75 = make_float2(acc[22], acc[23]); + a2_75 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[6], xv[7]), a2_75); + acc[22] = a2_75.x; + acc[23] = a2_75.y; + float2 m2_19 = make_float2(mv[13], mv[13]); + float2 a2_76 = make_float2(acc[24], acc[25]); + a2_76 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[0], xv[1]), a2_76); + acc[24] = a2_76.x; + acc[25] = a2_76.y; + float2 a2_77 = make_float2(acc[26], acc[27]); + a2_77 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[2], xv[3]), a2_77); + acc[26] = a2_77.x; + acc[27] = a2_77.y; + float2 a2_78 = make_float2(acc[28], acc[29]); + a2_78 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[4], xv[5]), a2_78); + acc[28] = a2_78.x; + acc[29] = a2_78.y; + float2 a2_79 = make_float2(acc[30], acc[31]); + a2_79 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[6], xv[7]), a2_79); + acc[30] = a2_79.x; + acc[31] = a2_79.y; + float2 m2_20 = make_float2(mv[17], mv[17]); + float2 a2_80 = make_float2(acc[32], acc[33]); + a2_80 = fma_f32x2_rn_ftz(m2_20, make_float2(xv[0], xv[1]), a2_80); + acc[32] = a2_80.x; + acc[33] = a2_80.y; + float2 a2_81 = make_float2(acc[34], acc[35]); + a2_81 = fma_f32x2_rn_ftz(m2_20, make_float2(xv[2], xv[3]), a2_81); + acc[34] = a2_81.x; + acc[35] = a2_81.y; + float2 a2_82 = make_float2(acc[36], acc[37]); + a2_82 = fma_f32x2_rn_ftz(m2_20, make_float2(xv[4], xv[5]), a2_82); + acc[36] = a2_82.x; + acc[37] = a2_82.y; + float2 a2_83 = make_float2(acc[38], acc[39]); + a2_83 = fma_f32x2_rn_ftz(m2_20, make_float2(xv[6], xv[7]), a2_83); + acc[38] = a2_83.x; + acc[39] = a2_83.y; + float2 m2_21 = make_float2(mv[21], mv[21]); + float2 a2_84 = make_float2(acc[40], acc[41]); + a2_84 = fma_f32x2_rn_ftz(m2_21, make_float2(xv[0], xv[1]), a2_84); + acc[40] = a2_84.x; + acc[41] = a2_84.y; + float2 a2_85 = make_float2(acc[42], acc[43]); + a2_85 = fma_f32x2_rn_ftz(m2_21, make_float2(xv[2], xv[3]), a2_85); + acc[42] = a2_85.x; + acc[43] = a2_85.y; + float2 a2_86 = make_float2(acc[44], acc[45]); + a2_86 = fma_f32x2_rn_ftz(m2_21, make_float2(xv[4], xv[5]), a2_86); + acc[44] = a2_86.x; + acc[45] = a2_86.y; + float2 a2_87 = make_float2(acc[46], acc[47]); + a2_87 = fma_f32x2_rn_ftz(m2_21, make_float2(xv[6], xv[7]), a2_87); + acc[46] = a2_87.x; + acc[47] = a2_87.y; + float2 m2_22 = make_float2(mv[25], mv[25]); + float2 a2_88 = make_float2(acc[48], acc[49]); + a2_88 = fma_f32x2_rn_ftz(m2_22, make_float2(xv[0], xv[1]), a2_88); + acc[48] = a2_88.x; + acc[49] = a2_88.y; + float2 a2_89 = make_float2(acc[50], acc[51]); + a2_89 = fma_f32x2_rn_ftz(m2_22, make_float2(xv[2], xv[3]), a2_89); + acc[50] = a2_89.x; + acc[51] = a2_89.y; + float2 a2_90 = make_float2(acc[52], acc[53]); + a2_90 = fma_f32x2_rn_ftz(m2_22, make_float2(xv[4], xv[5]), a2_90); + acc[52] = a2_90.x; + acc[53] = a2_90.y; + float2 a2_91 = make_float2(acc[54], acc[55]); + a2_91 = fma_f32x2_rn_ftz(m2_22, make_float2(xv[6], xv[7]), a2_91); + acc[54] = a2_91.x; + acc[55] = a2_91.y; + float2 m2_23 = make_float2(mv[29], mv[29]); + float2 a2_92 = make_float2(acc[56], acc[57]); + a2_92 = fma_f32x2_rn_ftz(m2_23, make_float2(xv[0], xv[1]), a2_92); + acc[56] = a2_92.x; + acc[57] = a2_92.y; + float2 a2_93 = make_float2(acc[58], acc[59]); + a2_93 = fma_f32x2_rn_ftz(m2_23, make_float2(xv[2], xv[3]), a2_93); + acc[58] = a2_93.x; + acc[59] = a2_93.y; + float2 a2_94 = make_float2(acc[60], acc[61]); + a2_94 = fma_f32x2_rn_ftz(m2_23, make_float2(xv[4], xv[5]), a2_94); + acc[60] = a2_94.x; + acc[61] = a2_94.y; + float2 a2_95 = make_float2(acc[62], acc[63]); + a2_95 = fma_f32x2_rn_ftz(m2_23, make_float2(xv[6], xv[7]), a2_95); + acc[62] = a2_95.x; + acc[63] = a2_95.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 544)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 672)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[2], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[2], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[2], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[2], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[2], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[2], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[2], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[2], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[6], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[6], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[6], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[6], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[6], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[6], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[6], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[6], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[10], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[10], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[10], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[10], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[10], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[10], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[10], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[10], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[14], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[14], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[14], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[14], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[14], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[14], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[14], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[14], xv[7], acc[31]); + acc[32] = __fmaf_rn(mv[18], xv[0], acc[32]); + acc[33] = __fmaf_rn(mv[18], xv[1], acc[33]); + acc[34] = __fmaf_rn(mv[18], xv[2], acc[34]); + acc[35] = __fmaf_rn(mv[18], xv[3], acc[35]); + acc[36] = __fmaf_rn(mv[18], xv[4], acc[36]); + acc[37] = __fmaf_rn(mv[18], xv[5], acc[37]); + acc[38] = __fmaf_rn(mv[18], xv[6], acc[38]); + acc[39] = __fmaf_rn(mv[18], xv[7], acc[39]); + acc[40] = __fmaf_rn(mv[22], xv[0], acc[40]); + acc[41] = __fmaf_rn(mv[22], xv[1], acc[41]); + acc[42] = __fmaf_rn(mv[22], xv[2], acc[42]); + acc[43] = __fmaf_rn(mv[22], xv[3], acc[43]); + acc[44] = __fmaf_rn(mv[22], xv[4], acc[44]); + acc[45] = __fmaf_rn(mv[22], xv[5], acc[45]); + acc[46] = __fmaf_rn(mv[22], xv[6], acc[46]); + acc[47] = __fmaf_rn(mv[22], xv[7], acc[47]); + acc[48] = __fmaf_rn(mv[26], xv[0], acc[48]); + acc[49] = __fmaf_rn(mv[26], xv[1], acc[49]); + acc[50] = __fmaf_rn(mv[26], xv[2], acc[50]); + acc[51] = __fmaf_rn(mv[26], xv[3], acc[51]); + acc[52] = __fmaf_rn(mv[26], xv[4], acc[52]); + acc[53] = __fmaf_rn(mv[26], xv[5], acc[53]); + acc[54] = __fmaf_rn(mv[26], xv[6], acc[54]); + acc[55] = __fmaf_rn(mv[26], xv[7], acc[55]); + acc[56] = __fmaf_rn(mv[30], xv[0], acc[56]); + acc[57] = __fmaf_rn(mv[30], xv[1], acc[57]); + acc[58] = __fmaf_rn(mv[30], xv[2], acc[58]); + acc[59] = __fmaf_rn(mv[30], xv[3], acc[59]); + acc[60] = __fmaf_rn(mv[30], xv[4], acc[60]); + acc[61] = __fmaf_rn(mv[30], xv[5], acc[61]); + acc[62] = __fmaf_rn(mv[30], xv[6], acc[62]); + acc[63] = __fmaf_rn(mv[30], xv[7], acc[63]); +#else + float2 m2_24 = make_float2(mv[2], mv[2]); + float2 a2_96 = make_float2(acc[0], acc[1]); + a2_96 = fma_f32x2_rn_ftz(m2_24, make_float2(xv[0], xv[1]), a2_96); + acc[0] = a2_96.x; + acc[1] = a2_96.y; + float2 a2_97 = make_float2(acc[2], acc[3]); + a2_97 = fma_f32x2_rn_ftz(m2_24, make_float2(xv[2], xv[3]), a2_97); + acc[2] = a2_97.x; + acc[3] = a2_97.y; + float2 a2_98 = make_float2(acc[4], acc[5]); + a2_98 = fma_f32x2_rn_ftz(m2_24, make_float2(xv[4], xv[5]), a2_98); + acc[4] = a2_98.x; + acc[5] = a2_98.y; + float2 a2_99 = make_float2(acc[6], acc[7]); + a2_99 = fma_f32x2_rn_ftz(m2_24, make_float2(xv[6], xv[7]), a2_99); + acc[6] = a2_99.x; + acc[7] = a2_99.y; + float2 m2_25 = make_float2(mv[6], mv[6]); + float2 a2_100 = make_float2(acc[8], acc[9]); + a2_100 = fma_f32x2_rn_ftz(m2_25, make_float2(xv[0], xv[1]), a2_100); + acc[8] = a2_100.x; + acc[9] = a2_100.y; + float2 a2_101 = make_float2(acc[10], acc[11]); + a2_101 = fma_f32x2_rn_ftz(m2_25, make_float2(xv[2], xv[3]), a2_101); + acc[10] = a2_101.x; + acc[11] = a2_101.y; + float2 a2_102 = make_float2(acc[12], acc[13]); + a2_102 = fma_f32x2_rn_ftz(m2_25, make_float2(xv[4], xv[5]), a2_102); + acc[12] = a2_102.x; + acc[13] = a2_102.y; + float2 a2_103 = make_float2(acc[14], acc[15]); + a2_103 = fma_f32x2_rn_ftz(m2_25, make_float2(xv[6], xv[7]), a2_103); + acc[14] = a2_103.x; + acc[15] = a2_103.y; + float2 m2_26 = make_float2(mv[10], mv[10]); + float2 a2_104 = make_float2(acc[16], acc[17]); + a2_104 = fma_f32x2_rn_ftz(m2_26, make_float2(xv[0], xv[1]), a2_104); + acc[16] = a2_104.x; + acc[17] = a2_104.y; + float2 a2_105 = make_float2(acc[18], acc[19]); + a2_105 = fma_f32x2_rn_ftz(m2_26, make_float2(xv[2], xv[3]), a2_105); + acc[18] = a2_105.x; + acc[19] = a2_105.y; + float2 a2_106 = make_float2(acc[20], acc[21]); + a2_106 = fma_f32x2_rn_ftz(m2_26, make_float2(xv[4], xv[5]), a2_106); + acc[20] = a2_106.x; + acc[21] = a2_106.y; + float2 a2_107 = make_float2(acc[22], acc[23]); + a2_107 = fma_f32x2_rn_ftz(m2_26, make_float2(xv[6], xv[7]), a2_107); + acc[22] = a2_107.x; + acc[23] = a2_107.y; + float2 m2_27 = make_float2(mv[14], mv[14]); + float2 a2_108 = make_float2(acc[24], acc[25]); + a2_108 = fma_f32x2_rn_ftz(m2_27, make_float2(xv[0], xv[1]), a2_108); + acc[24] = a2_108.x; + acc[25] = a2_108.y; + float2 a2_109 = make_float2(acc[26], acc[27]); + a2_109 = fma_f32x2_rn_ftz(m2_27, make_float2(xv[2], xv[3]), a2_109); + acc[26] = a2_109.x; + acc[27] = a2_109.y; + float2 a2_110 = make_float2(acc[28], acc[29]); + a2_110 = fma_f32x2_rn_ftz(m2_27, make_float2(xv[4], xv[5]), a2_110); + acc[28] = a2_110.x; + acc[29] = a2_110.y; + float2 a2_111 = make_float2(acc[30], acc[31]); + a2_111 = fma_f32x2_rn_ftz(m2_27, make_float2(xv[6], xv[7]), a2_111); + acc[30] = a2_111.x; + acc[31] = a2_111.y; + float2 m2_28 = make_float2(mv[18], mv[18]); + float2 a2_112 = make_float2(acc[32], acc[33]); + a2_112 = fma_f32x2_rn_ftz(m2_28, make_float2(xv[0], xv[1]), a2_112); + acc[32] = a2_112.x; + acc[33] = a2_112.y; + float2 a2_113 = make_float2(acc[34], acc[35]); + a2_113 = fma_f32x2_rn_ftz(m2_28, make_float2(xv[2], xv[3]), a2_113); + acc[34] = a2_113.x; + acc[35] = a2_113.y; + float2 a2_114 = make_float2(acc[36], acc[37]); + a2_114 = fma_f32x2_rn_ftz(m2_28, make_float2(xv[4], xv[5]), a2_114); + acc[36] = a2_114.x; + acc[37] = a2_114.y; + float2 a2_115 = make_float2(acc[38], acc[39]); + a2_115 = fma_f32x2_rn_ftz(m2_28, make_float2(xv[6], xv[7]), a2_115); + acc[38] = a2_115.x; + acc[39] = a2_115.y; + float2 m2_29 = make_float2(mv[22], mv[22]); + float2 a2_116 = make_float2(acc[40], acc[41]); + a2_116 = fma_f32x2_rn_ftz(m2_29, make_float2(xv[0], xv[1]), a2_116); + acc[40] = a2_116.x; + acc[41] = a2_116.y; + float2 a2_117 = make_float2(acc[42], acc[43]); + a2_117 = fma_f32x2_rn_ftz(m2_29, make_float2(xv[2], xv[3]), a2_117); + acc[42] = a2_117.x; + acc[43] = a2_117.y; + float2 a2_118 = make_float2(acc[44], acc[45]); + a2_118 = fma_f32x2_rn_ftz(m2_29, make_float2(xv[4], xv[5]), a2_118); + acc[44] = a2_118.x; + acc[45] = a2_118.y; + float2 a2_119 = make_float2(acc[46], acc[47]); + a2_119 = fma_f32x2_rn_ftz(m2_29, make_float2(xv[6], xv[7]), a2_119); + acc[46] = a2_119.x; + acc[47] = a2_119.y; + float2 m2_30 = make_float2(mv[26], mv[26]); + float2 a2_120 = make_float2(acc[48], acc[49]); + a2_120 = fma_f32x2_rn_ftz(m2_30, make_float2(xv[0], xv[1]), a2_120); + acc[48] = a2_120.x; + acc[49] = a2_120.y; + float2 a2_121 = make_float2(acc[50], acc[51]); + a2_121 = fma_f32x2_rn_ftz(m2_30, make_float2(xv[2], xv[3]), a2_121); + acc[50] = a2_121.x; + acc[51] = a2_121.y; + float2 a2_122 = make_float2(acc[52], acc[53]); + a2_122 = fma_f32x2_rn_ftz(m2_30, make_float2(xv[4], xv[5]), a2_122); + acc[52] = a2_122.x; + acc[53] = a2_122.y; + float2 a2_123 = make_float2(acc[54], acc[55]); + a2_123 = fma_f32x2_rn_ftz(m2_30, make_float2(xv[6], xv[7]), a2_123); + acc[54] = a2_123.x; + acc[55] = a2_123.y; + float2 m2_31 = make_float2(mv[30], mv[30]); + float2 a2_124 = make_float2(acc[56], acc[57]); + a2_124 = fma_f32x2_rn_ftz(m2_31, make_float2(xv[0], xv[1]), a2_124); + acc[56] = a2_124.x; + acc[57] = a2_124.y; + float2 a2_125 = make_float2(acc[58], acc[59]); + a2_125 = fma_f32x2_rn_ftz(m2_31, make_float2(xv[2], xv[3]), a2_125); + acc[58] = a2_125.x; + acc[59] = a2_125.y; + float2 a2_126 = make_float2(acc[60], acc[61]); + a2_126 = fma_f32x2_rn_ftz(m2_31, make_float2(xv[4], xv[5]), a2_126); + acc[60] = a2_126.x; + acc[61] = a2_126.y; + float2 a2_127 = make_float2(acc[62], acc[63]); + a2_127 = fma_f32x2_rn_ftz(m2_31, make_float2(xv[6], xv[7]), a2_127); + acc[62] = a2_127.x; + acc[63] = a2_127.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 816)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 944)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[3], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[3], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[3], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[3], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[3], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[3], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[3], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[3], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[7], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[7], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[7], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[7], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[7], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[7], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[7], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[7], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[11], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[11], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[11], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[11], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[11], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[11], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[11], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[11], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[15], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[15], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[15], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[15], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[15], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[15], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[15], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[15], xv[7], acc[31]); + acc[32] = __fmaf_rn(mv[19], xv[0], acc[32]); + acc[33] = __fmaf_rn(mv[19], xv[1], acc[33]); + acc[34] = __fmaf_rn(mv[19], xv[2], acc[34]); + acc[35] = __fmaf_rn(mv[19], xv[3], acc[35]); + acc[36] = __fmaf_rn(mv[19], xv[4], acc[36]); + acc[37] = __fmaf_rn(mv[19], xv[5], acc[37]); + acc[38] = __fmaf_rn(mv[19], xv[6], acc[38]); + acc[39] = __fmaf_rn(mv[19], xv[7], acc[39]); + acc[40] = __fmaf_rn(mv[23], xv[0], acc[40]); + acc[41] = __fmaf_rn(mv[23], xv[1], acc[41]); + acc[42] = __fmaf_rn(mv[23], xv[2], acc[42]); + acc[43] = __fmaf_rn(mv[23], xv[3], acc[43]); + acc[44] = __fmaf_rn(mv[23], xv[4], acc[44]); + acc[45] = __fmaf_rn(mv[23], xv[5], acc[45]); + acc[46] = __fmaf_rn(mv[23], xv[6], acc[46]); + acc[47] = __fmaf_rn(mv[23], xv[7], acc[47]); + acc[48] = __fmaf_rn(mv[27], xv[0], acc[48]); + acc[49] = __fmaf_rn(mv[27], xv[1], acc[49]); + acc[50] = __fmaf_rn(mv[27], xv[2], acc[50]); + acc[51] = __fmaf_rn(mv[27], xv[3], acc[51]); + acc[52] = __fmaf_rn(mv[27], xv[4], acc[52]); + acc[53] = __fmaf_rn(mv[27], xv[5], acc[53]); + acc[54] = __fmaf_rn(mv[27], xv[6], acc[54]); + acc[55] = __fmaf_rn(mv[27], xv[7], acc[55]); + acc[56] = __fmaf_rn(mv[31], xv[0], acc[56]); + acc[57] = __fmaf_rn(mv[31], xv[1], acc[57]); + acc[58] = __fmaf_rn(mv[31], xv[2], acc[58]); + acc[59] = __fmaf_rn(mv[31], xv[3], acc[59]); + acc[60] = __fmaf_rn(mv[31], xv[4], acc[60]); + acc[61] = __fmaf_rn(mv[31], xv[5], acc[61]); + acc[62] = __fmaf_rn(mv[31], xv[6], acc[62]); + acc[63] = __fmaf_rn(mv[31], xv[7], acc[63]); +#else + float2 m2_32 = make_float2(mv[3], mv[3]); + float2 a2_128 = make_float2(acc[0], acc[1]); + a2_128 = fma_f32x2_rn_ftz(m2_32, make_float2(xv[0], xv[1]), a2_128); + acc[0] = a2_128.x; + acc[1] = a2_128.y; + float2 a2_129 = make_float2(acc[2], acc[3]); + a2_129 = fma_f32x2_rn_ftz(m2_32, make_float2(xv[2], xv[3]), a2_129); + acc[2] = a2_129.x; + acc[3] = a2_129.y; + float2 a2_130 = make_float2(acc[4], acc[5]); + a2_130 = fma_f32x2_rn_ftz(m2_32, make_float2(xv[4], xv[5]), a2_130); + acc[4] = a2_130.x; + acc[5] = a2_130.y; + float2 a2_131 = make_float2(acc[6], acc[7]); + a2_131 = fma_f32x2_rn_ftz(m2_32, make_float2(xv[6], xv[7]), a2_131); + acc[6] = a2_131.x; + acc[7] = a2_131.y; + float2 m2_33 = make_float2(mv[7], mv[7]); + float2 a2_132 = make_float2(acc[8], acc[9]); + a2_132 = fma_f32x2_rn_ftz(m2_33, make_float2(xv[0], xv[1]), a2_132); + acc[8] = a2_132.x; + acc[9] = a2_132.y; + float2 a2_133 = make_float2(acc[10], acc[11]); + a2_133 = fma_f32x2_rn_ftz(m2_33, make_float2(xv[2], xv[3]), a2_133); + acc[10] = a2_133.x; + acc[11] = a2_133.y; + float2 a2_134 = make_float2(acc[12], acc[13]); + a2_134 = fma_f32x2_rn_ftz(m2_33, make_float2(xv[4], xv[5]), a2_134); + acc[12] = a2_134.x; + acc[13] = a2_134.y; + float2 a2_135 = make_float2(acc[14], acc[15]); + a2_135 = fma_f32x2_rn_ftz(m2_33, make_float2(xv[6], xv[7]), a2_135); + acc[14] = a2_135.x; + acc[15] = a2_135.y; + float2 m2_34 = make_float2(mv[11], mv[11]); + float2 a2_136 = make_float2(acc[16], acc[17]); + a2_136 = fma_f32x2_rn_ftz(m2_34, make_float2(xv[0], xv[1]), a2_136); + acc[16] = a2_136.x; + acc[17] = a2_136.y; + float2 a2_137 = make_float2(acc[18], acc[19]); + a2_137 = fma_f32x2_rn_ftz(m2_34, make_float2(xv[2], xv[3]), a2_137); + acc[18] = a2_137.x; + acc[19] = a2_137.y; + float2 a2_138 = make_float2(acc[20], acc[21]); + a2_138 = fma_f32x2_rn_ftz(m2_34, make_float2(xv[4], xv[5]), a2_138); + acc[20] = a2_138.x; + acc[21] = a2_138.y; + float2 a2_139 = make_float2(acc[22], acc[23]); + a2_139 = fma_f32x2_rn_ftz(m2_34, make_float2(xv[6], xv[7]), a2_139); + acc[22] = a2_139.x; + acc[23] = a2_139.y; + float2 m2_35 = make_float2(mv[15], mv[15]); + float2 a2_140 = make_float2(acc[24], acc[25]); + a2_140 = fma_f32x2_rn_ftz(m2_35, make_float2(xv[0], xv[1]), a2_140); + acc[24] = a2_140.x; + acc[25] = a2_140.y; + float2 a2_141 = make_float2(acc[26], acc[27]); + a2_141 = fma_f32x2_rn_ftz(m2_35, make_float2(xv[2], xv[3]), a2_141); + acc[26] = a2_141.x; + acc[27] = a2_141.y; + float2 a2_142 = make_float2(acc[28], acc[29]); + a2_142 = fma_f32x2_rn_ftz(m2_35, make_float2(xv[4], xv[5]), a2_142); + acc[28] = a2_142.x; + acc[29] = a2_142.y; + float2 a2_143 = make_float2(acc[30], acc[31]); + a2_143 = fma_f32x2_rn_ftz(m2_35, make_float2(xv[6], xv[7]), a2_143); + acc[30] = a2_143.x; + acc[31] = a2_143.y; + float2 m2_36 = make_float2(mv[19], mv[19]); + float2 a2_144 = make_float2(acc[32], acc[33]); + a2_144 = fma_f32x2_rn_ftz(m2_36, make_float2(xv[0], xv[1]), a2_144); + acc[32] = a2_144.x; + acc[33] = a2_144.y; + float2 a2_145 = make_float2(acc[34], acc[35]); + a2_145 = fma_f32x2_rn_ftz(m2_36, make_float2(xv[2], xv[3]), a2_145); + acc[34] = a2_145.x; + acc[35] = a2_145.y; + float2 a2_146 = make_float2(acc[36], acc[37]); + a2_146 = fma_f32x2_rn_ftz(m2_36, make_float2(xv[4], xv[5]), a2_146); + acc[36] = a2_146.x; + acc[37] = a2_146.y; + float2 a2_147 = make_float2(acc[38], acc[39]); + a2_147 = fma_f32x2_rn_ftz(m2_36, make_float2(xv[6], xv[7]), a2_147); + acc[38] = a2_147.x; + acc[39] = a2_147.y; + float2 m2_37 = make_float2(mv[23], mv[23]); + float2 a2_148 = make_float2(acc[40], acc[41]); + a2_148 = fma_f32x2_rn_ftz(m2_37, make_float2(xv[0], xv[1]), a2_148); + acc[40] = a2_148.x; + acc[41] = a2_148.y; + float2 a2_149 = make_float2(acc[42], acc[43]); + a2_149 = fma_f32x2_rn_ftz(m2_37, make_float2(xv[2], xv[3]), a2_149); + acc[42] = a2_149.x; + acc[43] = a2_149.y; + float2 a2_150 = make_float2(acc[44], acc[45]); + a2_150 = fma_f32x2_rn_ftz(m2_37, make_float2(xv[4], xv[5]), a2_150); + acc[44] = a2_150.x; + acc[45] = a2_150.y; + float2 a2_151 = make_float2(acc[46], acc[47]); + a2_151 = fma_f32x2_rn_ftz(m2_37, make_float2(xv[6], xv[7]), a2_151); + acc[46] = a2_151.x; + acc[47] = a2_151.y; + float2 m2_38 = make_float2(mv[27], mv[27]); + float2 a2_152 = make_float2(acc[48], acc[49]); + a2_152 = fma_f32x2_rn_ftz(m2_38, make_float2(xv[0], xv[1]), a2_152); + acc[48] = a2_152.x; + acc[49] = a2_152.y; + float2 a2_153 = make_float2(acc[50], acc[51]); + a2_153 = fma_f32x2_rn_ftz(m2_38, make_float2(xv[2], xv[3]), a2_153); + acc[50] = a2_153.x; + acc[51] = a2_153.y; + float2 a2_154 = make_float2(acc[52], acc[53]); + a2_154 = fma_f32x2_rn_ftz(m2_38, make_float2(xv[4], xv[5]), a2_154); + acc[52] = a2_154.x; + acc[53] = a2_154.y; + float2 a2_155 = make_float2(acc[54], acc[55]); + a2_155 = fma_f32x2_rn_ftz(m2_38, make_float2(xv[6], xv[7]), a2_155); + acc[54] = a2_155.x; + acc[55] = a2_155.y; + float2 m2_39 = make_float2(mv[31], mv[31]); + float2 a2_156 = make_float2(acc[56], acc[57]); + a2_156 = fma_f32x2_rn_ftz(m2_39, make_float2(xv[0], xv[1]), a2_156); + acc[56] = a2_156.x; + acc[57] = a2_156.y; + float2 a2_157 = make_float2(acc[58], acc[59]); + a2_157 = fma_f32x2_rn_ftz(m2_39, make_float2(xv[2], xv[3]), a2_157); + acc[58] = a2_157.x; + acc[59] = a2_157.y; + float2 a2_158 = make_float2(acc[60], acc[61]); + a2_158 = fma_f32x2_rn_ftz(m2_39, make_float2(xv[4], xv[5]), a2_158); + acc[60] = a2_158.x; + acc[61] = a2_158.y; + float2 a2_159 = make_float2(acc[62], acc[63]); + a2_159 = fma_f32x2_rn_ftz(m2_39, make_float2(xv[6], xv[7]), a2_159); + acc[62] = a2_159.x; + acc[63] = a2_159.y; +#endif + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 128), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 1088), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 1216), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 2176), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 2304), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 3264), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 3392), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 4352), "r"(__float_as_uint(acc[32])), "r"(__float_as_uint(acc[33])), "r"(__float_as_uint(acc[34])), "r"(__float_as_uint(acc[35])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 4480), "r"(__float_as_uint(acc[36])), "r"(__float_as_uint(acc[37])), "r"(__float_as_uint(acc[38])), "r"(__float_as_uint(acc[39])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 5440), "r"(__float_as_uint(acc[40])), "r"(__float_as_uint(acc[41])), "r"(__float_as_uint(acc[42])), "r"(__float_as_uint(acc[43])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 5568), "r"(__float_as_uint(acc[44])), "r"(__float_as_uint(acc[45])), "r"(__float_as_uint(acc[46])), "r"(__float_as_uint(acc[47])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 6528), "r"(__float_as_uint(acc[48])), "r"(__float_as_uint(acc[49])), "r"(__float_as_uint(acc[50])), "r"(__float_as_uint(acc[51])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 6656), "r"(__float_as_uint(acc[52])), "r"(__float_as_uint(acc[53])), "r"(__float_as_uint(acc[54])), "r"(__float_as_uint(acc[55])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 7616), "r"(__float_as_uint(acc[56])), "r"(__float_as_uint(acc[57])), "r"(__float_as_uint(acc[58])), "r"(__float_as_uint(acc[59])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816) + 7744), "r"(__float_as_uint(acc[60])), "r"(__float_as_uint(acc[61])), "r"(__float_as_uint(acc[62])), "r"(__float_as_uint(acc[63])) : "memory"); + if (y_pub_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 32) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 512) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 544) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 1056) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 1536) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 1568) = _v4; + } + { + float4 _v4 = make_float4(acc[32 + 0], acc[32 + 1], acc[32 + 2], acc[32 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[36 + 0], acc[36 + 1], acc[36 + 2], acc[36 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2080) = _v4; + } + { + float4 _v4 = make_float4(acc[40 + 0], acc[40 + 1], acc[40 + 2], acc[40 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2560) = _v4; + } + { + float4 _v4 = make_float4(acc[44 + 0], acc[44 + 1], acc[44 + 2], acc[44 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2592) = _v4; + } + { + float4 _v4 = make_float4(acc[48 + 0], acc[48 + 1], acc[48 + 2], acc[48 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 3072) = _v4; + } + { + float4 _v4 = make_float4(acc[52 + 0], acc[52 + 1], acc[52 + 2], acc[52 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 3104) = _v4; + } + { + float4 _v4 = make_float4(acc[56 + 0], acc[56 + 1], acc[56 + 2], acc[56 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 3584) = _v4; + } + { + float4 _v4 = make_float4(acc[60 + 0], acc[60 + 1], acc[60 + 2], acc[60 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 3616) = _v4; + } + } + __syncthreads(); + buf_1 = 1 - buf_1; + m_cur_1 = m_next_1; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } +} + +} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu deleted file mode 100644 index 82b12fcc758..00000000000 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_89f5ff59f03eb170e3bf_kcp_kernel.cu +++ /dev/null @@ -1,1312 +0,0 @@ -/* - * Copyright (c) 2026 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define NUM_MAIN_STAGES 1 -#define SMEM_X_BLOCK_OFF 0 -#define SMEM_X_BLOCK_STAGE_BYTES 8192 -#define SMEM_X_BLOCK_STRIDE 8192 -#define SMEM_SLABS_OFF 8192 -#define SMEM_SLABS_STAGE_BYTES 33024 -#define SMEM_SLABS_STRIDE 33024 -#define SMEM_TOTAL 41216 -#define THREADS 128 -#define LAUNCH_MIN_BLOCKS 2 - -#include - -extern "C" { - -__global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_kda_det_train_89f5ff59f03eb170e3bf_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(128) char smem_raw[]; - int smem; -#if __CUDA_ARCH__ == 1000 - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); -#else - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); -#endif - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - float* x_block = reinterpret_cast(smem_raw + 0); - const int x_block_addr = smem + 0; - float* slabs = reinterpret_cast(smem_raw + 8192); - const int slabs_addr = smem + 8192; - - // === Task calls (declaration order) === - int nblk = ncols / 16; - int chain = blockIdx.x / nblk; - int col0 = blockIdx.x % nblk * 16; - int col_tile = col0 / 128; - int v0 = col0 % 128; - int k = tid; - long long row_off = (long long)col_tile * 16384 + (long long)(k * 128 + v0); - int e_begin = chain_start[chain]; - int e_end = chain_start[chain + 1]; - unsigned int x_addr = x_block_addr + (unsigned int)(k * 64); - int w4r = warp * 4 + lane / 8; - int q4 = lane % 8 * 4; - float vals[16]; - float acc[16]; - float xv[4]; - float ring[64]; - float pf[32]; - long long x0_off = x0_index[chain]; - if (x0_off >= 0) { - { - float4 _v4 = *reinterpret_cast(x0 + x0_off + row_off); - vals[0 + 0] = _v4.x; - vals[0 + 1] = _v4.y; - vals[0 + 2] = _v4.z; - vals[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(x0 + x0_off + row_off + 4); - vals[4 + 0] = _v4.x; - vals[4 + 1] = _v4.y; - vals[4 + 2] = _v4.z; - vals[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(x0 + x0_off + row_off + 8); - vals[8 + 0] = _v4.x; - vals[8 + 1] = _v4.y; - vals[8 + 2] = _v4.z; - vals[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(x0 + x0_off + row_off + 12); - vals[12 + 0] = _v4.x; - vals[12 + 1] = _v4.y; - vals[12 + 2] = _v4.z; - vals[12 + 3] = _v4.w; - } - } else { - int diag = k - v0; - if (col_tile == 0) { - diag = -1; - } - #pragma unroll - for (int v = 0; v < 16; v++) { - vals[v] = ((v == diag) ? 1.0f : 0.0f); - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr), "r"(__float_as_uint(vals[0])), "r"(__float_as_uint(vals[1])), "r"(__float_as_uint(vals[2])), "r"(__float_as_uint(vals[3])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 16), "r"(__float_as_uint(vals[4])), "r"(__float_as_uint(vals[5])), "r"(__float_as_uint(vals[6])), "r"(__float_as_uint(vals[7])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 32), "r"(__float_as_uint(vals[8])), "r"(__float_as_uint(vals[9])), "r"(__float_as_uint(vals[10])), "r"(__float_as_uint(vals[11])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 48), "r"(__float_as_uint(vals[12])), "r"(__float_as_uint(vals[13])), "r"(__float_as_uint(vals[14])), "r"(__float_as_uint(vals[15])) : "memory"); - if (e_begin < e_end) { - if (write_states != 0) { - long long y_first = elem_y[e_begin]; - if (y_first >= 0) { - { - float4 _v4 = make_float4(vals[0 + 0], vals[0 + 1], vals[0 + 2], vals[0 + 3]); - *reinterpret_cast(y + y_first + row_off) = _v4; - } - { - float4 _v4 = make_float4(vals[4 + 0], vals[4 + 1], vals[4 + 2], vals[4 + 3]); - *reinterpret_cast(y + y_first + row_off + 4) = _v4; - } - { - float4 _v4 = make_float4(vals[8 + 0], vals[8 + 1], vals[8 + 2], vals[8 + 3]); - *reinterpret_cast(y + y_first + row_off + 8) = _v4; - } - { - float4 _v4 = make_float4(vals[12 + 0], vals[12 + 1], vals[12 + 2], vals[12 + 3]); - *reinterpret_cast(y + y_first + row_off + 12) = _v4; - } - } - } - } else { - long long y_last = y_final[chain]; - if (y_last >= 0) { - { - float4 _v4 = make_float4(vals[0 + 0], vals[0 + 1], vals[0 + 2], vals[0 + 3]); - *reinterpret_cast(y + y_last + row_off) = _v4; - } - { - float4 _v4 = make_float4(vals[4 + 0], vals[4 + 1], vals[4 + 2], vals[4 + 3]); - *reinterpret_cast(y + y_last + row_off + 4) = _v4; - } - { - float4 _v4 = make_float4(vals[8 + 0], vals[8 + 1], vals[8 + 2], vals[8 + 3]); - *reinterpret_cast(y + y_last + row_off + 8) = _v4; - } - { - float4 _v4 = make_float4(vals[12 + 0], vals[12 + 1], vals[12 + 2], vals[12 + 3]); - *reinterpret_cast(y + y_last + row_off + 12) = _v4; - } - } - } - if (transpose_m == 0) { - if (e_begin < e_end) { - long long m_first = elem_m[e_begin]; - #pragma unroll - for (int f = 0; f < 16; f++) { - ring[f] = m[m_first + (long long)(f * 128 + k)]; - } - #pragma unroll - for (int f_1 = 0; f_1 < 16; f_1++) { - ring[16 + f_1] = m[m_first + (long long)((16 + f_1) * 128 + k)]; - } - #pragma unroll - for (int f_2 = 0; f_2 < 16; f_2++) { - ring[32 + f_2] = m[m_first + (long long)((32 + f_2) * 128 + k)]; - } - } - __syncthreads(); - #pragma unroll 1 - for (int e = e_begin; e < e_end; e++) { - long long cur = elem_m[e]; - int has_next = 0; - long long nxt = 0; - if (e_end > e + 1) { - has_next = 1; - nxt = elem_m[e + 1]; - } - long long c_off = elem_c[e]; - if (col_tile != 0) { - c_off = -1; - } - if (c_off >= 0) { - { - float4 _v4 = *reinterpret_cast(c + c_off + row_off); - acc[0 + 0] = _v4.x; - acc[0 + 1] = _v4.y; - acc[0 + 2] = _v4.z; - acc[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off + row_off + 4); - acc[4 + 0] = _v4.x; - acc[4 + 1] = _v4.y; - acc[4 + 2] = _v4.z; - acc[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off + row_off + 8); - acc[8 + 0] = _v4.x; - acc[8 + 1] = _v4.y; - acc[8 + 2] = _v4.z; - acc[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off + row_off + 12); - acc[12 + 0] = _v4.x; - acc[12 + 1] = _v4.y; - acc[12 + 2] = _v4.z; - acc[12 + 3] = _v4.w; - } - } else { - acc[0] = 0.0f; - acc[1] = 0.0f; - acc[2] = 0.0f; - acc[3] = 0.0f; - acc[4] = 0.0f; - acc[5] = 0.0f; - acc[6] = 0.0f; - acc[7] = 0.0f; - acc[8] = 0.0f; - acc[9] = 0.0f; - acc[10] = 0.0f; - acc[11] = 0.0f; - acc[12] = 0.0f; - acc[13] = 0.0f; - acc[14] = 0.0f; - acc[15] = 0.0f; - } - #pragma unroll - for (int f_3 = 0; f_3 < 16; f_3++) { - ring[48 + f_3] = m[cur + (long long)((48 + f_3) * 128 + k)]; - } - #pragma unroll - for (int f_4 = 0; f_4 < 16; f_4++) { - float mv = ring[f_4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(f_4 * 16 * 4))); - acc[0] = __fmaf_rn(mv, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((f_4 * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((f_4 * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((f_4 * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv, xv[3], acc[15]); - } - #pragma unroll - for (int f_5 = 0; f_5 < 16; f_5++) { - ring[f_5] = m[cur + (long long)((64 + f_5) * 128 + k)]; - } - #pragma unroll - for (int f_6 = 0; f_6 < 16; f_6++) { - float mv_1 = ring[16 + f_6]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((16 + f_6) * 16 * 4))); - acc[0] = __fmaf_rn(mv_1, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_1, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_1, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_1, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((16 + f_6) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_1, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_1, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_1, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_1, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((16 + f_6) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_1, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_1, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_1, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_1, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((16 + f_6) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_1, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_1, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_1, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_1, xv[3], acc[15]); - } - #pragma unroll - for (int f_7 = 0; f_7 < 16; f_7++) { - ring[16 + f_7] = m[cur + (long long)((80 + f_7) * 128 + k)]; - } - #pragma unroll - for (int f_8 = 0; f_8 < 16; f_8++) { - float mv_2 = ring[32 + f_8]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((32 + f_8) * 16 * 4))); - acc[0] = __fmaf_rn(mv_2, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_2, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_2, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_2, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + f_8) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_2, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_2, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_2, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_2, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + f_8) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_2, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_2, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_2, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_2, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + f_8) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_2, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_2, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_2, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_2, xv[3], acc[15]); - } - #pragma unroll - for (int f_9 = 0; f_9 < 16; f_9++) { - ring[32 + f_9] = m[cur + (long long)((96 + f_9) * 128 + k)]; - } - #pragma unroll - for (int f_10 = 0; f_10 < 16; f_10++) { - float mv_3 = ring[48 + f_10]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((48 + f_10) * 16 * 4))); - acc[0] = __fmaf_rn(mv_3, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_3, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_3, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_3, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((48 + f_10) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_3, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_3, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_3, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_3, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((48 + f_10) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_3, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_3, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_3, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_3, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((48 + f_10) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_3, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_3, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_3, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_3, xv[3], acc[15]); - } - #pragma unroll - for (int f_11 = 0; f_11 < 16; f_11++) { - ring[48 + f_11] = m[cur + (long long)((112 + f_11) * 128 + k)]; - } - #pragma unroll - for (int f_12 = 0; f_12 < 16; f_12++) { - float mv_4 = ring[f_12]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((64 + f_12) * 16 * 4))); - acc[0] = __fmaf_rn(mv_4, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_4, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_4, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_4, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + f_12) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_4, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_4, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_4, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_4, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + f_12) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_4, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_4, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_4, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_4, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + f_12) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_4, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_4, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_4, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_4, xv[3], acc[15]); - } - if (has_next != 0) { - #pragma unroll - for (int f_13 = 0; f_13 < 16; f_13++) { - ring[f_13] = m[nxt + (long long)(f_13 * 128 + k)]; - } - } - #pragma unroll - for (int f_14 = 0; f_14 < 16; f_14++) { - float mv_5 = ring[16 + f_14]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((80 + f_14) * 16 * 4))); - acc[0] = __fmaf_rn(mv_5, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_5, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_5, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_5, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((80 + f_14) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_5, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_5, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_5, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_5, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((80 + f_14) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_5, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_5, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_5, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_5, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((80 + f_14) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_5, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_5, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_5, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_5, xv[3], acc[15]); - } - if (has_next != 0) { - #pragma unroll - for (int f_15 = 0; f_15 < 16; f_15++) { - ring[16 + f_15] = m[nxt + (long long)((16 + f_15) * 128 + k)]; - } - } - #pragma unroll - for (int f_16 = 0; f_16 < 16; f_16++) { - float mv_6 = ring[32 + f_16]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((96 + f_16) * 16 * 4))); - acc[0] = __fmaf_rn(mv_6, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_6, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_6, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_6, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + f_16) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_6, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_6, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_6, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_6, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + f_16) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_6, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_6, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_6, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_6, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + f_16) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_6, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_6, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_6, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_6, xv[3], acc[15]); - } - if (has_next != 0) { - #pragma unroll - for (int f_17 = 0; f_17 < 16; f_17++) { - ring[32 + f_17] = m[nxt + (long long)((32 + f_17) * 128 + k)]; - } - } - #pragma unroll - for (int f_18 = 0; f_18 < 16; f_18++) { - float mv_7 = ring[48 + f_18]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((112 + f_18) * 16 * 4))); - acc[0] = __fmaf_rn(mv_7, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_7, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_7, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_7, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((112 + f_18) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_7, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_7, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_7, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_7, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((112 + f_18) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_7, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_7, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_7, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_7, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((112 + f_18) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_7, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_7, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_7, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_7, xv[3], acc[15]); - } - __syncthreads(); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 16), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 32), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 48), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); - if (has_next != 0) { - if (write_states != 0) { - long long y_next = elem_y[e + 1]; - if (y_next >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_next + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_next + row_off + 4) = _v4; - } - { - float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); - *reinterpret_cast(y + y_next + row_off + 8) = _v4; - } - { - float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); - *reinterpret_cast(y + y_next + row_off + 12) = _v4; - } - } - } - } else { - long long y_out = y_final[chain]; - if (y_out >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_out + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_out + row_off + 4) = _v4; - } - { - float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); - *reinterpret_cast(y + y_out + row_off + 8) = _v4; - } - { - float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); - *reinterpret_cast(y + y_out + row_off + 12) = _v4; - } - } - } - __syncthreads(); - } - } else { - if (e_begin < e_end) { - long long m_first_1 = elem_m[e_begin]; - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)(w4r * 128 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((16 + w4r) * 128 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((32 + w4r) * 128 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((48 + w4r) * 128 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((64 + w4r) * 128 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((80 + w4r) * 128 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((96 + w4r) * 128 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((112 + w4r) * 128 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - slabs[q4 * 129 + w4r] = pf[0]; - slabs[(q4 + 1) * 129 + w4r] = pf[1]; - slabs[(q4 + 2) * 129 + w4r] = pf[2]; - slabs[(q4 + 3) * 129 + w4r] = pf[3]; - slabs[q4 * 129 + 16 + w4r] = pf[4]; - slabs[(q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[(q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[(q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[q4 * 129 + 32 + w4r] = pf[8]; - slabs[(q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[(q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[(q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[q4 * 129 + 48 + w4r] = pf[12]; - slabs[(q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[(q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[(q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[q4 * 129 + 64 + w4r] = pf[16]; - slabs[(q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[(q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[(q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[q4 * 129 + 80 + w4r] = pf[20]; - slabs[(q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[(q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[(q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[q4 * 129 + 96 + w4r] = pf[24]; - slabs[(q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[(q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[(q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[q4 * 129 + 112 + w4r] = pf[28]; - slabs[(q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[(q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[(q4 + 3) * 129 + 112 + w4r] = pf[31]; - } - __syncthreads(); - #pragma unroll 1 - for (int e_1 = e_begin; e_1 < e_end; e_1++) { - long long cur_1 = elem_m[e_1]; - int has_next_1 = 0; - long long nxt_1 = 0; - if (e_end > e_1 + 1) { - has_next_1 = 1; - nxt_1 = elem_m[e_1 + 1]; - } - long long c_off_1 = elem_c[e_1]; - if (col_tile != 0) { - c_off_1 = -1; - } - if (c_off_1 >= 0) { - { - float4 _v4 = *reinterpret_cast(c + c_off_1 + row_off); - acc[0 + 0] = _v4.x; - acc[0 + 1] = _v4.y; - acc[0 + 2] = _v4.z; - acc[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off_1 + row_off + 4); - acc[4 + 0] = _v4.x; - acc[4 + 1] = _v4.y; - acc[4 + 2] = _v4.z; - acc[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off_1 + row_off + 8); - acc[8 + 0] = _v4.x; - acc[8 + 1] = _v4.y; - acc[8 + 2] = _v4.z; - acc[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off_1 + row_off + 12); - acc[12 + 0] = _v4.x; - acc[12 + 1] = _v4.y; - acc[12 + 2] = _v4.z; - acc[12 + 3] = _v4.w; - } - } else { - acc[0] = 0.0f; - acc[1] = 0.0f; - acc[2] = 0.0f; - acc[3] = 0.0f; - acc[4] = 0.0f; - acc[5] = 0.0f; - acc[6] = 0.0f; - acc[7] = 0.0f; - acc[8] = 0.0f; - acc[9] = 0.0f; - acc[10] = 0.0f; - acc[11] = 0.0f; - acc[12] = 0.0f; - acc[13] = 0.0f; - acc[14] = 0.0f; - acc[15] = 0.0f; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)(w4r * 128 + 32 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((16 + w4r) * 128 + 32 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((32 + w4r) * 128 + 32 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((48 + w4r) * 128 + 32 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((64 + w4r) * 128 + 32 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((80 + w4r) * 128 + 32 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((96 + w4r) * 128 + 32 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((112 + w4r) * 128 + 32 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - #pragma unroll - for (int kk = 0; kk < 32; kk++) { - float mv_8 = slabs[kk * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(kk * 16 * 4))); - acc[0] = __fmaf_rn(mv_8, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_8, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_8, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_8, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((kk * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_8, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_8, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_8, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_8, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((kk * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_8, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_8, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_8, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_8, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((kk * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_8, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_8, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_8, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_8, xv[3], acc[15]); - } - slabs[4128 + q4 * 129 + w4r] = pf[0]; - slabs[4128 + (q4 + 1) * 129 + w4r] = pf[1]; - slabs[4128 + (q4 + 2) * 129 + w4r] = pf[2]; - slabs[4128 + (q4 + 3) * 129 + w4r] = pf[3]; - slabs[4128 + q4 * 129 + 16 + w4r] = pf[4]; - slabs[4128 + (q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[4128 + (q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[4128 + (q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[4128 + q4 * 129 + 32 + w4r] = pf[8]; - slabs[4128 + (q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[4128 + (q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[4128 + (q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[4128 + q4 * 129 + 48 + w4r] = pf[12]; - slabs[4128 + (q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[4128 + (q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[4128 + (q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[4128 + q4 * 129 + 64 + w4r] = pf[16]; - slabs[4128 + (q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[4128 + (q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[4128 + (q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[4128 + q4 * 129 + 80 + w4r] = pf[20]; - slabs[4128 + (q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[4128 + (q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[4128 + (q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[4128 + q4 * 129 + 96 + w4r] = pf[24]; - slabs[4128 + (q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[4128 + (q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[4128 + (q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[4128 + q4 * 129 + 112 + w4r] = pf[28]; - slabs[4128 + (q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[4128 + (q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[4128 + (q4 + 3) * 129 + 112 + w4r] = pf[31]; - __syncthreads(); - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)(w4r * 128 + 64 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((16 + w4r) * 128 + 64 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((32 + w4r) * 128 + 64 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((48 + w4r) * 128 + 64 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((64 + w4r) * 128 + 64 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((80 + w4r) * 128 + 64 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((96 + w4r) * 128 + 64 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((112 + w4r) * 128 + 64 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - #pragma unroll - for (int kk_1 = 0; kk_1 < 32; kk_1++) { - float mv_9 = slabs[4128 + kk_1 * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((32 + kk_1) * 16 * 4))); - acc[0] = __fmaf_rn(mv_9, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_9, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_9, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_9, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + kk_1) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_9, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_9, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_9, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_9, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + kk_1) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_9, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_9, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_9, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_9, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + kk_1) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_9, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_9, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_9, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_9, xv[3], acc[15]); - } - slabs[q4 * 129 + w4r] = pf[0]; - slabs[(q4 + 1) * 129 + w4r] = pf[1]; - slabs[(q4 + 2) * 129 + w4r] = pf[2]; - slabs[(q4 + 3) * 129 + w4r] = pf[3]; - slabs[q4 * 129 + 16 + w4r] = pf[4]; - slabs[(q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[(q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[(q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[q4 * 129 + 32 + w4r] = pf[8]; - slabs[(q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[(q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[(q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[q4 * 129 + 48 + w4r] = pf[12]; - slabs[(q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[(q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[(q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[q4 * 129 + 64 + w4r] = pf[16]; - slabs[(q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[(q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[(q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[q4 * 129 + 80 + w4r] = pf[20]; - slabs[(q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[(q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[(q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[q4 * 129 + 96 + w4r] = pf[24]; - slabs[(q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[(q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[(q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[q4 * 129 + 112 + w4r] = pf[28]; - slabs[(q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[(q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[(q4 + 3) * 129 + 112 + w4r] = pf[31]; - __syncthreads(); - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)(w4r * 128 + 96 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((16 + w4r) * 128 + 96 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((32 + w4r) * 128 + 96 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((48 + w4r) * 128 + 96 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((64 + w4r) * 128 + 96 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((80 + w4r) * 128 + 96 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((96 + w4r) * 128 + 96 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((112 + w4r) * 128 + 96 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - #pragma unroll - for (int kk_2 = 0; kk_2 < 32; kk_2++) { - float mv_10 = slabs[kk_2 * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((64 + kk_2) * 16 * 4))); - acc[0] = __fmaf_rn(mv_10, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_10, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_10, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_10, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + kk_2) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_10, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_10, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_10, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_10, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + kk_2) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_10, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_10, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_10, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_10, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + kk_2) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_10, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_10, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_10, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_10, xv[3], acc[15]); - } - slabs[4128 + q4 * 129 + w4r] = pf[0]; - slabs[4128 + (q4 + 1) * 129 + w4r] = pf[1]; - slabs[4128 + (q4 + 2) * 129 + w4r] = pf[2]; - slabs[4128 + (q4 + 3) * 129 + w4r] = pf[3]; - slabs[4128 + q4 * 129 + 16 + w4r] = pf[4]; - slabs[4128 + (q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[4128 + (q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[4128 + (q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[4128 + q4 * 129 + 32 + w4r] = pf[8]; - slabs[4128 + (q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[4128 + (q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[4128 + (q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[4128 + q4 * 129 + 48 + w4r] = pf[12]; - slabs[4128 + (q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[4128 + (q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[4128 + (q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[4128 + q4 * 129 + 64 + w4r] = pf[16]; - slabs[4128 + (q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[4128 + (q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[4128 + (q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[4128 + q4 * 129 + 80 + w4r] = pf[20]; - slabs[4128 + (q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[4128 + (q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[4128 + (q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[4128 + q4 * 129 + 96 + w4r] = pf[24]; - slabs[4128 + (q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[4128 + (q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[4128 + (q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[4128 + q4 * 129 + 112 + w4r] = pf[28]; - slabs[4128 + (q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[4128 + (q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[4128 + (q4 + 3) * 129 + 112 + w4r] = pf[31]; - __syncthreads(); - if (has_next_1 != 0) { - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)(w4r * 128 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((16 + w4r) * 128 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((32 + w4r) * 128 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((48 + w4r) * 128 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((64 + w4r) * 128 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((80 + w4r) * 128 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((96 + w4r) * 128 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((112 + w4r) * 128 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - } - #pragma unroll - for (int kk_3 = 0; kk_3 < 32; kk_3++) { - float mv_11 = slabs[4128 + kk_3 * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((96 + kk_3) * 16 * 4))); - acc[0] = __fmaf_rn(mv_11, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_11, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_11, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_11, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + kk_3) * 16 + 4) * 4))); - acc[4] = __fmaf_rn(mv_11, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_11, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_11, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_11, xv[3], acc[7]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + kk_3) * 16 + 8) * 4))); - acc[8] = __fmaf_rn(mv_11, xv[0], acc[8]); - acc[9] = __fmaf_rn(mv_11, xv[1], acc[9]); - acc[10] = __fmaf_rn(mv_11, xv[2], acc[10]); - acc[11] = __fmaf_rn(mv_11, xv[3], acc[11]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + kk_3) * 16 + 12) * 4))); - acc[12] = __fmaf_rn(mv_11, xv[0], acc[12]); - acc[13] = __fmaf_rn(mv_11, xv[1], acc[13]); - acc[14] = __fmaf_rn(mv_11, xv[2], acc[14]); - acc[15] = __fmaf_rn(mv_11, xv[3], acc[15]); - } - if (has_next_1 != 0) { - slabs[q4 * 129 + w4r] = pf[0]; - slabs[(q4 + 1) * 129 + w4r] = pf[1]; - slabs[(q4 + 2) * 129 + w4r] = pf[2]; - slabs[(q4 + 3) * 129 + w4r] = pf[3]; - slabs[q4 * 129 + 16 + w4r] = pf[4]; - slabs[(q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[(q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[(q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[q4 * 129 + 32 + w4r] = pf[8]; - slabs[(q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[(q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[(q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[q4 * 129 + 48 + w4r] = pf[12]; - slabs[(q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[(q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[(q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[q4 * 129 + 64 + w4r] = pf[16]; - slabs[(q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[(q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[(q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[q4 * 129 + 80 + w4r] = pf[20]; - slabs[(q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[(q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[(q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[q4 * 129 + 96 + w4r] = pf[24]; - slabs[(q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[(q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[(q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[q4 * 129 + 112 + w4r] = pf[28]; - slabs[(q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[(q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[(q4 + 3) * 129 + 112 + w4r] = pf[31]; - } - __syncthreads(); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 16), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 32), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 48), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); - if (has_next_1 != 0) { - if (write_states != 0) { - long long y_next_1 = elem_y[e_1 + 1]; - if (y_next_1 >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_next_1 + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_next_1 + row_off + 4) = _v4; - } - { - float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); - *reinterpret_cast(y + y_next_1 + row_off + 8) = _v4; - } - { - float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); - *reinterpret_cast(y + y_next_1 + row_off + 12) = _v4; - } - } - } - } else { - long long y_out_1 = y_final[chain]; - if (y_out_1 >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_out_1 + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_out_1 + row_off + 4) = _v4; - } - { - float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); - *reinterpret_cast(y + y_out_1 + row_off + 8) = _v4; - } - { - float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); - *reinterpret_cast(y + y_out_1 + row_off + 12) = _v4; - } - } - } - __syncthreads(); - } - } -} - -} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu similarity index 85% rename from flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu rename to flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu index b5e0a27a5e0..c21f4c8ebe2 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_binding.cu +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); +extern "C" __global__ void kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); -namespace cake_binding_kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp { +namespace cake_binding_kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp"; +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -100,8 +100,8 @@ void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, Tensor check_same_device(a_x0_index, a_m, "x0_index", "m"); check_contiguous(a_y_final, "y_final"); check_same_device(a_y_final, a_m, "y_final", "m"); - TVM_FFI_CHECK((a_cols_per_cta == 8LL), ValueError) - << KERNEL_NAME << ": requires (cols_per_cta == 8)"; + TVM_FFI_CHECK((a_cols_per_cta == 32LL), ValueError) + << KERNEL_NAME << ": requires (cols_per_cta == 32)"; TVM_FFI_CHECK((CMod(a_ncols, 128LL) == 0LL), ValueError) << KERNEL_NAME << ": requires (CMod(ncols, 128) == 0)"; TVM_FFI_CHECK((a_ncols >= 128LL), ValueError) @@ -117,6 +117,15 @@ void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, Tensor TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 110592), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } void* v_m = DataPtr(a_m); void* v_c = DataPtr(a_c); void* v_x0 = DataPtr(a_x0); @@ -135,13 +144,13 @@ void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, Tensor cudaLaunchConfig_t config{}; config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); config.blockDim = dim3(128, 1, 1); - config.dynamicSmemBytes = 37120; + config.dynamicSmemBytes = 110592; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp +} // namespace cake_binding_kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu new file mode 100644 index 00000000000..ff9051e768d --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp_kernel.cu @@ -0,0 +1,1854 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LAUNCH_MIN_BLOCKS 2 + +#include +#if __CUDA_ARCH__ == 1070 +#else + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} +#endif + +extern "C" { + +__global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(128) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Prologue + float* xs = reinterpret_cast(smem_raw + 0); + const int xs_addr = smem + 0; + float* ms = reinterpret_cast(smem_raw + 36864); + const int ms_addr = smem + 36864; + + // === Task calls (declaration order) === + int nblk = ncols / 32; + int chain = blockIdx.x / nblk; + int col0 = blockIdx.x % nblk * 32; + int col_tile = col0 / 128; + int v0 = col0 % 128; + int warp_1 = warp; + int g = lane / 4; + int cg = lane % 4; + int e_begin = chain_start[chain]; + int e_end = chain_start[chain + 1]; + unsigned int ring_w = ms_addr + (unsigned int)(warp_1 * 18432); + float acc[32]; + float cn[32]; + float mv[16]; + float xv[8]; + if (transpose_m == 0) { + int row0 = warp_1 * 32 + g * 4; + long long g_off = (long long)col_tile * 16384 + (long long)(row0 * 128 + v0 + cg * 4); + unsigned int x_thr = xs_addr + (unsigned int)((row0 * 36 + cg * 4) * 4); + unsigned int x_rd = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd = ring_w + (unsigned int)(g * 16); + int cp_i = lane / 8; + int cp_j = lane % 8; + int cp_src = cp_i * 128 + warp_1 * 32 + 4 * cp_j; + unsigned int cp_dst = ring_w + (unsigned int)((cp_i * 32 + 4 * cp_j) * 4); + long long m_cur = 0; + if (e_begin < e_end) { + m_cur = elem_m[e_begin]; + long long src0 = m_cur + (long long)cp_src; + unsigned int dst0 = cp_dst + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0), "l"(m + src0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 512), "l"(m + (src0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1024), "l"(m + (src0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1536), "l"(m + (src0 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 2048), "l"(m + (src0 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 2560), "l"(m + (src0 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 3072), "l"(m + (src0 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 3584), "l"(m + (src0 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_0 = m_cur + (long long)(cp_src + 4096); + unsigned int dst0_0 = cp_dst + 4608; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0), "l"(m + src0_0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 512), "l"(m + (src0_0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1024), "l"(m + (src0_0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1536), "l"(m + (src0_0 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 2048), "l"(m + (src0_0 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 2560), "l"(m + (src0_0 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 3072), "l"(m + (src0_0 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 3584), "l"(m + (src0_0 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_1 = m_cur + (long long)(cp_src + 8192); + unsigned int dst0_1 = cp_dst + 9216; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1), "l"(m + src0_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 512), "l"(m + (src0_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1024), "l"(m + (src0_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1536), "l"(m + (src0_1 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 2048), "l"(m + (src0_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 2560), "l"(m + (src0_1 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 3072), "l"(m + (src0_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 3584), "l"(m + (src0_1 + 3584))); + asm volatile("cp.async.commit_group;"); + long long c_first = -1; + if (col_tile == 0) { + c_first = elem_c[e_begin]; + } + if (c_first >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 16); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 128); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 144); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 256); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 272); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 384); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 400); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + } + } + long long x0_off = x0_index[chain]; + if (x0_off >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 16); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 128); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 144); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 256); + acc[16 + 0] = _v4.x; + acc[16 + 1] = _v4.y; + acc[16 + 2] = _v4.z; + acc[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 272); + acc[20 + 0] = _v4.x; + acc[20 + 1] = _v4.y; + acc[20 + 2] = _v4.z; + acc[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 384); + acc[24 + 0] = _v4.x; + acc[24 + 1] = _v4.y; + acc[24 + 2] = _v4.z; + acc[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 400); + acc[28 + 0] = _v4.x; + acc[28 + 1] = _v4.y; + acc[28 + 2] = _v4.z; + acc[28 + 3] = _v4.w; + } + } else { + int diag = row0 - v0 - cg * 4; + if (col_tile == 0) { + diag = -4096; + } + acc[0] = ((diag == 0) ? 1.0f : 0.0f); + acc[1] = ((diag == 1) ? 1.0f : 0.0f); + acc[2] = ((diag == 2) ? 1.0f : 0.0f); + acc[3] = ((diag == 3) ? 1.0f : 0.0f); + acc[4] = ((diag == 16) ? 1.0f : 0.0f); + acc[5] = ((diag == 17) ? 1.0f : 0.0f); + acc[6] = ((diag == 18) ? 1.0f : 0.0f); + acc[7] = ((diag == 19) ? 1.0f : 0.0f); + acc[8] = ((diag + 1 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag + 1 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag + 1 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag + 1 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag + 1 == 16) ? 1.0f : 0.0f); + acc[13] = ((diag + 1 == 17) ? 1.0f : 0.0f); + acc[14] = ((diag + 1 == 18) ? 1.0f : 0.0f); + acc[15] = ((diag + 1 == 19) ? 1.0f : 0.0f); + acc[16] = ((diag + 2 == 0) ? 1.0f : 0.0f); + acc[17] = ((diag + 2 == 1) ? 1.0f : 0.0f); + acc[18] = ((diag + 2 == 2) ? 1.0f : 0.0f); + acc[19] = ((diag + 2 == 3) ? 1.0f : 0.0f); + acc[20] = ((diag + 2 == 16) ? 1.0f : 0.0f); + acc[21] = ((diag + 2 == 17) ? 1.0f : 0.0f); + acc[22] = ((diag + 2 == 18) ? 1.0f : 0.0f); + acc[23] = ((diag + 2 == 19) ? 1.0f : 0.0f); + acc[24] = ((diag + 3 == 0) ? 1.0f : 0.0f); + acc[25] = ((diag + 3 == 1) ? 1.0f : 0.0f); + acc[26] = ((diag + 3 == 2) ? 1.0f : 0.0f); + acc[27] = ((diag + 3 == 3) ? 1.0f : 0.0f); + acc[28] = ((diag + 3 == 16) ? 1.0f : 0.0f); + acc[29] = ((diag + 3 == 17) ? 1.0f : 0.0f); + acc[30] = ((diag + 3 == 18) ? 1.0f : 0.0f); + acc[31] = ((diag + 3 == 19) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 64), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 144), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 208), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 288), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 352), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 432), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 496), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first = elem_y[e_begin]; + if (y_first >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first + g_off + 16) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first + g_off + 144) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_first + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_first + g_off + 272) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_first + g_off + 384) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_first + g_off + 400) = _v4; + } + } + } + } else { + long long y_last = y_final[chain]; + if (y_last >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last + g_off + 16) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last + g_off + 144) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_last + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_last + g_off + 272) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_last + g_off + 384) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_last + g_off + 400) = _v4; + } + } + } + __syncthreads(); + unsigned int buf = 0; + #pragma unroll 1 + for (int e = e_begin; e < e_end; e++) { + int has_next = 0; + long long m_next = 0; + long long c_next = -1; + long long y_pub = -1; + if (e + 1 < e_end) { + has_next = 1; + m_next = elem_m[e + 1]; + if (col_tile == 0) { + c_next = elem_c[e + 1]; + } + if (write_states != 0) { + y_pub = elem_y[e + 1]; + } + } else { + y_pub = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + acc[16] = cn[16]; + acc[17] = cn[17]; + acc[18] = cn[18]; + acc[19] = cn[19]; + acc[20] = cn[20]; + acc[21] = cn[21]; + acc[22] = cn[22]; + acc[23] = cn[23]; + acc[24] = cn[24]; + acc[25] = cn[25]; + acc[26] = cn[26]; + acc[27] = cn[27]; + acc[28] = cn[28]; + acc[29] = cn[29]; + acc[30] = cn[30]; + acc[31] = cn[31]; + unsigned int x_rd_e = x_rd + buf * 18432; + #pragma unroll 1 + for (int ch = 0; ch < 4; ch++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc = ch + 3; + if (nc < 4) { + long long src0_2 = m_cur + (long long)(cp_src + nc * 4096); + unsigned int dst0_2 = cp_dst + (unsigned int)(nc % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2), "l"(m + src0_2)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 512), "l"(m + (src0_2 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1024), "l"(m + (src0_2 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1536), "l"(m + (src0_2 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 2048), "l"(m + (src0_2 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 2560), "l"(m + (src0_2 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 3072), "l"(m + (src0_2 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 3584), "l"(m + (src0_2 + 3584))); + } else if (has_next != 0) { + long long src0_3 = m_next + (long long)(cp_src + (nc - 4) * 4096); + unsigned int dst0_3 = cp_dst + (unsigned int)(nc % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3), "l"(m + src0_3)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 512), "l"(m + (src0_3 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1024), "l"(m + (src0_3 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1536), "l"(m + (src0_3 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 2048), "l"(m + (src0_3 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 2560), "l"(m + (src0_3 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 3072), "l"(m + (src0_3 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 3584), "l"(m + (src0_3 + 3584))); + } + asm volatile("cp.async.commit_group;"); + if (ch == 1) { + if (c_next >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 16); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 128); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 144); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 256); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 272); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 384); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 400); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + } + } + #pragma unroll + for (int kk = 0; kk < 32; kk++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd + (unsigned int)(ch % 4 * 4608) + (unsigned int)(kk * 128))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e + (unsigned int)(ch * 4608) + (unsigned int)(kk * 144))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e + (unsigned int)(ch * 4608) + (unsigned int)(kk * 144) + 64)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[0], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[0], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[0], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[0], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[1], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[1], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[1], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[1], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[1], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[1], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[1], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[1], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[2], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[2], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[2], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[2], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[2], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[2], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[2], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[2], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[3], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[3], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[3], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[3], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[3], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[3], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[3], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[3], xv[7], acc[31]); +#else + float2 m2_0 = make_float2(mv[0], mv[0]); + float2 a2_0 = make_float2(acc[0], acc[1]); + a2_0 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[0], xv[1]), a2_0); + acc[0] = a2_0.x; + acc[1] = a2_0.y; + float2 a2_1 = make_float2(acc[2], acc[3]); + a2_1 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[2], xv[3]), a2_1); + acc[2] = a2_1.x; + acc[3] = a2_1.y; + float2 a2_2 = make_float2(acc[4], acc[5]); + a2_2 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[4], xv[5]), a2_2); + acc[4] = a2_2.x; + acc[5] = a2_2.y; + float2 a2_3 = make_float2(acc[6], acc[7]); + a2_3 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[6], xv[7]), a2_3); + acc[6] = a2_3.x; + acc[7] = a2_3.y; + float2 m2_1 = make_float2(mv[1], mv[1]); + float2 a2_4 = make_float2(acc[8], acc[9]); + a2_4 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[0], xv[1]), a2_4); + acc[8] = a2_4.x; + acc[9] = a2_4.y; + float2 a2_5 = make_float2(acc[10], acc[11]); + a2_5 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[2], xv[3]), a2_5); + acc[10] = a2_5.x; + acc[11] = a2_5.y; + float2 a2_6 = make_float2(acc[12], acc[13]); + a2_6 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[4], xv[5]), a2_6); + acc[12] = a2_6.x; + acc[13] = a2_6.y; + float2 a2_7 = make_float2(acc[14], acc[15]); + a2_7 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[6], xv[7]), a2_7); + acc[14] = a2_7.x; + acc[15] = a2_7.y; + float2 m2_2 = make_float2(mv[2], mv[2]); + float2 a2_8 = make_float2(acc[16], acc[17]); + a2_8 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[0], xv[1]), a2_8); + acc[16] = a2_8.x; + acc[17] = a2_8.y; + float2 a2_9 = make_float2(acc[18], acc[19]); + a2_9 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[2], xv[3]), a2_9); + acc[18] = a2_9.x; + acc[19] = a2_9.y; + float2 a2_10 = make_float2(acc[20], acc[21]); + a2_10 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[4], xv[5]), a2_10); + acc[20] = a2_10.x; + acc[21] = a2_10.y; + float2 a2_11 = make_float2(acc[22], acc[23]); + a2_11 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[6], xv[7]), a2_11); + acc[22] = a2_11.x; + acc[23] = a2_11.y; + float2 m2_3 = make_float2(mv[3], mv[3]); + float2 a2_12 = make_float2(acc[24], acc[25]); + a2_12 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[0], xv[1]), a2_12); + acc[24] = a2_12.x; + acc[25] = a2_12.y; + float2 a2_13 = make_float2(acc[26], acc[27]); + a2_13 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[2], xv[3]), a2_13); + acc[26] = a2_13.x; + acc[27] = a2_13.y; + float2 a2_14 = make_float2(acc[28], acc[29]); + a2_14 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[4], xv[5]), a2_14); + acc[28] = a2_14.x; + acc[29] = a2_14.y; + float2 a2_15 = make_float2(acc[30], acc[31]); + a2_15 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[6], xv[7]), a2_15); + acc[30] = a2_15.x; + acc[31] = a2_15.y; +#endif + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432) + 64), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432) + 144), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432) + 208), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432) + 288), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432) + 352), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432) + 432), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (18432 - buf * 18432) + 496), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + if (y_pub >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub + g_off + 16) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub + g_off + 144) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_pub + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_pub + g_off + 272) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_pub + g_off + 384) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_pub + g_off + 400) = _v4; + } + } + __syncthreads(); + buf = 1 - buf; + m_cur = m_next; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } else { + int row0_1 = warp_1 * 32 + g; + long long g_off_1 = (long long)col_tile * 16384 + (long long)(row0_1 * 128 + v0 + cg * 4); + unsigned int x_thr_1 = xs_addr + (unsigned int)((row0_1 * 36 + cg * 4) * 4); + unsigned int x_rd_1 = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd_1 = ring_w + (unsigned int)(g * 144); + int cp_i_1 = lane / 8; + int cp_j_1 = lane % 8; + int cp_src_1 = cp_i_1 * 128 + warp_1 * 4096 + 4 * cp_j_1; + unsigned int cp_dst_1 = ring_w + (unsigned int)((cp_i_1 * 36 + 4 * cp_j_1) * 4); + long long m_cur_1 = 0; + if (e_begin < e_end) { + m_cur_1 = elem_m[e_begin]; + long long src0_4 = m_cur_1 + (long long)cp_src_1; + unsigned int dst0_4 = cp_dst_1 + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4), "l"(m + src0_4)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 576), "l"(m + (src0_4 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1152), "l"(m + (src0_4 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1728), "l"(m + (src0_4 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 2304), "l"(m + (src0_4 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 2880), "l"(m + (src0_4 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 3456), "l"(m + (src0_4 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 4032), "l"(m + (src0_4 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_2_1 = m_cur_1 + (long long)(cp_src_1 + 32); + unsigned int dst0_2_1 = cp_dst_1 + 4608; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1), "l"(m + src0_2_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 576), "l"(m + (src0_2_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1152), "l"(m + (src0_2_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1728), "l"(m + (src0_2_1 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 2304), "l"(m + (src0_2_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 2880), "l"(m + (src0_2_1 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 3456), "l"(m + (src0_2_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 4032), "l"(m + (src0_2_1 + 3584))); + asm volatile("cp.async.commit_group;"); + long long src0_3_1 = m_cur_1 + (long long)(cp_src_1 + 64); + unsigned int dst0_3_1 = cp_dst_1 + 9216; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1), "l"(m + src0_3_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 576), "l"(m + (src0_3_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1152), "l"(m + (src0_3_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1728), "l"(m + (src0_3_1 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 2304), "l"(m + (src0_3_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 2880), "l"(m + (src0_3_1 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 3456), "l"(m + (src0_3_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 4032), "l"(m + (src0_3_1 + 3584))); + asm volatile("cp.async.commit_group;"); + long long c_first_1 = -1; + if (col_tile == 0) { + c_first_1 = elem_c[e_begin]; + } + if (c_first_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 16); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 1024); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 1040); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2048); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2064); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 3072); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 3088); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + } + } + long long x0_off_1 = x0_index[chain]; + if (x0_off_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 16); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 1024); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 1040); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2048); + acc[16 + 0] = _v4.x; + acc[16 + 1] = _v4.y; + acc[16 + 2] = _v4.z; + acc[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2064); + acc[20 + 0] = _v4.x; + acc[20 + 1] = _v4.y; + acc[20 + 2] = _v4.z; + acc[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 3072); + acc[24 + 0] = _v4.x; + acc[24 + 1] = _v4.y; + acc[24 + 2] = _v4.z; + acc[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 3088); + acc[28 + 0] = _v4.x; + acc[28 + 1] = _v4.y; + acc[28 + 2] = _v4.z; + acc[28 + 3] = _v4.w; + } + } else { + int diag_1 = row0_1 - v0 - cg * 4; + if (col_tile == 0) { + diag_1 = -4096; + } + acc[0] = ((diag_1 == 0) ? 1.0f : 0.0f); + acc[1] = ((diag_1 == 1) ? 1.0f : 0.0f); + acc[2] = ((diag_1 == 2) ? 1.0f : 0.0f); + acc[3] = ((diag_1 == 3) ? 1.0f : 0.0f); + acc[4] = ((diag_1 == 16) ? 1.0f : 0.0f); + acc[5] = ((diag_1 == 17) ? 1.0f : 0.0f); + acc[6] = ((diag_1 == 18) ? 1.0f : 0.0f); + acc[7] = ((diag_1 == 19) ? 1.0f : 0.0f); + acc[8] = ((diag_1 + 8 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag_1 + 8 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag_1 + 8 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag_1 + 8 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag_1 + 8 == 16) ? 1.0f : 0.0f); + acc[13] = ((diag_1 + 8 == 17) ? 1.0f : 0.0f); + acc[14] = ((diag_1 + 8 == 18) ? 1.0f : 0.0f); + acc[15] = ((diag_1 + 8 == 19) ? 1.0f : 0.0f); + acc[16] = ((diag_1 + 16 == 0) ? 1.0f : 0.0f); + acc[17] = ((diag_1 + 16 == 1) ? 1.0f : 0.0f); + acc[18] = ((diag_1 + 16 == 2) ? 1.0f : 0.0f); + acc[19] = ((diag_1 + 16 == 3) ? 1.0f : 0.0f); + acc[20] = ((diag_1 + 16 == 16) ? 1.0f : 0.0f); + acc[21] = ((diag_1 + 16 == 17) ? 1.0f : 0.0f); + acc[22] = ((diag_1 + 16 == 18) ? 1.0f : 0.0f); + acc[23] = ((diag_1 + 16 == 19) ? 1.0f : 0.0f); + acc[24] = ((diag_1 + 24 == 0) ? 1.0f : 0.0f); + acc[25] = ((diag_1 + 24 == 1) ? 1.0f : 0.0f); + acc[26] = ((diag_1 + 24 == 2) ? 1.0f : 0.0f); + acc[27] = ((diag_1 + 24 == 3) ? 1.0f : 0.0f); + acc[28] = ((diag_1 + 24 == 16) ? 1.0f : 0.0f); + acc[29] = ((diag_1 + 24 == 17) ? 1.0f : 0.0f); + acc[30] = ((diag_1 + 24 == 18) ? 1.0f : 0.0f); + acc[31] = ((diag_1 + 24 == 19) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 64), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 1152), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 1216), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 2304), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 2368), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 3456), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 3520), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first_1 = elem_y[e_begin]; + if (y_first_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 16) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 1040) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2064) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 3072) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 3088) = _v4; + } + } + } + } else { + long long y_last_1 = y_final[chain]; + if (y_last_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 16) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 1040) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2064) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 3072) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 3088) = _v4; + } + } + } + __syncthreads(); + unsigned int buf_1 = 0; + #pragma unroll 1 + for (int e_1 = e_begin; e_1 < e_end; e_1++) { + int has_next_1 = 0; + long long m_next_1 = 0; + long long c_next_1 = -1; + long long y_pub_1 = -1; + if (e_1 + 1 < e_end) { + has_next_1 = 1; + m_next_1 = elem_m[e_1 + 1]; + if (col_tile == 0) { + c_next_1 = elem_c[e_1 + 1]; + } + if (write_states != 0) { + y_pub_1 = elem_y[e_1 + 1]; + } + } else { + y_pub_1 = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + acc[16] = cn[16]; + acc[17] = cn[17]; + acc[18] = cn[18]; + acc[19] = cn[19]; + acc[20] = cn[20]; + acc[21] = cn[21]; + acc[22] = cn[22]; + acc[23] = cn[23]; + acc[24] = cn[24]; + acc[25] = cn[25]; + acc[26] = cn[26]; + acc[27] = cn[27]; + acc[28] = cn[28]; + acc[29] = cn[29]; + acc[30] = cn[30]; + acc[31] = cn[31]; + unsigned int x_rd_e_1 = x_rd_1 + buf_1 * 18432; + #pragma unroll 1 + for (int ch_1 = 0; ch_1 < 4; ch_1++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc_1 = ch_1 + 3; + if (nc_1 < 4) { + long long src0_5 = m_cur_1 + (long long)(cp_src_1 + nc_1 * 32); + unsigned int dst0_5 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5), "l"(m + src0_5)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 576), "l"(m + (src0_5 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1152), "l"(m + (src0_5 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1728), "l"(m + (src0_5 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 2304), "l"(m + (src0_5 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 2880), "l"(m + (src0_5 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 3456), "l"(m + (src0_5 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 4032), "l"(m + (src0_5 + 3584))); + } else if (has_next_1 != 0) { + long long src0_6 = m_next_1 + (long long)(cp_src_1 + (nc_1 - 4) * 32); + unsigned int dst0_6 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 4608); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6), "l"(m + src0_6)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 576), "l"(m + (src0_6 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1152), "l"(m + (src0_6 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1728), "l"(m + (src0_6 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 2304), "l"(m + (src0_6 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 2880), "l"(m + (src0_6 + 2560))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 3456), "l"(m + (src0_6 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 4032), "l"(m + (src0_6 + 3584))); + } + asm volatile("cp.async.commit_group;"); + if (ch_1 == 1) { + if (c_next_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 16); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 1024); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 1040); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2048); + cn[16 + 0] = _v4.x; + cn[16 + 1] = _v4.y; + cn[16 + 2] = _v4.z; + cn[16 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2064); + cn[20 + 0] = _v4.x; + cn[20 + 1] = _v4.y; + cn[20 + 2] = _v4.z; + cn[20 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 3072); + cn[24 + 0] = _v4.x; + cn[24 + 1] = _v4.y; + cn[24 + 2] = _v4.z; + cn[24 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 3088); + cn[28 + 0] = _v4.x; + cn[28 + 1] = _v4.y; + cn[28 + 2] = _v4.z; + cn[28 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + cn[16] = 0.0f; + cn[17] = 0.0f; + cn[18] = 0.0f; + cn[19] = 0.0f; + cn[20] = 0.0f; + cn[21] = 0.0f; + cn[22] = 0.0f; + cn[23] = 0.0f; + cn[24] = 0.0f; + cn[25] = 0.0f; + cn[26] = 0.0f; + cn[27] = 0.0f; + cn[28] = 0.0f; + cn[29] = 0.0f; + cn[30] = 0.0f; + cn[31] = 0.0f; + } + } + #pragma unroll + for (int q = 0; q < 8; q++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[4])), "=r"(*reinterpret_cast(&mv[(4) + 1])), "=r"(*reinterpret_cast(&mv[(4) + 2])), "=r"(*reinterpret_cast(&mv[(4) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 1152)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[8])), "=r"(*reinterpret_cast(&mv[(8) + 1])), "=r"(*reinterpret_cast(&mv[(8) + 2])), "=r"(*reinterpret_cast(&mv[(8) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 2304)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[12])), "=r"(*reinterpret_cast(&mv[(12) + 1])), "=r"(*reinterpret_cast(&mv[(12) + 2])), "=r"(*reinterpret_cast(&mv[(12) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 4608) + (unsigned int)(q * 16) + 3456)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576) + 64)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[0], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[0], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[0], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[0], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[4], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[4], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[4], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[4], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[4], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[4], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[4], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[4], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[8], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[8], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[8], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[8], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[8], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[8], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[8], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[8], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[12], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[12], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[12], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[12], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[12], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[12], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[12], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[12], xv[7], acc[31]); +#else + float2 m2_4 = make_float2(mv[0], mv[0]); + float2 a2_16 = make_float2(acc[0], acc[1]); + a2_16 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[0], xv[1]), a2_16); + acc[0] = a2_16.x; + acc[1] = a2_16.y; + float2 a2_17 = make_float2(acc[2], acc[3]); + a2_17 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[2], xv[3]), a2_17); + acc[2] = a2_17.x; + acc[3] = a2_17.y; + float2 a2_18 = make_float2(acc[4], acc[5]); + a2_18 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[4], xv[5]), a2_18); + acc[4] = a2_18.x; + acc[5] = a2_18.y; + float2 a2_19 = make_float2(acc[6], acc[7]); + a2_19 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[6], xv[7]), a2_19); + acc[6] = a2_19.x; + acc[7] = a2_19.y; + float2 m2_5 = make_float2(mv[4], mv[4]); + float2 a2_20 = make_float2(acc[8], acc[9]); + a2_20 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[0], xv[1]), a2_20); + acc[8] = a2_20.x; + acc[9] = a2_20.y; + float2 a2_21 = make_float2(acc[10], acc[11]); + a2_21 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[2], xv[3]), a2_21); + acc[10] = a2_21.x; + acc[11] = a2_21.y; + float2 a2_22 = make_float2(acc[12], acc[13]); + a2_22 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[4], xv[5]), a2_22); + acc[12] = a2_22.x; + acc[13] = a2_22.y; + float2 a2_23 = make_float2(acc[14], acc[15]); + a2_23 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[6], xv[7]), a2_23); + acc[14] = a2_23.x; + acc[15] = a2_23.y; + float2 m2_6 = make_float2(mv[8], mv[8]); + float2 a2_24 = make_float2(acc[16], acc[17]); + a2_24 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[0], xv[1]), a2_24); + acc[16] = a2_24.x; + acc[17] = a2_24.y; + float2 a2_25 = make_float2(acc[18], acc[19]); + a2_25 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[2], xv[3]), a2_25); + acc[18] = a2_25.x; + acc[19] = a2_25.y; + float2 a2_26 = make_float2(acc[20], acc[21]); + a2_26 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[4], xv[5]), a2_26); + acc[20] = a2_26.x; + acc[21] = a2_26.y; + float2 a2_27 = make_float2(acc[22], acc[23]); + a2_27 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[6], xv[7]), a2_27); + acc[22] = a2_27.x; + acc[23] = a2_27.y; + float2 m2_7 = make_float2(mv[12], mv[12]); + float2 a2_28 = make_float2(acc[24], acc[25]); + a2_28 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[0], xv[1]), a2_28); + acc[24] = a2_28.x; + acc[25] = a2_28.y; + float2 a2_29 = make_float2(acc[26], acc[27]); + a2_29 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[2], xv[3]), a2_29); + acc[26] = a2_29.x; + acc[27] = a2_29.y; + float2 a2_30 = make_float2(acc[28], acc[29]); + a2_30 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[4], xv[5]), a2_30); + acc[28] = a2_30.x; + acc[29] = a2_30.y; + float2 a2_31 = make_float2(acc[30], acc[31]); + a2_31 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[6], xv[7]), a2_31); + acc[30] = a2_31.x; + acc[31] = a2_31.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576) + 144)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576) + 208)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[1], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[1], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[1], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[1], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[1], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[1], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[1], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[1], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[5], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[5], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[5], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[5], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[5], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[5], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[5], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[5], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[9], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[9], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[9], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[9], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[9], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[9], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[9], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[9], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[13], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[13], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[13], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[13], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[13], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[13], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[13], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[13], xv[7], acc[31]); +#else + float2 m2_8 = make_float2(mv[1], mv[1]); + float2 a2_32 = make_float2(acc[0], acc[1]); + a2_32 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[0], xv[1]), a2_32); + acc[0] = a2_32.x; + acc[1] = a2_32.y; + float2 a2_33 = make_float2(acc[2], acc[3]); + a2_33 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[2], xv[3]), a2_33); + acc[2] = a2_33.x; + acc[3] = a2_33.y; + float2 a2_34 = make_float2(acc[4], acc[5]); + a2_34 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[4], xv[5]), a2_34); + acc[4] = a2_34.x; + acc[5] = a2_34.y; + float2 a2_35 = make_float2(acc[6], acc[7]); + a2_35 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[6], xv[7]), a2_35); + acc[6] = a2_35.x; + acc[7] = a2_35.y; + float2 m2_9 = make_float2(mv[5], mv[5]); + float2 a2_36 = make_float2(acc[8], acc[9]); + a2_36 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[0], xv[1]), a2_36); + acc[8] = a2_36.x; + acc[9] = a2_36.y; + float2 a2_37 = make_float2(acc[10], acc[11]); + a2_37 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[2], xv[3]), a2_37); + acc[10] = a2_37.x; + acc[11] = a2_37.y; + float2 a2_38 = make_float2(acc[12], acc[13]); + a2_38 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[4], xv[5]), a2_38); + acc[12] = a2_38.x; + acc[13] = a2_38.y; + float2 a2_39 = make_float2(acc[14], acc[15]); + a2_39 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[6], xv[7]), a2_39); + acc[14] = a2_39.x; + acc[15] = a2_39.y; + float2 m2_10 = make_float2(mv[9], mv[9]); + float2 a2_40 = make_float2(acc[16], acc[17]); + a2_40 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[0], xv[1]), a2_40); + acc[16] = a2_40.x; + acc[17] = a2_40.y; + float2 a2_41 = make_float2(acc[18], acc[19]); + a2_41 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[2], xv[3]), a2_41); + acc[18] = a2_41.x; + acc[19] = a2_41.y; + float2 a2_42 = make_float2(acc[20], acc[21]); + a2_42 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[4], xv[5]), a2_42); + acc[20] = a2_42.x; + acc[21] = a2_42.y; + float2 a2_43 = make_float2(acc[22], acc[23]); + a2_43 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[6], xv[7]), a2_43); + acc[22] = a2_43.x; + acc[23] = a2_43.y; + float2 m2_11 = make_float2(mv[13], mv[13]); + float2 a2_44 = make_float2(acc[24], acc[25]); + a2_44 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[0], xv[1]), a2_44); + acc[24] = a2_44.x; + acc[25] = a2_44.y; + float2 a2_45 = make_float2(acc[26], acc[27]); + a2_45 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[2], xv[3]), a2_45); + acc[26] = a2_45.x; + acc[27] = a2_45.y; + float2 a2_46 = make_float2(acc[28], acc[29]); + a2_46 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[4], xv[5]), a2_46); + acc[28] = a2_46.x; + acc[29] = a2_46.y; + float2 a2_47 = make_float2(acc[30], acc[31]); + a2_47 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[6], xv[7]), a2_47); + acc[30] = a2_47.x; + acc[31] = a2_47.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576) + 288)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576) + 352)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[2], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[2], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[2], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[2], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[2], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[2], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[2], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[2], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[6], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[6], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[6], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[6], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[6], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[6], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[6], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[6], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[10], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[10], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[10], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[10], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[10], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[10], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[10], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[10], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[14], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[14], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[14], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[14], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[14], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[14], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[14], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[14], xv[7], acc[31]); +#else + float2 m2_12 = make_float2(mv[2], mv[2]); + float2 a2_48 = make_float2(acc[0], acc[1]); + a2_48 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[0], xv[1]), a2_48); + acc[0] = a2_48.x; + acc[1] = a2_48.y; + float2 a2_49 = make_float2(acc[2], acc[3]); + a2_49 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[2], xv[3]), a2_49); + acc[2] = a2_49.x; + acc[3] = a2_49.y; + float2 a2_50 = make_float2(acc[4], acc[5]); + a2_50 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[4], xv[5]), a2_50); + acc[4] = a2_50.x; + acc[5] = a2_50.y; + float2 a2_51 = make_float2(acc[6], acc[7]); + a2_51 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[6], xv[7]), a2_51); + acc[6] = a2_51.x; + acc[7] = a2_51.y; + float2 m2_13 = make_float2(mv[6], mv[6]); + float2 a2_52 = make_float2(acc[8], acc[9]); + a2_52 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[0], xv[1]), a2_52); + acc[8] = a2_52.x; + acc[9] = a2_52.y; + float2 a2_53 = make_float2(acc[10], acc[11]); + a2_53 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[2], xv[3]), a2_53); + acc[10] = a2_53.x; + acc[11] = a2_53.y; + float2 a2_54 = make_float2(acc[12], acc[13]); + a2_54 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[4], xv[5]), a2_54); + acc[12] = a2_54.x; + acc[13] = a2_54.y; + float2 a2_55 = make_float2(acc[14], acc[15]); + a2_55 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[6], xv[7]), a2_55); + acc[14] = a2_55.x; + acc[15] = a2_55.y; + float2 m2_14 = make_float2(mv[10], mv[10]); + float2 a2_56 = make_float2(acc[16], acc[17]); + a2_56 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[0], xv[1]), a2_56); + acc[16] = a2_56.x; + acc[17] = a2_56.y; + float2 a2_57 = make_float2(acc[18], acc[19]); + a2_57 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[2], xv[3]), a2_57); + acc[18] = a2_57.x; + acc[19] = a2_57.y; + float2 a2_58 = make_float2(acc[20], acc[21]); + a2_58 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[4], xv[5]), a2_58); + acc[20] = a2_58.x; + acc[21] = a2_58.y; + float2 a2_59 = make_float2(acc[22], acc[23]); + a2_59 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[6], xv[7]), a2_59); + acc[22] = a2_59.x; + acc[23] = a2_59.y; + float2 m2_15 = make_float2(mv[14], mv[14]); + float2 a2_60 = make_float2(acc[24], acc[25]); + a2_60 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[0], xv[1]), a2_60); + acc[24] = a2_60.x; + acc[25] = a2_60.y; + float2 a2_61 = make_float2(acc[26], acc[27]); + a2_61 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[2], xv[3]), a2_61); + acc[26] = a2_61.x; + acc[27] = a2_61.y; + float2 a2_62 = make_float2(acc[28], acc[29]); + a2_62 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[4], xv[5]), a2_62); + acc[28] = a2_62.x; + acc[29] = a2_62.y; + float2 a2_63 = make_float2(acc[30], acc[31]); + a2_63 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[6], xv[7]), a2_63); + acc[30] = a2_63.x; + acc[31] = a2_63.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576) + 432)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 4608) + (unsigned int)(q * 576) + 496)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[3], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[3], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[3], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[3], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[3], xv[4], acc[4]); + acc[5] = __fmaf_rn(mv[3], xv[5], acc[5]); + acc[6] = __fmaf_rn(mv[3], xv[6], acc[6]); + acc[7] = __fmaf_rn(mv[3], xv[7], acc[7]); + acc[8] = __fmaf_rn(mv[7], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[7], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[7], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[7], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[7], xv[4], acc[12]); + acc[13] = __fmaf_rn(mv[7], xv[5], acc[13]); + acc[14] = __fmaf_rn(mv[7], xv[6], acc[14]); + acc[15] = __fmaf_rn(mv[7], xv[7], acc[15]); + acc[16] = __fmaf_rn(mv[11], xv[0], acc[16]); + acc[17] = __fmaf_rn(mv[11], xv[1], acc[17]); + acc[18] = __fmaf_rn(mv[11], xv[2], acc[18]); + acc[19] = __fmaf_rn(mv[11], xv[3], acc[19]); + acc[20] = __fmaf_rn(mv[11], xv[4], acc[20]); + acc[21] = __fmaf_rn(mv[11], xv[5], acc[21]); + acc[22] = __fmaf_rn(mv[11], xv[6], acc[22]); + acc[23] = __fmaf_rn(mv[11], xv[7], acc[23]); + acc[24] = __fmaf_rn(mv[15], xv[0], acc[24]); + acc[25] = __fmaf_rn(mv[15], xv[1], acc[25]); + acc[26] = __fmaf_rn(mv[15], xv[2], acc[26]); + acc[27] = __fmaf_rn(mv[15], xv[3], acc[27]); + acc[28] = __fmaf_rn(mv[15], xv[4], acc[28]); + acc[29] = __fmaf_rn(mv[15], xv[5], acc[29]); + acc[30] = __fmaf_rn(mv[15], xv[6], acc[30]); + acc[31] = __fmaf_rn(mv[15], xv[7], acc[31]); +#else + float2 m2_16 = make_float2(mv[3], mv[3]); + float2 a2_64 = make_float2(acc[0], acc[1]); + a2_64 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[0], xv[1]), a2_64); + acc[0] = a2_64.x; + acc[1] = a2_64.y; + float2 a2_65 = make_float2(acc[2], acc[3]); + a2_65 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[2], xv[3]), a2_65); + acc[2] = a2_65.x; + acc[3] = a2_65.y; + float2 a2_66 = make_float2(acc[4], acc[5]); + a2_66 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[4], xv[5]), a2_66); + acc[4] = a2_66.x; + acc[5] = a2_66.y; + float2 a2_67 = make_float2(acc[6], acc[7]); + a2_67 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[6], xv[7]), a2_67); + acc[6] = a2_67.x; + acc[7] = a2_67.y; + float2 m2_17 = make_float2(mv[7], mv[7]); + float2 a2_68 = make_float2(acc[8], acc[9]); + a2_68 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[0], xv[1]), a2_68); + acc[8] = a2_68.x; + acc[9] = a2_68.y; + float2 a2_69 = make_float2(acc[10], acc[11]); + a2_69 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[2], xv[3]), a2_69); + acc[10] = a2_69.x; + acc[11] = a2_69.y; + float2 a2_70 = make_float2(acc[12], acc[13]); + a2_70 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[4], xv[5]), a2_70); + acc[12] = a2_70.x; + acc[13] = a2_70.y; + float2 a2_71 = make_float2(acc[14], acc[15]); + a2_71 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[6], xv[7]), a2_71); + acc[14] = a2_71.x; + acc[15] = a2_71.y; + float2 m2_18 = make_float2(mv[11], mv[11]); + float2 a2_72 = make_float2(acc[16], acc[17]); + a2_72 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[0], xv[1]), a2_72); + acc[16] = a2_72.x; + acc[17] = a2_72.y; + float2 a2_73 = make_float2(acc[18], acc[19]); + a2_73 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[2], xv[3]), a2_73); + acc[18] = a2_73.x; + acc[19] = a2_73.y; + float2 a2_74 = make_float2(acc[20], acc[21]); + a2_74 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[4], xv[5]), a2_74); + acc[20] = a2_74.x; + acc[21] = a2_74.y; + float2 a2_75 = make_float2(acc[22], acc[23]); + a2_75 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[6], xv[7]), a2_75); + acc[22] = a2_75.x; + acc[23] = a2_75.y; + float2 m2_19 = make_float2(mv[15], mv[15]); + float2 a2_76 = make_float2(acc[24], acc[25]); + a2_76 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[0], xv[1]), a2_76); + acc[24] = a2_76.x; + acc[25] = a2_76.y; + float2 a2_77 = make_float2(acc[26], acc[27]); + a2_77 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[2], xv[3]), a2_77); + acc[26] = a2_77.x; + acc[27] = a2_77.y; + float2 a2_78 = make_float2(acc[28], acc[29]); + a2_78 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[4], xv[5]), a2_78); + acc[28] = a2_78.x; + acc[29] = a2_78.y; + float2 a2_79 = make_float2(acc[30], acc[31]); + a2_79 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[6], xv[7]), a2_79); + acc[30] = a2_79.x; + acc[31] = a2_79.y; +#endif + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432) + 64), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432) + 1152), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432) + 1216), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432) + 2304), "r"(__float_as_uint(acc[16])), "r"(__float_as_uint(acc[17])), "r"(__float_as_uint(acc[18])), "r"(__float_as_uint(acc[19])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432) + 2368), "r"(__float_as_uint(acc[20])), "r"(__float_as_uint(acc[21])), "r"(__float_as_uint(acc[22])), "r"(__float_as_uint(acc[23])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432) + 3456), "r"(__float_as_uint(acc[24])), "r"(__float_as_uint(acc[25])), "r"(__float_as_uint(acc[26])), "r"(__float_as_uint(acc[27])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (18432 - buf_1 * 18432) + 3520), "r"(__float_as_uint(acc[28])), "r"(__float_as_uint(acc[29])), "r"(__float_as_uint(acc[30])), "r"(__float_as_uint(acc[31])) : "memory"); + if (y_pub_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 16) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 1024) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 1040) = _v4; + } + { + float4 _v4 = make_float4(acc[16 + 0], acc[16 + 1], acc[16 + 2], acc[16 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[20 + 0], acc[20 + 1], acc[20 + 2], acc[20 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2064) = _v4; + } + { + float4 _v4 = make_float4(acc[24 + 0], acc[24 + 1], acc[24 + 2], acc[24 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 3072) = _v4; + } + { + float4 _v4 = make_float4(acc[28 + 0], acc[28 + 1], acc[28 + 2], acc[28 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 3088) = _v4; + } + } + __syncthreads(); + buf_1 = 1 - buf_1; + m_cur_1 = m_next_1; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } +} + +} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu deleted file mode 100644 index be9301dbfa4..00000000000 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp_kernel.cu +++ /dev/null @@ -1,1032 +0,0 @@ -/* - * Copyright (c) 2026 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define NUM_MAIN_STAGES 1 -#define SMEM_X_BLOCK_OFF 0 -#define SMEM_X_BLOCK_STAGE_BYTES 4096 -#define SMEM_X_BLOCK_STRIDE 4096 -#define SMEM_SLABS_OFF 4096 -#define SMEM_SLABS_STAGE_BYTES 33024 -#define SMEM_SLABS_STRIDE 33024 -#define SMEM_TOTAL 37120 -#define THREADS 128 -#define LAUNCH_MIN_BLOCKS 2 - -#include - -extern "C" { - -__global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_kda_det_train_c6958a0a4ed5cb55a9a1_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(128) char smem_raw[]; - int smem; -#if __CUDA_ARCH__ == 1000 - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); -#else - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); -#endif - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - float* x_block = reinterpret_cast(smem_raw + 0); - const int x_block_addr = smem + 0; - float* slabs = reinterpret_cast(smem_raw + 4096); - const int slabs_addr = smem + 4096; - - // === Task calls (declaration order) === - int nblk = ncols / 8; - int chain = blockIdx.x / nblk; - int col0 = blockIdx.x % nblk * 8; - int col_tile = col0 / 128; - int v0 = col0 % 128; - int k = tid; - long long row_off = (long long)col_tile * 16384 + (long long)(k * 128 + v0); - int e_begin = chain_start[chain]; - int e_end = chain_start[chain + 1]; - unsigned int x_addr = x_block_addr + (unsigned int)(k * 32); - int w4r = warp * 4 + lane / 8; - int q4 = lane % 8 * 4; - float vals[8]; - float acc[8]; - float xv[4]; - float ring[64]; - float pf[32]; - long long x0_off = x0_index[chain]; - if (x0_off >= 0) { - { - float4 _v4 = *reinterpret_cast(x0 + x0_off + row_off); - vals[0 + 0] = _v4.x; - vals[0 + 1] = _v4.y; - vals[0 + 2] = _v4.z; - vals[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(x0 + x0_off + row_off + 4); - vals[4 + 0] = _v4.x; - vals[4 + 1] = _v4.y; - vals[4 + 2] = _v4.z; - vals[4 + 3] = _v4.w; - } - } else { - int diag = k - v0; - if (col_tile == 0) { - diag = -1; - } - #pragma unroll - for (int v = 0; v < 8; v++) { - vals[v] = ((v == diag) ? 1.0f : 0.0f); - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr), "r"(__float_as_uint(vals[0])), "r"(__float_as_uint(vals[1])), "r"(__float_as_uint(vals[2])), "r"(__float_as_uint(vals[3])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 16), "r"(__float_as_uint(vals[4])), "r"(__float_as_uint(vals[5])), "r"(__float_as_uint(vals[6])), "r"(__float_as_uint(vals[7])) : "memory"); - if (e_begin < e_end) { - if (write_states != 0) { - long long y_first = elem_y[e_begin]; - if (y_first >= 0) { - { - float4 _v4 = make_float4(vals[0 + 0], vals[0 + 1], vals[0 + 2], vals[0 + 3]); - *reinterpret_cast(y + y_first + row_off) = _v4; - } - { - float4 _v4 = make_float4(vals[4 + 0], vals[4 + 1], vals[4 + 2], vals[4 + 3]); - *reinterpret_cast(y + y_first + row_off + 4) = _v4; - } - } - } - } else { - long long y_last = y_final[chain]; - if (y_last >= 0) { - { - float4 _v4 = make_float4(vals[0 + 0], vals[0 + 1], vals[0 + 2], vals[0 + 3]); - *reinterpret_cast(y + y_last + row_off) = _v4; - } - { - float4 _v4 = make_float4(vals[4 + 0], vals[4 + 1], vals[4 + 2], vals[4 + 3]); - *reinterpret_cast(y + y_last + row_off + 4) = _v4; - } - } - } - if (transpose_m == 0) { - if (e_begin < e_end) { - long long m_first = elem_m[e_begin]; - #pragma unroll - for (int f = 0; f < 16; f++) { - ring[f] = m[m_first + (long long)(f * 128 + k)]; - } - #pragma unroll - for (int f_1 = 0; f_1 < 16; f_1++) { - ring[16 + f_1] = m[m_first + (long long)((16 + f_1) * 128 + k)]; - } - #pragma unroll - for (int f_2 = 0; f_2 < 16; f_2++) { - ring[32 + f_2] = m[m_first + (long long)((32 + f_2) * 128 + k)]; - } - } - __syncthreads(); - #pragma unroll 1 - for (int e = e_begin; e < e_end; e++) { - long long cur = elem_m[e]; - int has_next = 0; - long long nxt = 0; - if (e_end > e + 1) { - has_next = 1; - nxt = elem_m[e + 1]; - } - long long c_off = elem_c[e]; - if (col_tile != 0) { - c_off = -1; - } - if (c_off >= 0) { - { - float4 _v4 = *reinterpret_cast(c + c_off + row_off); - acc[0 + 0] = _v4.x; - acc[0 + 1] = _v4.y; - acc[0 + 2] = _v4.z; - acc[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off + row_off + 4); - acc[4 + 0] = _v4.x; - acc[4 + 1] = _v4.y; - acc[4 + 2] = _v4.z; - acc[4 + 3] = _v4.w; - } - } else { - acc[0] = 0.0f; - acc[1] = 0.0f; - acc[2] = 0.0f; - acc[3] = 0.0f; - acc[4] = 0.0f; - acc[5] = 0.0f; - acc[6] = 0.0f; - acc[7] = 0.0f; - } - #pragma unroll - for (int f_3 = 0; f_3 < 16; f_3++) { - ring[48 + f_3] = m[cur + (long long)((48 + f_3) * 128 + k)]; - } - #pragma unroll - for (int f_4 = 0; f_4 < 16; f_4++) { - float mv = ring[f_4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(f_4 * 8 * 4))); - acc[0] = __fmaf_rn(mv, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((f_4 * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv, xv[3], acc[7]); - } - #pragma unroll - for (int f_5 = 0; f_5 < 16; f_5++) { - ring[f_5] = m[cur + (long long)((64 + f_5) * 128 + k)]; - } - #pragma unroll - for (int f_6 = 0; f_6 < 16; f_6++) { - float mv_1 = ring[16 + f_6]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((16 + f_6) * 8 * 4))); - acc[0] = __fmaf_rn(mv_1, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_1, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_1, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_1, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((16 + f_6) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_1, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_1, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_1, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_1, xv[3], acc[7]); - } - #pragma unroll - for (int f_7 = 0; f_7 < 16; f_7++) { - ring[16 + f_7] = m[cur + (long long)((80 + f_7) * 128 + k)]; - } - #pragma unroll - for (int f_8 = 0; f_8 < 16; f_8++) { - float mv_2 = ring[32 + f_8]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((32 + f_8) * 8 * 4))); - acc[0] = __fmaf_rn(mv_2, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_2, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_2, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_2, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + f_8) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_2, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_2, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_2, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_2, xv[3], acc[7]); - } - #pragma unroll - for (int f_9 = 0; f_9 < 16; f_9++) { - ring[32 + f_9] = m[cur + (long long)((96 + f_9) * 128 + k)]; - } - #pragma unroll - for (int f_10 = 0; f_10 < 16; f_10++) { - float mv_3 = ring[48 + f_10]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((48 + f_10) * 8 * 4))); - acc[0] = __fmaf_rn(mv_3, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_3, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_3, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_3, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((48 + f_10) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_3, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_3, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_3, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_3, xv[3], acc[7]); - } - #pragma unroll - for (int f_11 = 0; f_11 < 16; f_11++) { - ring[48 + f_11] = m[cur + (long long)((112 + f_11) * 128 + k)]; - } - #pragma unroll - for (int f_12 = 0; f_12 < 16; f_12++) { - float mv_4 = ring[f_12]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((64 + f_12) * 8 * 4))); - acc[0] = __fmaf_rn(mv_4, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_4, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_4, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_4, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + f_12) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_4, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_4, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_4, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_4, xv[3], acc[7]); - } - if (has_next != 0) { - #pragma unroll - for (int f_13 = 0; f_13 < 16; f_13++) { - ring[f_13] = m[nxt + (long long)(f_13 * 128 + k)]; - } - } - #pragma unroll - for (int f_14 = 0; f_14 < 16; f_14++) { - float mv_5 = ring[16 + f_14]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((80 + f_14) * 8 * 4))); - acc[0] = __fmaf_rn(mv_5, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_5, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_5, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_5, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((80 + f_14) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_5, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_5, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_5, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_5, xv[3], acc[7]); - } - if (has_next != 0) { - #pragma unroll - for (int f_15 = 0; f_15 < 16; f_15++) { - ring[16 + f_15] = m[nxt + (long long)((16 + f_15) * 128 + k)]; - } - } - #pragma unroll - for (int f_16 = 0; f_16 < 16; f_16++) { - float mv_6 = ring[32 + f_16]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((96 + f_16) * 8 * 4))); - acc[0] = __fmaf_rn(mv_6, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_6, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_6, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_6, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + f_16) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_6, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_6, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_6, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_6, xv[3], acc[7]); - } - if (has_next != 0) { - #pragma unroll - for (int f_17 = 0; f_17 < 16; f_17++) { - ring[32 + f_17] = m[nxt + (long long)((32 + f_17) * 128 + k)]; - } - } - #pragma unroll - for (int f_18 = 0; f_18 < 16; f_18++) { - float mv_7 = ring[48 + f_18]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((112 + f_18) * 8 * 4))); - acc[0] = __fmaf_rn(mv_7, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_7, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_7, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_7, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((112 + f_18) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_7, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_7, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_7, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_7, xv[3], acc[7]); - } - __syncthreads(); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 16), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); - if (has_next != 0) { - if (write_states != 0) { - long long y_next = elem_y[e + 1]; - if (y_next >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_next + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_next + row_off + 4) = _v4; - } - } - } - } else { - long long y_out = y_final[chain]; - if (y_out >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_out + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_out + row_off + 4) = _v4; - } - } - } - __syncthreads(); - } - } else { - if (e_begin < e_end) { - long long m_first_1 = elem_m[e_begin]; - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)(w4r * 128 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((16 + w4r) * 128 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((32 + w4r) * 128 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((48 + w4r) * 128 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((64 + w4r) * 128 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((80 + w4r) * 128 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((96 + w4r) * 128 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + m_first_1 + (long long)((112 + w4r) * 128 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - slabs[q4 * 129 + w4r] = pf[0]; - slabs[(q4 + 1) * 129 + w4r] = pf[1]; - slabs[(q4 + 2) * 129 + w4r] = pf[2]; - slabs[(q4 + 3) * 129 + w4r] = pf[3]; - slabs[q4 * 129 + 16 + w4r] = pf[4]; - slabs[(q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[(q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[(q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[q4 * 129 + 32 + w4r] = pf[8]; - slabs[(q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[(q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[(q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[q4 * 129 + 48 + w4r] = pf[12]; - slabs[(q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[(q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[(q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[q4 * 129 + 64 + w4r] = pf[16]; - slabs[(q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[(q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[(q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[q4 * 129 + 80 + w4r] = pf[20]; - slabs[(q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[(q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[(q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[q4 * 129 + 96 + w4r] = pf[24]; - slabs[(q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[(q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[(q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[q4 * 129 + 112 + w4r] = pf[28]; - slabs[(q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[(q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[(q4 + 3) * 129 + 112 + w4r] = pf[31]; - } - __syncthreads(); - #pragma unroll 1 - for (int e_1 = e_begin; e_1 < e_end; e_1++) { - long long cur_1 = elem_m[e_1]; - int has_next_1 = 0; - long long nxt_1 = 0; - if (e_end > e_1 + 1) { - has_next_1 = 1; - nxt_1 = elem_m[e_1 + 1]; - } - long long c_off_1 = elem_c[e_1]; - if (col_tile != 0) { - c_off_1 = -1; - } - if (c_off_1 >= 0) { - { - float4 _v4 = *reinterpret_cast(c + c_off_1 + row_off); - acc[0 + 0] = _v4.x; - acc[0 + 1] = _v4.y; - acc[0 + 2] = _v4.z; - acc[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(c + c_off_1 + row_off + 4); - acc[4 + 0] = _v4.x; - acc[4 + 1] = _v4.y; - acc[4 + 2] = _v4.z; - acc[4 + 3] = _v4.w; - } - } else { - acc[0] = 0.0f; - acc[1] = 0.0f; - acc[2] = 0.0f; - acc[3] = 0.0f; - acc[4] = 0.0f; - acc[5] = 0.0f; - acc[6] = 0.0f; - acc[7] = 0.0f; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)(w4r * 128 + 32 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((16 + w4r) * 128 + 32 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((32 + w4r) * 128 + 32 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((48 + w4r) * 128 + 32 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((64 + w4r) * 128 + 32 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((80 + w4r) * 128 + 32 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((96 + w4r) * 128 + 32 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((112 + w4r) * 128 + 32 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - #pragma unroll - for (int kk = 0; kk < 32; kk++) { - float mv_8 = slabs[kk * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(kk * 8 * 4))); - acc[0] = __fmaf_rn(mv_8, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_8, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_8, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_8, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((kk * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_8, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_8, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_8, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_8, xv[3], acc[7]); - } - slabs[4128 + q4 * 129 + w4r] = pf[0]; - slabs[4128 + (q4 + 1) * 129 + w4r] = pf[1]; - slabs[4128 + (q4 + 2) * 129 + w4r] = pf[2]; - slabs[4128 + (q4 + 3) * 129 + w4r] = pf[3]; - slabs[4128 + q4 * 129 + 16 + w4r] = pf[4]; - slabs[4128 + (q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[4128 + (q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[4128 + (q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[4128 + q4 * 129 + 32 + w4r] = pf[8]; - slabs[4128 + (q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[4128 + (q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[4128 + (q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[4128 + q4 * 129 + 48 + w4r] = pf[12]; - slabs[4128 + (q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[4128 + (q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[4128 + (q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[4128 + q4 * 129 + 64 + w4r] = pf[16]; - slabs[4128 + (q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[4128 + (q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[4128 + (q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[4128 + q4 * 129 + 80 + w4r] = pf[20]; - slabs[4128 + (q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[4128 + (q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[4128 + (q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[4128 + q4 * 129 + 96 + w4r] = pf[24]; - slabs[4128 + (q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[4128 + (q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[4128 + (q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[4128 + q4 * 129 + 112 + w4r] = pf[28]; - slabs[4128 + (q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[4128 + (q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[4128 + (q4 + 3) * 129 + 112 + w4r] = pf[31]; - __syncthreads(); - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)(w4r * 128 + 64 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((16 + w4r) * 128 + 64 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((32 + w4r) * 128 + 64 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((48 + w4r) * 128 + 64 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((64 + w4r) * 128 + 64 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((80 + w4r) * 128 + 64 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((96 + w4r) * 128 + 64 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((112 + w4r) * 128 + 64 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - #pragma unroll - for (int kk_1 = 0; kk_1 < 32; kk_1++) { - float mv_9 = slabs[4128 + kk_1 * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((32 + kk_1) * 8 * 4))); - acc[0] = __fmaf_rn(mv_9, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_9, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_9, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_9, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((32 + kk_1) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_9, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_9, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_9, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_9, xv[3], acc[7]); - } - slabs[q4 * 129 + w4r] = pf[0]; - slabs[(q4 + 1) * 129 + w4r] = pf[1]; - slabs[(q4 + 2) * 129 + w4r] = pf[2]; - slabs[(q4 + 3) * 129 + w4r] = pf[3]; - slabs[q4 * 129 + 16 + w4r] = pf[4]; - slabs[(q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[(q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[(q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[q4 * 129 + 32 + w4r] = pf[8]; - slabs[(q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[(q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[(q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[q4 * 129 + 48 + w4r] = pf[12]; - slabs[(q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[(q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[(q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[q4 * 129 + 64 + w4r] = pf[16]; - slabs[(q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[(q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[(q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[q4 * 129 + 80 + w4r] = pf[20]; - slabs[(q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[(q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[(q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[q4 * 129 + 96 + w4r] = pf[24]; - slabs[(q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[(q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[(q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[q4 * 129 + 112 + w4r] = pf[28]; - slabs[(q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[(q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[(q4 + 3) * 129 + 112 + w4r] = pf[31]; - __syncthreads(); - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)(w4r * 128 + 96 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((16 + w4r) * 128 + 96 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((32 + w4r) * 128 + 96 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((48 + w4r) * 128 + 96 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((64 + w4r) * 128 + 96 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((80 + w4r) * 128 + 96 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((96 + w4r) * 128 + 96 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + cur_1 + (long long)((112 + w4r) * 128 + 96 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - #pragma unroll - for (int kk_2 = 0; kk_2 < 32; kk_2++) { - float mv_10 = slabs[kk_2 * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((64 + kk_2) * 8 * 4))); - acc[0] = __fmaf_rn(mv_10, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_10, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_10, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_10, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((64 + kk_2) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_10, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_10, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_10, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_10, xv[3], acc[7]); - } - slabs[4128 + q4 * 129 + w4r] = pf[0]; - slabs[4128 + (q4 + 1) * 129 + w4r] = pf[1]; - slabs[4128 + (q4 + 2) * 129 + w4r] = pf[2]; - slabs[4128 + (q4 + 3) * 129 + w4r] = pf[3]; - slabs[4128 + q4 * 129 + 16 + w4r] = pf[4]; - slabs[4128 + (q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[4128 + (q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[4128 + (q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[4128 + q4 * 129 + 32 + w4r] = pf[8]; - slabs[4128 + (q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[4128 + (q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[4128 + (q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[4128 + q4 * 129 + 48 + w4r] = pf[12]; - slabs[4128 + (q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[4128 + (q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[4128 + (q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[4128 + q4 * 129 + 64 + w4r] = pf[16]; - slabs[4128 + (q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[4128 + (q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[4128 + (q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[4128 + q4 * 129 + 80 + w4r] = pf[20]; - slabs[4128 + (q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[4128 + (q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[4128 + (q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[4128 + q4 * 129 + 96 + w4r] = pf[24]; - slabs[4128 + (q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[4128 + (q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[4128 + (q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[4128 + q4 * 129 + 112 + w4r] = pf[28]; - slabs[4128 + (q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[4128 + (q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[4128 + (q4 + 3) * 129 + 112 + w4r] = pf[31]; - __syncthreads(); - if (has_next_1 != 0) { - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)(w4r * 128 + q4)); - pf[0 + 0] = _v4.x; - pf[0 + 1] = _v4.y; - pf[0 + 2] = _v4.z; - pf[0 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((16 + w4r) * 128 + q4)); - pf[4 + 0] = _v4.x; - pf[4 + 1] = _v4.y; - pf[4 + 2] = _v4.z; - pf[4 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((32 + w4r) * 128 + q4)); - pf[8 + 0] = _v4.x; - pf[8 + 1] = _v4.y; - pf[8 + 2] = _v4.z; - pf[8 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((48 + w4r) * 128 + q4)); - pf[12 + 0] = _v4.x; - pf[12 + 1] = _v4.y; - pf[12 + 2] = _v4.z; - pf[12 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((64 + w4r) * 128 + q4)); - pf[16 + 0] = _v4.x; - pf[16 + 1] = _v4.y; - pf[16 + 2] = _v4.z; - pf[16 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((80 + w4r) * 128 + q4)); - pf[20 + 0] = _v4.x; - pf[20 + 1] = _v4.y; - pf[20 + 2] = _v4.z; - pf[20 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((96 + w4r) * 128 + q4)); - pf[24 + 0] = _v4.x; - pf[24 + 1] = _v4.y; - pf[24 + 2] = _v4.z; - pf[24 + 3] = _v4.w; - } - { - float4 _v4 = *reinterpret_cast(m + nxt_1 + (long long)((112 + w4r) * 128 + q4)); - pf[28 + 0] = _v4.x; - pf[28 + 1] = _v4.y; - pf[28 + 2] = _v4.z; - pf[28 + 3] = _v4.w; - } - } - #pragma unroll - for (int kk_3 = 0; kk_3 < 32; kk_3++) { - float mv_11 = slabs[4128 + kk_3 * 129 + k]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)((96 + kk_3) * 8 * 4))); - acc[0] = __fmaf_rn(mv_11, xv[0], acc[0]); - acc[1] = __fmaf_rn(mv_11, xv[1], acc[1]); - acc[2] = __fmaf_rn(mv_11, xv[2], acc[2]); - acc[3] = __fmaf_rn(mv_11, xv[3], acc[3]); - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) - : "r"(x_block_addr + (unsigned int)(((96 + kk_3) * 8 + 4) * 4))); - acc[4] = __fmaf_rn(mv_11, xv[0], acc[4]); - acc[5] = __fmaf_rn(mv_11, xv[1], acc[5]); - acc[6] = __fmaf_rn(mv_11, xv[2], acc[6]); - acc[7] = __fmaf_rn(mv_11, xv[3], acc[7]); - } - if (has_next_1 != 0) { - slabs[q4 * 129 + w4r] = pf[0]; - slabs[(q4 + 1) * 129 + w4r] = pf[1]; - slabs[(q4 + 2) * 129 + w4r] = pf[2]; - slabs[(q4 + 3) * 129 + w4r] = pf[3]; - slabs[q4 * 129 + 16 + w4r] = pf[4]; - slabs[(q4 + 1) * 129 + 16 + w4r] = pf[5]; - slabs[(q4 + 2) * 129 + 16 + w4r] = pf[6]; - slabs[(q4 + 3) * 129 + 16 + w4r] = pf[7]; - slabs[q4 * 129 + 32 + w4r] = pf[8]; - slabs[(q4 + 1) * 129 + 32 + w4r] = pf[9]; - slabs[(q4 + 2) * 129 + 32 + w4r] = pf[10]; - slabs[(q4 + 3) * 129 + 32 + w4r] = pf[11]; - slabs[q4 * 129 + 48 + w4r] = pf[12]; - slabs[(q4 + 1) * 129 + 48 + w4r] = pf[13]; - slabs[(q4 + 2) * 129 + 48 + w4r] = pf[14]; - slabs[(q4 + 3) * 129 + 48 + w4r] = pf[15]; - slabs[q4 * 129 + 64 + w4r] = pf[16]; - slabs[(q4 + 1) * 129 + 64 + w4r] = pf[17]; - slabs[(q4 + 2) * 129 + 64 + w4r] = pf[18]; - slabs[(q4 + 3) * 129 + 64 + w4r] = pf[19]; - slabs[q4 * 129 + 80 + w4r] = pf[20]; - slabs[(q4 + 1) * 129 + 80 + w4r] = pf[21]; - slabs[(q4 + 2) * 129 + 80 + w4r] = pf[22]; - slabs[(q4 + 3) * 129 + 80 + w4r] = pf[23]; - slabs[q4 * 129 + 96 + w4r] = pf[24]; - slabs[(q4 + 1) * 129 + 96 + w4r] = pf[25]; - slabs[(q4 + 2) * 129 + 96 + w4r] = pf[26]; - slabs[(q4 + 3) * 129 + 96 + w4r] = pf[27]; - slabs[q4 * 129 + 112 + w4r] = pf[28]; - slabs[(q4 + 1) * 129 + 112 + w4r] = pf[29]; - slabs[(q4 + 2) * 129 + 112 + w4r] = pf[30]; - slabs[(q4 + 3) * 129 + 112 + w4r] = pf[31]; - } - __syncthreads(); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr + 16), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); - if (has_next_1 != 0) { - if (write_states != 0) { - long long y_next_1 = elem_y[e_1 + 1]; - if (y_next_1 >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_next_1 + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_next_1 + row_off + 4) = _v4; - } - } - } - } else { - long long y_out_1 = y_final[chain]; - if (y_out_1 >= 0) { - { - float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); - *reinterpret_cast(y + y_out_1 + row_off) = _v4; - } - { - float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); - *reinterpret_cast(y + y_out_1 + row_off + 4) = _v4; - } - } - } - __syncthreads(); - } - } -} - -} // extern "C" From ce2923ba2e0c99079c1a3e4a4bbbccb6fcca07dc Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 17:36:28 -0800 Subject: [PATCH 03/10] cake_kda_det_train_kcp: the affine merge follows a per-architecture rule table MERGE_RULE names the merge record of a launch from T = n_chains * ncols / 128 (its 128-column chain blocks) per architecture: the 4 x 4 builds over 8 and 16 columns, the 4 x 8 build over 32 columns and the 8 x 8 build over 64 columns (packed FMA pairs on sm_100a / sm_103a; sm_107a ships the scalar builds and no 8 x 8 build). The bands are the fastest build per T measured on GB300. The "merge_rule" select kind replaces the region rule; merge_record() takes the architecture. --- .../cake_kda_det_train_kcp/README.md | 12 +- .../cake_kda_det_train_kcp/cake_kcp.py | 20 +- .../cake_kda_det_train_kcp/cake_kcp_jit.py | 235 ++++++++++++------ .../test_cake_kda_det_train_kcp_build.py | 90 ++++--- 4 files changed, 237 insertions(+), 120 deletions(-) diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/README.md b/flashinfer/experimental/cake_kda_det_train_kcp/README.md index 6eeb09f431e..6e9e05630e1 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/README.md +++ b/flashinfer/experimental/cake_kda_det_train_kcp/README.md @@ -472,15 +472,15 @@ families of a phase in one launch; off by default). The files of this package: -* `cake_kcp_jit.py` -- the `MODULES` registry of the twelve shipped programs (the +* `cake_kcp_jit.py` -- the `MODULES` registry of the thirteen shipped programs (the fragment walks of the forward and backward state recurrences in their `offset` and `tape` forms, the BF16x3 transfer walk (two CTAs per segment and head, each over one 64-column half of the state), the affine merge in one record per register-blocked build, and the two chunk stages of the context-parallel route, `kcp_fwd_o` (the output from the K1 tapes) and `kcp_dav` (`dA_qk` and the `dv1` K2 reads); selection by `always`, by - `walk_form`, or by `chains` on the merge's `ncols`, `transpose_m` and chain - count, grids by `kcp_segments_heads`, + `walk_form`, or by `merge_rule` (the `MERGE_RULE` bands of the merge + launch's 128-column chain blocks per architecture), grids by `kcp_segments_heads`, `kcp_chains`, `kcp_tile_items` or `kcp_pairs`), the logical vocabulary of the KCP arguments, `run_program(name, args, device_index=)`, the lowering of the host backend's walk and merge requests to the registry's launches, @@ -500,18 +500,18 @@ The files of this package: tests; without the file those tests skip. The package reads no environment variable. * `cake_build.py` + `build.json` -- the loader shared by every Cake-exported - package (byte-identical copies) and the manifest of the twelve shipped + package (byte-identical copies) and the manifest of the thirteen shipped programs (schema `cake-build/1`, one entry per `_`); `aot/` is never tracked (`.gitignore`). * `csrc/cake_ffi_checks.h` -- the TVM-FFI helper header the bindings include. -* `csrc/cake_kda_det_train__kcp_{kernel,binding}.cu` of the sixteen +* `csrc/cake_kda_det_train__kcp_{kernel,binding}.cu` of the seventeen context-parallel-only programs and `csrc/cake_kda_det_train__{kernel,binding}.cu` of the two chunk stages (`kcp_fwd_o`, `kcp_dav`), rendered like the single-rank ones (digests in `build.json` and `build.test.json`). ## Status -The registry holds twelve shipped programs (`cake_kcp_jit.MODULES`, in +The registry holds thirteen shipped programs (`cake_kcp_jit.MODULES`, in `build.json`) and six test-only programs (`cake_kcp_jit.TEST_MODULES`, in `build.test.json`) for `sm_100a`, `sm_103a` and `sm_107a`. diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py index c0763f96407..3cb6e756b9e 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py @@ -1378,9 +1378,7 @@ def _merge_args( cols = ( self.cols_per_cta if self.cols_per_cta is not None - else _registry_merge_cols( - ch.ncols, ch.n_chains, int(args["transpose_m"]) - ) + else _registry_merge_cols(ch.ncols, ch.n_chains, call.device) ) if ch.ncols % cols: raise ValueError( @@ -1415,11 +1413,17 @@ def _merge_args( return abi, ch.n_chains * (ch.ncols // cols) -def _registry_merge_cols(ncols: int, n_chains: int, transpose_m: int) -> int: - """``cols_per_cta`` of the K4 record the registry selects for the launch (:func:`.cake_kcp_jit.merge_record`).""" +def _registry_merge_cols(ncols: int, n_chains: int, device: torch.device) -> int: + """``cols_per_cta`` of the K4 record the registry selects for the launch on ``device`` + (:func:`.cake_kcp_jit.merge_record`; a device that is not CUDA: the record every architecture's rule names).""" from . import cake_kcp_jit # noqa: PLC0415 - H3's registry; optional at import time - record = cake_kcp_jit.merge_record(ncols, n_chains, transpose_m) + arch = ( + cake_kcp_jit.device_arch(device.index if device.index is not None else 0) + if device.type == "cuda" + else None + ) + record = cake_kcp_jit.merge_record(ncols, n_chains, arch) return int(cake_kcp_jit.MODULES[record]["grid_rule"]["cols_per_cta"]) @@ -2412,9 +2416,7 @@ def _launch( ncols = int(abi[params.index("ncols")]) cols = int(abi[params.index("cols_per_cta")]) record = jit.merge_record( - ncols, - int(grid) // (ncols // cols), - int(abi[params.index("transpose_m")]), + ncols, int(grid) // (ncols // cols), jit.device_arch(device_index) ) if int(jit.MODULES[record]["grid_rule"]["cols_per_cta"]) != cols: raise ValueError( diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index 2519bccb64d..49ea2957aa4 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -56,8 +56,9 @@ one record per register-blocked build (column block ``cols_per_cta``, thread tile, packed FMA pairs on sm_100a / sm_103a), every build giving the same bits (each output is one FP32 FMA chain in a fixed order); the - ``select_rule`` ``chains`` picks the build per launch from its column - count, operator layout and chain count (:func:`merge_record`); + ``select_rule`` ``merge_rule`` picks the build per launch from its + 128-column chain blocks and the architecture (:data:`MERGE_RULE`, + :func:`merge_record`); * ``kcp_fwd_o`` -- the chunk output stage of the forward at the rank's shape: ``out`` from the ``h`` / ``v_new`` tapes the forward walk writes and the intra stage's ``qg`` / ``Aqk`` rows, one CTA per (chunk, head) tile of @@ -250,6 +251,37 @@ "cols_per_cta", ) +# The ``kcp_affine_merge`` record of a launch (:func:`merge_record`), by architecture: ``T = n_chains * ncols / 128`` +# (the launch's 128-column chain blocks; one rule for both operator layouts) -> the record of the first band whose +# upper bound holds ``T`` (``None``: no bound). Every record gives the same bits. The bands are the fastest build per +# ``T`` in the r2611 K4 A/B on GB300 (sm_103a, 152 SMs: 17 launch rows x 11 builds, every row bitwise equal to the +# previous build); sm_100a takes the GB300 bands (not measured on sm_100a). sm_107a has no packed FMA pairs and ships +# no 8 x 8 build: its bands are the GB300 bands with the scalar builds and the 4 x 8 build above ``T`` 56 (a proxy, not +# measured on sm_107a). +MERGE_RULE: dict[str, tuple[tuple[int | None, str], ...]] = { + "sm_100a": ( + (12, "kcp_affine_merge_c8"), + (20, "kcp_affine_merge_c16"), + (40, "kcp_affine_merge_c32"), + (56, "kcp_affine_merge_c16"), + (None, "kcp_affine_merge_c64"), + ), + "sm_103a": ( + (12, "kcp_affine_merge_c8"), + (20, "kcp_affine_merge_c16"), + (40, "kcp_affine_merge_c32"), + (56, "kcp_affine_merge_c16"), + (None, "kcp_affine_merge_c64"), + ), + "sm_107a": ( + (12, "kcp_affine_merge_c8"), + (20, "kcp_affine_merge_c16"), + (40, "kcp_affine_merge_c32"), + (56, "kcp_affine_merge_c16"), + (None, "kcp_affine_merge_c32"), + ), +} + # Filled mechanically by the exporter (the derived leaves ``kernel_symbol``, ``launch``, ``arg_plan``, # ``logical``); the host plan (``program``, ``variant``, ``select_rule``, ``grid_rule``, ``library``, ``arches``) # is written by hand. Every value is a literal: FlashInfer's manifest checker and the exporter read this @@ -703,7 +735,7 @@ "select_rule": {"kind": "always"}, "variant": None, }, - "kcp_affine_merge_c64": { + "kcp_affine_merge_c8": { "arches": ["sm_100a", "sm_103a", "sm_107a"], "arg_plan": [ ["buffer", "m"], @@ -725,18 +757,72 @@ ["grid", "grid_z"], ], "ffi_entry": "run", - # The 8 x 8 thread tile over 64 columns (packed FMA pairs on sm_100a / sm_103a): the widest - # launches (provisional rule). - "grid_rule": {"cols_per_cta": 64, "kind": "kcp_chains"}, - "kernel_symbol": "kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp", + # The 4 x 4 thread tile over 8 columns (packed FMA pairs on sm_100a / sm_103a): T <= 12 + # (MERGE_RULE). + "grid_rule": {"cols_per_cta": 8, "kind": "kcp_chains"}, + "kernel_symbol": "kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp", + "launch": { + "block": (64, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 0, + "use_pdl": False, + }, + "library": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp", + "logical": { + "c": "c", + "chain_start": "chain_start", + "cols_per_cta": "cols_per_cta", + "elem_c": "elem_c", + "elem_m": "elem_m", + "elem_y": "elem_y", + "m": "m", + "ncols": "ncols", + "transpose_m": "transpose_m", + "write_states": "write_states", + "x0": "x0", + "x0_index": "x0_index", + "y": "y", + "y_final": "y_final", + }, + "program": "kcp_affine_merge", + "select_rule": {"kind": "merge_rule"}, + "variant": "c8", + }, + "kcp_affine_merge_c16": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "m"], + ["buffer", "c"], + ["buffer", "x0"], + ["buffer", "y"], + ["buffer", "chain_start"], + ["buffer", "elem_m"], + ["buffer", "elem_c"], + ["buffer", "elem_y"], + ["buffer", "x0_index"], + ["buffer", "y_final"], + ["parameter", "ncols"], + ["parameter", "transpose_m"], + ["parameter", "write_states"], + ["parameter", "cols_per_cta"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # The 4 x 4 thread tile over 16 columns (packed FMA pairs on sm_100a / sm_103a): T 13-20 and + # 41-56 (MERGE_RULE). + "grid_rule": {"cols_per_cta": 16, "kind": "kcp_chains"}, + "kernel_symbol": "kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp", "launch": { "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 143360, + "dynamic_smem_bytes": 0, "use_pdl": False, }, - "library": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp", + "library": "cake_kda_det_train_9ab870ac43ee331f2dd3_kcp", "logical": { "c": "c", "chain_start": "chain_start", @@ -754,8 +840,8 @@ "y_final": "y_final", }, "program": "kcp_affine_merge", - "select_rule": {"kind": "chains", "regions": [[128, None, 48, None]]}, - "variant": "c64", + "select_rule": {"kind": "merge_rule"}, + "variant": "c16", }, "kcp_affine_merge_c32": { "arches": ["sm_100a", "sm_103a", "sm_107a"], @@ -779,14 +865,15 @@ ["grid", "grid_z"], ], "ffi_entry": "run", - # The 4 x 8 thread tile over 32 columns (packed FMA pairs on sm_100a / sm_103a) (provisional rule). + # The 4 x 8 thread tile over 32 columns (packed FMA pairs on sm_100a / sm_103a): T 21-40, and + # above 56 on sm_107a (MERGE_RULE). "grid_rule": {"cols_per_cta": 32, "kind": "kcp_chains"}, "kernel_symbol": "kernel_cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp", "launch": { "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 110592, + "dynamic_smem_bytes": 0, "use_pdl": False, }, "library": "cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp", @@ -807,11 +894,11 @@ "y_final": "y_final", }, "program": "kcp_affine_merge", - "select_rule": {"kind": "chains", "regions": [[128, None, 17, 47]]}, + "select_rule": {"kind": "merge_rule"}, "variant": "c32", }, - "kcp_affine_merge_c16": { - "arches": ["sm_100a", "sm_103a", "sm_107a"], + "kcp_affine_merge_c64": { + "arches": ["sm_100a", "sm_103a"], "arg_plan": [ ["buffer", "m"], ["buffer", "c"], @@ -832,17 +919,18 @@ ["grid", "grid_z"], ], "ffi_entry": "run", - # The 4 x 4 thread tile over 16 columns: few chains and the map chains (provisional rule). - "grid_rule": {"cols_per_cta": 16, "kind": "kcp_chains"}, - "kernel_symbol": "kernel_cake_kda_det_train_15e81414b27b4b831808_kcp", + # The 8 x 8 thread tile over 64 columns with packed FMA pairs (sm_100a / sm_103a only): T above + # 56 (MERGE_RULE). + "grid_rule": {"cols_per_cta": 64, "kind": "kcp_chains"}, + "kernel_symbol": "kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp", "launch": { "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 61440, + "dynamic_smem_bytes": 0, "use_pdl": False, }, - "library": "cake_kda_det_train_15e81414b27b4b831808_kcp", + "library": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp", "logical": { "c": "c", "chain_start": "chain_start", @@ -860,11 +948,8 @@ "y_final": "y_final", }, "program": "kcp_affine_merge", - "select_rule": { - "kind": "chains", - "regions": [[128, None, 1, 16], [256, None, 1, None]], - }, - "variant": "c16", + "select_rule": {"kind": "merge_rule"}, + "variant": "c64", }, "kcp_fwd_o": { "arches": ["sm_100a", "sm_103a", "sm_107a"], @@ -1342,27 +1427,14 @@ def records_of_program(program: str) -> tuple[str, ...]: return tuple(name for name, rec in MODULES.items() if rec["program"] == program) -def _in_region( - region: Sequence[Any], ncols: int, n_chains: int, transpose_m: int -) -> bool: - """Whether a ``chains`` region ``[ncols, transpose_m, min_chains, max_chains]`` admits the launch - (``transpose_m`` / ``max_chains`` ``None``: any).""" - width, layout, low, high = region - return ( - ncols == int(width) - and (layout is None or transpose_m == int(layout)) - and n_chains >= int(low) - and (high is None or n_chains <= int(high)) - ) - - def _admits( + name: str, rule: Mapping[str, Any], *, ncols: int | None, form: str | None = None, n_chains: int | None = None, - transpose_m: int | None = None, + arch: str | None = None, ) -> bool: kind = rule["kind"] if kind == "always": @@ -1371,13 +1443,10 @@ def _admits( return ncols is not None and int(ncols) == int(rule["ncols"]) if kind == "walk_form": return form is not None and str(form) == str(rule["form"]) - if kind == "chains": - if ncols is None or n_chains is None or transpose_m is None: + if kind == "merge_rule": + if ncols is None or n_chains is None: return False - return any( - _in_region(region, int(ncols), int(n_chains), int(transpose_m)) - for region in rule["regions"] - ) + return _merge_rule_record(int(ncols), int(n_chains), arch) == name raise ValueError(f"unknown Cake KDA KCP selection rule: {kind!r}") @@ -1387,53 +1456,74 @@ def select_record( ncols: int | None = None, form: str | None = None, n_chains: int | None = None, - transpose_m: int | None = None, + arch: str | None = None, ) -> str: """The one record of ``program`` the launch's scalars select. ``always`` is the rule of a single-record program; ``columns`` (``{"kind": "columns", "ncols": N}``) picks the variant whose column count is the - chain's ``ncols``; ``chains`` (``{"kind": "chains", "regions": [[ncols, - transpose_m, min_chains, max_chains], ...]}``) picks the merge build whose - regions hold the launch's ``(ncols, transpose_m, n_chains)`` - (``transpose_m`` / ``max_chains`` ``None``: any; :func:`merge_record`); - ``walk_form`` (``{"kind": "walk_form", "form": F}``) picks the K1 / K2 - walk record of form ``F`` -- ``"full"`` (one CTA per segment and head) or - ``"split"`` (the two-CTA record over 64 value columns each), the - launch's form from :func:`walk_form`; a program under a fixed rule ships - its one ``always`` record, the other form being test-only - (:func:`test_walk_record`). Exactly one record admits every launch, or - the registry is inconsistent. + chain's ``ncols``; ``merge_rule`` (``{"kind": "merge_rule"}``) picks the + merge record :data:`MERGE_RULE` names for the launch's ``ncols``, + ``n_chains`` and architecture (:func:`merge_record`); ``walk_form`` + (``{"kind": "walk_form", "form": F}``) picks the K1 / K2 walk record of + form ``F`` -- ``"full"`` (one CTA per segment and head) or ``"split"`` + (the two-CTA record over 64 value columns each), the launch's form from + :func:`walk_form`; a program under a fixed rule ships its one ``always`` + record, the other form being test-only (:func:`test_walk_record`). + Exactly one record admits every launch, or the registry is inconsistent. """ admitted = [ name for name in records_of_program(program) if _admits( + name, MODULES[name]["select_rule"], ncols=ncols, form=form, n_chains=n_chains, - transpose_m=transpose_m, + arch=arch, ) ] if len(admitted) != 1: raise ValueError( f"program {program!r}: {len(admitted)} records admit ncols={ncols} form={form} n_chains={n_chains} " - f"transpose_m={transpose_m} ({admitted}); expected one" + f"arch={arch} ({admitted}); expected one" ) return admitted[0] -def merge_record(ncols: int, n_chains: int, transpose_m: int) -> str: - """The ``kcp_affine_merge`` record of a launch of ``n_chains`` chains over ``ncols`` columns with operator - layout ``transpose_m``: the build the records' ``chains`` regions name (every build gives the same bits, the FMA - order of every output being fixed; the rule picks the fastest measured build per launch shape).""" - return select_record( - "kcp_affine_merge", - ncols=int(ncols), - n_chains=int(n_chains), - transpose_m=int(transpose_m), - ) +def _merge_rule_record(ncols: int, n_chains: int, arch: str | None) -> str | None: + """The record :data:`MERGE_RULE` names for ``T = n_chains * ncols / 128`` on ``arch``; ``arch`` ``None``: the + record every architecture's bands name (``None`` where they differ or ``T`` is not a positive whole count).""" + if ncols < 128 or ncols % 128 or n_chains < 1: + return None + blocks = n_chains * (ncols // 128) + named = { + next( + record + for bound, record in MERGE_RULE[a] + if bound is None or blocks <= int(bound) + ) + for a in (MERGE_RULE if arch is None else (arch,)) + } + return named.pop() if len(named) == 1 else None + + +def merge_record(ncols: int, n_chains: int, arch: str | None = None) -> str: + """The ``kcp_affine_merge`` record of a launch of ``n_chains`` chains over ``ncols`` columns on ``arch`` + (:data:`MERGE_RULE`; every record gives the same bits, each output being one FP32 FMA chain in a fixed order). + ``arch`` ``None`` asks for the record every architecture's rule names (a host-side query without a device).""" + if arch is not None and arch not in MERGE_RULE: + raise ValueError( + f"the affine merge ships no build for {arch}; its architectures: {tuple(MERGE_RULE)}" + ) + record = _merge_rule_record(int(ncols), int(n_chains), arch) + if record is None: + raise ValueError( + f"no affine merge record for ncols={ncols} n_chains={n_chains} arch={arch}: ncols must be a positive " + "multiple of 128 and n_chains >= 1 (and, without an architecture, every architecture's rule must agree)" + ) + return record def walk_form( @@ -2104,7 +2194,7 @@ def _merge(args: Mapping[str, Any], call: Any, device_index: int) -> None: return device = str(call.device) ncols = int(chains.ncols) - record = merge_record(ncols, n_chains, int(args["transpose_m"])) + record = merge_record(ncols, n_chains, _arch_of(device_index)) dummy = _dummy_tile(device) values: dict[str, Any] = { "m": _flat_pool(args["m"], "kcp_affine_merge.m"), @@ -2176,6 +2266,7 @@ def run_program(name: str, args: Mapping[str, Any], *, device_index: int) -> Non "ARCHES", "LOGICAL_SCALARS", "LOGICAL_TENSORS", + "MERGE_RULE", "MODULES", "PRECISION_MODES", "PROGRAMS", diff --git a/tests/experimental/test_cake_kda_det_train_kcp_build.py b/tests/experimental/test_cake_kda_det_train_kcp_build.py index 1be6495c836..9f49a6dbf51 100644 --- a/tests/experimental/test_cake_kda_det_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_det_train_kcp_build.py @@ -74,34 +74,47 @@ ["buffer", "do_tma"], ["buffer", "vnew_tma"], ["buffer", "aqk_tma"], ["buffer", "dAqk"], ["buffer", "dv1"], ["parameter", "num_heads"], ["parameter", "scale"], *GRID, ] # fmt: skip -# The K4 merge records (the register-blocked v2 builds, one per CONFIGS entry the select rule names): record -> -# (build keywords, threads per CTA = 128 * cols_per_cta / the thread tile's outputs). +# The K4 merge records (the register-blocked v2 builds the rule names): record -> (build keywords, threads per CTA = +# 128 * cols_per_cta / the thread tile's outputs, architectures). The pairs builds are scalar on sm_107a (the +# exporter's per-architecture keywords; the program name hashes the keywords above). MERGE_V2_MODULE = "tart.kda_chunk_kcp_affine_merge.sm100_sm103_sm107.kda_chunk_det_affine_merge_fp32_v2_kcp" MERGE_RECORDS = { - "kcp_affine_merge_c64": ( - {"chunk_k": 32, "cols_per_cta": 64, "pairs": True, "ring": 4, "tile": "8x8"}, + "kcp_affine_merge_c8": ( + {"chunk_k": 16, "cols_per_cta": 8, "pairs": True, "ring": 4, "tile": "4x4"}, + 64, + ARCHES, + ), + "kcp_affine_merge_c16": ( + {"chunk_k": 16, "cols_per_cta": 16, "pairs": True, "ring": 4, "tile": "4x4"}, 128, + ARCHES, ), "kcp_affine_merge_c32": ( {"chunk_k": 32, "cols_per_cta": 32, "pairs": True, "ring": 4, "tile": "4x8"}, 128, + ARCHES, ), - "kcp_affine_merge_c16": ( - {"chunk_k": 16, "cols_per_cta": 16, "pairs": False, "ring": 4, "tile": "4x4"}, + "kcp_affine_merge_c64": ( + {"chunk_k": 32, "cols_per_cta": 64, "pairs": True, "ring": 4, "tile": "8x8"}, 128, + ("sm_100a", "sm_103a"), ), } -# The K4 select rule at sample launches: (ncols, transpose_m, n_chains) -> record. +# The K4 rule at sample launches: (arch, ncols, n_chains) -> record; T = n_chains * ncols / 128 (the W4 merge launches +# are T = 64; a 256-column map chain counts twice). MERGE_SELECT = { - (128, 0, 64): "kcp_affine_merge_c64", - (128, 1, 64): "kcp_affine_merge_c64", - (128, 0, 48): "kcp_affine_merge_c64", - (128, 1, 47): "kcp_affine_merge_c32", - (128, 0, 17): "kcp_affine_merge_c32", - (128, 0, 16): "kcp_affine_merge_c16", - (128, 1, 1): "kcp_affine_merge_c16", - (256, 0, 24): "kcp_affine_merge_c16", - (256, 1, 96): "kcp_affine_merge_c16", + ("sm_103a", 128, 64): "kcp_affine_merge_c64", + ("sm_100a", 128, 57): "kcp_affine_merge_c64", + ("sm_107a", 128, 64): "kcp_affine_merge_c32", + ("sm_103a", 128, 56): "kcp_affine_merge_c16", + ("sm_103a", 256, 21): "kcp_affine_merge_c16", + ("sm_103a", 128, 40): "kcp_affine_merge_c32", + ("sm_107a", 128, 21): "kcp_affine_merge_c32", + ("sm_103a", 256, 10): "kcp_affine_merge_c16", + ("sm_103a", 128, 13): "kcp_affine_merge_c16", + ("sm_103a", 128, 12): "kcp_affine_merge_c8", + ("sm_107a", 256, 6): "kcp_affine_merge_c8", + ("sm_100a", 128, 1): "kcp_affine_merge_c8", } PLANS = { "kcp_fwdh_offset": FWDH_PLAN, @@ -155,7 +168,7 @@ f"tart.kda_chunk_kcp_transfer_walk.{_W}.kda_chunk_kcp_transfer_walk_bf16x3_v1", {"vh": 64}, ), - **{name: (MERGE_V2_MODULE, kwargs) for name, (kwargs, _) in MERGE_RECORDS.items()}, + **{name: (MERGE_V2_MODULE, row[0]) for name, row in MERGE_RECORDS.items()}, "kcp_fwd_o": (f"tart.kda_chunk_kcp_fwd_o.{_W}.kda_chunk_det_fwd_o_v1_kcp", {}), "kcp_dav": (f"tart.kda_chunk_bwd_dav.{_W}.kda_chunk_train_bwd_dav_v1", {}), "kcp_transfer_walk_three_mma": ( @@ -223,7 +236,9 @@ def test_the_tables_are_disjoint_and_complete() -> None: def test_record_matches_the_frozen_abi(name: str) -> None: record = _records()[name] assert record["arg_plan"] == PLANS[name], name - assert record["arches"] == list(ARCHES) + assert record["arches"] == list( + MERGE_RECORDS[name][2] if name in MERGE_RECORDS else ARCHES + ) assert record["ffi_entry"] == "run" assert record["kernel_symbol"] == f"kernel_{record['library']}" assert ( @@ -276,27 +291,36 @@ def test_program_names_follow_the_identity_rule() -> None: def test_merge_selection_covers_every_launch_once() -> None: - for (ncols, transpose_m, n_chains), record in MERGE_SELECT.items(): + assert tuple(kcp_jit.MERGE_RULE) == ARCHES + for (arch, ncols, n_chains), record in MERGE_SELECT.items(): assert ( kcp_jit.select_record( - "kcp_affine_merge", - ncols=ncols, - n_chains=n_chains, - transpose_m=transpose_m, + "kcp_affine_merge", ncols=ncols, n_chains=n_chains, arch=arch ) == record - ), (ncols, transpose_m, n_chains) - assert kcp_jit.merge_record(ncols, n_chains, transpose_m) == record - # every state / map launch shape is admitted by exactly one record (select_record raises otherwise) - for ncols in (128, 256): - for transpose_m in (0, 1): + ), (arch, ncols, n_chains) + assert kcp_jit.merge_record(ncols, n_chains, arch) == record + # every state / map launch is admitted by exactly one record built for its architecture (select_record raises + # otherwise), and the rule names every merge record on every architecture it ships for + named: dict[str, set[str]] = {} + for arch in ARCHES: + for ncols in (128, 256): for n_chains in range(1, 300): - kcp_jit.merge_record(ncols, n_chains, transpose_m) + record = kcp_jit.merge_record(ncols, n_chains, arch) + assert arch in kcp_jit.MODULES[record]["arches"], (arch, record) + named.setdefault(record, set()).add(arch) + assert named == {name: set(MERGE_RECORDS[name][2]) for name in MERGE_RECORDS} + # without an architecture: the record every architecture names, refused where they differ (T > 56) + assert kcp_jit.merge_record(128, 8) == "kcp_affine_merge_c8" + with pytest.raises(ValueError): + kcp_jit.merge_record(128, 64) with pytest.raises(ValueError): - kcp_jit.merge_record(64, 4, 0) - # the chains rule needs the chain count and the layout + kcp_jit.merge_record(64, 4, "sm_103a") with pytest.raises(ValueError): - kcp_jit.select_record("kcp_affine_merge", ncols=128) + kcp_jit.merge_record(128, 4, "sm_90a") + # the rule needs the chain count + with pytest.raises(ValueError): + kcp_jit.select_record("kcp_affine_merge", ncols=128, arch="sm_103a") def test_selection_and_grids() -> None: @@ -452,7 +476,7 @@ def test_shipped_manifest_lists_the_kcp_programs_and_no_test_only_one() -> None: manifest = cake_build.manifest() libraries = cake_build.libraries_by_name() for name, record in kcp_jit.MODULES.items(): - for arch in ARCHES: + for arch in record["arches"]: library = kcp_jit.library_name(name, arch) assert library == f"{record['library']}_{arch}" entry = libraries[library] From 789d13c35e54e96f0cfe5fa42c7ae7f4587cc981 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 17:49:02 -0800 Subject: [PATCH 04/10] cake_kda_det_train_kcp: render the affine merge builds the rule table names The merge ships four register-blocked builds: c8 4x4, c16 4x4 and c32 4x8 on sm_100a, sm_103a and sm_107a (column pairs on sm_100a/sm_103a), and c64 8x8 with column pairs on sm_100a/sm_103a only. --- .../cake_kda_det_train_kcp/build.json | 204 +-- .../cake_kda_det_train_kcp/cake_kcp_jit.py | 8 +- ..._train_7b6b22d3423f3138baf7_kcp_binding.cu | 156 +++ ...t_train_7b6b22d3423f3138baf7_kcp_kernel.cu | 1191 +++++++++++++++++ ...t_train_7ddde81b137fb82d27d5_kcp_kernel.cu | 338 ----- ...train_9ab870ac43ee331f2dd3_kcp_binding.cu} | 16 +- ..._train_9ab870ac43ee331f2dd3_kcp_kernel.cu} | 213 ++- 7 files changed, 1694 insertions(+), 432 deletions(-) create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu rename flashinfer/experimental/cake_kda_det_train_kcp/csrc/{cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu => cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu} (92%) rename flashinfer/experimental/cake_kda_det_train_kcp/csrc/{cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu => cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu} (80%) diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/build.json b/flashinfer/experimental/cake_kda_det_train_kcp/build.json index 4e2d1bbae06..5a5e5b4e801 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/build.json +++ b/flashinfer/experimental/cake_kda_det_train_kcp/build.json @@ -60,69 +60,6 @@ "rendered_at": "2026-11-03T23:32:53Z" }, "libraries": [ - { - "arch": "sm_100a", - "closure_sha256": "1ebfa9e2976961a1ca7742ef0e8e9b8429e5ed298f2b8de61e555c40ac0945f2", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_kda_det_train_15e81414b27b4b831808_kcp_sm_100a", - "nvcc_flags": [ - "--use_fast_math" - ], - "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_15e81414b27b4b831808_kcp_sm_100a.so", - "sha256": { - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu": "a65f44614f617eabf7a6ac5635a00ffb5547d177460e3f66a4a014600d263288", - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu": "29d89e83f021d69047583d2008293e6321209ec0b6d37b86e3663d6f8bbafa6f" - }, - "sources": [ - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu", - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "3e88ddd0f81033a7cca8ea00e0634de10f88d88a41b0744f9d4b43b676a6b02b", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_kda_det_train_15e81414b27b4b831808_kcp_sm_103a", - "nvcc_flags": [ - "--use_fast_math" - ], - "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_15e81414b27b4b831808_kcp_sm_103a.so", - "sha256": { - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu": "a65f44614f617eabf7a6ac5635a00ffb5547d177460e3f66a4a014600d263288", - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu": "29d89e83f021d69047583d2008293e6321209ec0b6d37b86e3663d6f8bbafa6f" - }, - "sources": [ - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu", - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "208de082eb390338fcae73c52912bd052041f0c3d74cd310cae5eda5e624e2ed", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_kda_det_train_15e81414b27b4b831808_kcp_sm_107a", - "nvcc_flags": [ - "--use_fast_math" - ], - "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_15e81414b27b4b831808_kcp_sm_107a.so", - "sha256": { - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu": "a65f44614f617eabf7a6ac5635a00ffb5547d177460e3f66a4a014600d263288", - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu": "29d89e83f021d69047583d2008293e6321209ec0b6d37b86e3663d6f8bbafa6f" - }, - "sources": [ - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu", - "csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu" - ] - }, { "arch": "sm_100a", "closure_sha256": "a4751e59d5ad550996fe0b1aaaa85ecebb03a5e962a4b5b1a14c2571469df361", @@ -440,61 +377,103 @@ }, { "arch": "sm_100a", - "closure_sha256": "c53a166bc23db11af242f70da51f2b6e7b80d0c54b7f785485bba05d16bc81d8", + "closure_sha256": "1abe6d29eeab5be3d96e0dc0a4e6b7a646a89817892fbb3b7704ab1ad5883ea0", "entry": "run", "gencode": "-gencode=arch=compute_100a,code=sm_100a", "min_cuda": "12.8", - "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_100a", + "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_100a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_100a.so", + "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_100a.so", "sha256": { - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu": "7d580a68978de3029cdcb0b44c43361f37acdc101cd7a074907dc8d5604f0eba", - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "7e30683f1adde5744209052239eb51e83b0e2dd6537acd74624d4f0e33f3a5a5" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" }, "sources": [ - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu", - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" ] }, { "arch": "sm_103a", - "closure_sha256": "26e4143bc8f2d06f256f5af75a6fad89f8da44971d505e0143b90f7a1f458116", + "closure_sha256": "c523a00b55ca9add18bd82866a5bd53223caf3a15ccc9a98fc2a6d4932e53125", "entry": "run", "gencode": "-gencode=arch=compute_103a,code=sm_103a", "min_cuda": "12.9", - "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_103a", + "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_103a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_103a.so", + "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_103a.so", "sha256": { - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu": "7d580a68978de3029cdcb0b44c43361f37acdc101cd7a074907dc8d5604f0eba", - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "7e30683f1adde5744209052239eb51e83b0e2dd6537acd74624d4f0e33f3a5a5" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" }, "sources": [ - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu", - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" ] }, { "arch": "sm_107a", - "closure_sha256": "3f0f09169653ee87e8757c7f767e100e63d18d6e644e558434136ebefed78f9f", + "closure_sha256": "b6b3daf2d4fccef3711a403214d28c2c6386100a7dc5fd3fe0942ba7fb7fd89a", "entry": "run", "gencode": "-gencode=arch=compute_107a,code=sm_107a", "min_cuda": "13.4", - "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_107a", + "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_107a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_107a.so", + "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_107a.so", + "sha256": { + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" + }, + "sources": [ + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "14863db81ff73ea8326b93cdb64f8c0bc5520285fd187b65b5e5b3f0f23b21c0", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_100a.so", + "sha256": { + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu": "7d580a68978de3029cdcb0b44c43361f37acdc101cd7a074907dc8d5604f0eba", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "21c8ff52f48c4f338c2ab603903f7c7cafd7d6b8a0269c7c20f041ba5b9f726d" + }, + "sources": [ + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu", + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "2a6e28df442548ce7930950cb33ae61efd4f33882ce2c5f0864b091888e7fc00", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_sm_103a.so", "sha256": { "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_binding.cu": "7d580a68978de3029cdcb0b44c43361f37acdc101cd7a074907dc8d5604f0eba", - "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "7e30683f1adde5744209052239eb51e83b0e2dd6537acd74624d4f0e33f3a5a5" + "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu": "21c8ff52f48c4f338c2ab603903f7c7cafd7d6b8a0269c7c20f041ba5b9f726d" }, "sources": [ "csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu", @@ -564,6 +543,69 @@ "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu" ] }, + { + "arch": "sm_100a", + "closure_sha256": "6e64c1a93e5a3dfa7800f17efb909c97e91ce9a00ce88c0c7f1751f8c17f66f8", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_sm_100a.so", + "sha256": { + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu": "5a4539283e1d2e3bb512e48ef962e5def78a1b8d2059efdbcfa74f7dcf30f231", + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu": "4459b78abd52240e43b3705b363463b29e4781a577d19560adf19f2c3b1f1dbf" + }, + "sources": [ + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu", + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "df1caed18b42212f906f1f9ffe5ef4a2f07eab4eb69c6a8a6fa46043259e3de4", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_sm_103a.so", + "sha256": { + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu": "5a4539283e1d2e3bb512e48ef962e5def78a1b8d2059efdbcfa74f7dcf30f231", + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu": "4459b78abd52240e43b3705b363463b29e4781a577d19560adf19f2c3b1f1dbf" + }, + "sources": [ + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu", + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "9f20815ad5c578e9c921e5209efc80d9c8f43bde96860a3a03eb6f2b74c3eae7", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_sm_107a.so", + "sha256": { + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu": "5a4539283e1d2e3bb512e48ef962e5def78a1b8d2059efdbcfa74f7dcf30f231", + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu": "4459b78abd52240e43b3705b363463b29e4781a577d19560adf19f2c3b1f1dbf" + }, + "sources": [ + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu", + "csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "3942bb3852909199331c9c4c897239294d4bdfb1cc392c05eec89beb0306986d", diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index 49ea2957aa4..5973ffa40d8 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -765,7 +765,7 @@ "block": (64, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 53248, "use_pdl": False, }, "library": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp", @@ -819,7 +819,7 @@ "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 61440, "use_pdl": False, }, "library": "cake_kda_det_train_9ab870ac43ee331f2dd3_kcp", @@ -873,7 +873,7 @@ "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 110592, "use_pdl": False, }, "library": "cake_kda_det_train_c2ad3b63ad48e23ccbdb_kcp", @@ -927,7 +927,7 @@ "block": (128, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 143360, "use_pdl": False, }, "library": "cake_kda_det_train_7ddde81b137fb82d27d5_kcp", diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu new file mode 100644 index 00000000000..a2f34d122b0 --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu @@ -0,0 +1,156 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp, for the device source that defines it. Do not edit. +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); + +namespace cake_binding_kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) { + TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b + << " has no int64 value in a host check"; +} + +int64_t CMod(int64_t a, int64_t b) { + if (b == 0) NoInt64(a, "%", b); + return b == -1 ? 0 : a % b; +} + +void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, TensorView a_chain_start, TensorView a_elem_m, TensorView a_elem_c, TensorView a_elem_y, TensorView a_x0_index, TensorView a_y_final, int64_t a_ncols, int64_t a_transpose_m, int64_t a_write_states, int64_t a_cols_per_cta, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_m.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_m, "m"); + check_dtype(a_m, DLDataType{kDLFloat, 32, 1}, "m"); + check_cuda_tensor(a_c, "c"); + check_dtype(a_c, DLDataType{kDLFloat, 32, 1}, "c"); + check_cuda_tensor(a_x0, "x0"); + check_dtype(a_x0, DLDataType{kDLFloat, 32, 1}, "x0"); + check_cuda_tensor(a_y, "y"); + check_dtype(a_y, DLDataType{kDLFloat, 32, 1}, "y"); + check_cuda_tensor(a_chain_start, "chain_start"); + check_dtype(a_chain_start, DLDataType{kDLInt, 32, 1}, "chain_start"); + check_cuda_tensor(a_elem_m, "elem_m"); + check_dtype(a_elem_m, DLDataType{kDLInt, 64, 1}, "elem_m"); + check_cuda_tensor(a_elem_c, "elem_c"); + check_dtype(a_elem_c, DLDataType{kDLInt, 64, 1}, "elem_c"); + check_cuda_tensor(a_elem_y, "elem_y"); + check_dtype(a_elem_y, DLDataType{kDLInt, 64, 1}, "elem_y"); + check_cuda_tensor(a_x0_index, "x0_index"); + check_dtype(a_x0_index, DLDataType{kDLInt, 64, 1}, "x0_index"); + check_cuda_tensor(a_y_final, "y_final"); + check_dtype(a_y_final, DLDataType{kDLInt, 64, 1}, "y_final"); + TVM_FFI_CHECK(a_ncols >= -2147483648LL && a_ncols <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'ncols' = " << a_ncols + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_transpose_m >= -2147483648LL && a_transpose_m <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'transpose_m' = " << a_transpose_m + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_write_states >= -2147483648LL && a_write_states <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'write_states' = " << a_write_states + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_cols_per_cta >= -2147483648LL && a_cols_per_cta <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'cols_per_cta' = " << a_cols_per_cta + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_m, "m"); + check_contiguous(a_c, "c"); + check_same_device(a_c, a_m, "c", "m"); + check_contiguous(a_x0, "x0"); + check_same_device(a_x0, a_m, "x0", "m"); + check_contiguous(a_y, "y"); + check_same_device(a_y, a_m, "y", "m"); + check_contiguous(a_chain_start, "chain_start"); + check_same_device(a_chain_start, a_m, "chain_start", "m"); + check_contiguous(a_elem_m, "elem_m"); + check_same_device(a_elem_m, a_m, "elem_m", "m"); + check_contiguous(a_elem_c, "elem_c"); + check_same_device(a_elem_c, a_m, "elem_c", "m"); + check_contiguous(a_elem_y, "elem_y"); + check_same_device(a_elem_y, a_m, "elem_y", "m"); + check_contiguous(a_x0_index, "x0_index"); + check_same_device(a_x0_index, a_m, "x0_index", "m"); + check_contiguous(a_y_final, "y_final"); + check_same_device(a_y_final, a_m, "y_final", "m"); + TVM_FFI_CHECK((a_cols_per_cta == 8LL), ValueError) + << KERNEL_NAME << ": requires (cols_per_cta == 8)"; + TVM_FFI_CHECK((CMod(a_ncols, 128LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(ncols, 128) == 0)"; + TVM_FFI_CHECK((a_ncols >= 128LL), ValueError) + << KERNEL_NAME << ": requires (ncols >= 128)"; + TVM_FFI_CHECK((a_ncols <= 256LL), ValueError) + << KERNEL_NAME << ": requires (ncols <= 256)"; + TVM_FFI_CHECK((a_transpose_m >= 0LL), ValueError) + << KERNEL_NAME << ": requires (transpose_m >= 0)"; + TVM_FFI_CHECK((a_transpose_m <= 1LL), ValueError) + << KERNEL_NAME << ": requires (transpose_m <= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 53248), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + void* v_m = DataPtr(a_m); + void* v_c = DataPtr(a_c); + void* v_x0 = DataPtr(a_x0); + void* v_y = DataPtr(a_y); + void* v_chain_start = DataPtr(a_chain_start); + void* v_elem_m = DataPtr(a_elem_m); + void* v_elem_c = DataPtr(a_elem_c); + void* v_elem_y = DataPtr(a_elem_y); + void* v_x0_index = DataPtr(a_x0_index); + void* v_y_final = DataPtr(a_y_final); + int32_t v_ncols = static_cast(a_ncols); + int32_t v_transpose_m = static_cast(a_transpose_m); + int32_t v_write_states = static_cast(a_write_states); + int32_t v_cols_per_cta = static_cast(a_cols_per_cta); + void* kernel_args[] = {&v_m, &v_c, &v_x0, &v_y, &v_chain_start, &v_elem_m, &v_elem_c, &v_elem_y, &v_x0_index, &v_y_final, &v_ncols, &v_transpose_m, &v_write_states, &v_cols_per_cta}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(64, 1, 1); + config.dynamicSmemBytes = 53248; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu new file mode 100644 index 00000000000..98765540d9e --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu @@ -0,0 +1,1191 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LAUNCH_MIN_BLOCKS 4 + +#include +#if __CUDA_ARCH__ == 1070 +#else + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} +#endif + +extern "C" { + +__global__ __launch_bounds__(64, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_det_train_7b6b22d3423f3138baf7_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(128) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Prologue + float* xs = reinterpret_cast(smem_raw + 0); + const int xs_addr = smem + 0; + float* ms = reinterpret_cast(smem_raw + 12288); + const int ms_addr = smem + 12288; + + // === Task calls (declaration order) === + int nblk = ncols / 8; + int chain = blockIdx.x / nblk; + int col0 = blockIdx.x % nblk * 8; + int col_tile = col0 / 128; + int v0 = col0 % 128; + int warp_1 = warp; + int g = lane / 2; + int cg = lane % 2; + int e_begin = chain_start[chain]; + int e_end = chain_start[chain + 1]; + unsigned int ring_w = ms_addr + (unsigned int)(warp_1 * 20480); + float acc[16]; + float cn[16]; + float mv[16]; + float xv[4]; + if (transpose_m == 0) { + int row0 = warp_1 * 64 + g * 4; + long long g_off = (long long)col_tile * 16384 + (long long)(row0 * 128 + v0 + cg * 4); + unsigned int x_thr = xs_addr + (unsigned int)((row0 * 12 + cg * 4) * 4); + unsigned int x_rd = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd = ring_w + (unsigned int)(g * 16); + int cp_i = lane / 16; + int cp_j = lane % 16; + int cp_src = cp_i * 128 + warp_1 * 64 + 4 * cp_j; + unsigned int cp_dst = ring_w + (unsigned int)((cp_i * 64 + 4 * cp_j) * 4); + long long m_cur = 0; + if (e_begin < e_end) { + m_cur = elem_m[e_begin]; + long long src0 = m_cur + (long long)cp_src; + unsigned int dst0 = cp_dst + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0), "l"(m + src0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 512), "l"(m + (src0 + 256))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1024), "l"(m + (src0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 1536), "l"(m + (src0 + 768))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 2048), "l"(m + (src0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 2560), "l"(m + (src0 + 1280))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 3072), "l"(m + (src0 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0 + 3584), "l"(m + (src0 + 1792))); + asm volatile("cp.async.commit_group;"); + long long src0_0 = m_cur + (long long)(cp_src + 2048); + unsigned int dst0_0 = cp_dst + 5120; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0), "l"(m + src0_0)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 512), "l"(m + (src0_0 + 256))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1024), "l"(m + (src0_0 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 1536), "l"(m + (src0_0 + 768))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 2048), "l"(m + (src0_0 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 2560), "l"(m + (src0_0 + 1280))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 3072), "l"(m + (src0_0 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_0 + 3584), "l"(m + (src0_0 + 1792))); + asm volatile("cp.async.commit_group;"); + long long src0_1 = m_cur + (long long)(cp_src + 4096); + unsigned int dst0_1 = cp_dst + 10240; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1), "l"(m + src0_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 512), "l"(m + (src0_1 + 256))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1024), "l"(m + (src0_1 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 1536), "l"(m + (src0_1 + 768))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 2048), "l"(m + (src0_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 2560), "l"(m + (src0_1 + 1280))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 3072), "l"(m + (src0_1 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_1 + 3584), "l"(m + (src0_1 + 1792))); + asm volatile("cp.async.commit_group;"); + long long c_first = -1; + if (col_tile == 0) { + c_first = elem_c[e_begin]; + } + if (c_first >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 128); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 256); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first + g_off + 384); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + long long x0_off = x0_index[chain]; + if (x0_off >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 128); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 256); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off + g_off + 384); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + } else { + int diag = row0 - v0 - cg * 4; + if (col_tile == 0) { + diag = -4096; + } + acc[0] = ((diag == 0) ? 1.0f : 0.0f); + acc[1] = ((diag == 1) ? 1.0f : 0.0f); + acc[2] = ((diag == 2) ? 1.0f : 0.0f); + acc[3] = ((diag == 3) ? 1.0f : 0.0f); + acc[4] = ((diag + 1 == 0) ? 1.0f : 0.0f); + acc[5] = ((diag + 1 == 1) ? 1.0f : 0.0f); + acc[6] = ((diag + 1 == 2) ? 1.0f : 0.0f); + acc[7] = ((diag + 1 == 3) ? 1.0f : 0.0f); + acc[8] = ((diag + 2 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag + 2 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag + 2 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag + 2 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag + 3 == 0) ? 1.0f : 0.0f); + acc[13] = ((diag + 3 == 1) ? 1.0f : 0.0f); + acc[14] = ((diag + 3 == 2) ? 1.0f : 0.0f); + acc[15] = ((diag + 3 == 3) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 48), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 96), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + 144), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first = elem_y[e_begin]; + if (y_first >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first + g_off + 384) = _v4; + } + } + } + } else { + long long y_last = y_final[chain]; + if (y_last >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last + g_off + 384) = _v4; + } + } + } + __syncthreads(); + unsigned int buf = 0; + #pragma unroll 1 + for (int e = e_begin; e < e_end; e++) { + int has_next = 0; + long long m_next = 0; + long long c_next = -1; + long long y_pub = -1; + if (e + 1 < e_end) { + has_next = 1; + m_next = elem_m[e + 1]; + if (col_tile == 0) { + c_next = elem_c[e + 1]; + } + if (write_states != 0) { + y_pub = elem_y[e + 1]; + } + } else { + y_pub = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + unsigned int x_rd_e = x_rd + buf * 6144; + #pragma unroll 1 + for (int ch = 0; ch < 8; ch++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc = ch + 3; + if (nc < 8) { + long long src0_2 = m_cur + (long long)(cp_src + nc * 2048); + unsigned int dst0_2 = cp_dst + (unsigned int)(nc % 4 * 5120); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2), "l"(m + src0_2)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 512), "l"(m + (src0_2 + 256))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1024), "l"(m + (src0_2 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 1536), "l"(m + (src0_2 + 768))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 2048), "l"(m + (src0_2 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 2560), "l"(m + (src0_2 + 1280))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 3072), "l"(m + (src0_2 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2 + 3584), "l"(m + (src0_2 + 1792))); + } else if (has_next != 0) { + long long src0_3 = m_next + (long long)(cp_src + (nc - 8) * 2048); + unsigned int dst0_3 = cp_dst + (unsigned int)(nc % 4 * 5120); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3), "l"(m + src0_3)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 512), "l"(m + (src0_3 + 256))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1024), "l"(m + (src0_3 + 512))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 1536), "l"(m + (src0_3 + 768))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 2048), "l"(m + (src0_3 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 2560), "l"(m + (src0_3 + 1280))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 3072), "l"(m + (src0_3 + 1536))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3 + 3584), "l"(m + (src0_3 + 1792))); + } + asm volatile("cp.async.commit_group;"); + if (ch == 1) { + if (c_next >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 128); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 256); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next + g_off + 384); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + #pragma unroll + for (int kk = 0; kk < 16; kk++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd + (unsigned int)(ch % 4 * 5120) + (unsigned int)(kk * 256))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e + (unsigned int)(ch * 768) + (unsigned int)(kk * 48))); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[1], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[1], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[1], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[1], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[2], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[2], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[2], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[2], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[3], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[3], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[3], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[3], xv[3], acc[15]); +#else + float2 m2_0 = make_float2(mv[0], mv[0]); + float2 a2_0 = make_float2(acc[0], acc[1]); + a2_0 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[0], xv[1]), a2_0); + acc[0] = a2_0.x; + acc[1] = a2_0.y; + float2 a2_1 = make_float2(acc[2], acc[3]); + a2_1 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[2], xv[3]), a2_1); + acc[2] = a2_1.x; + acc[3] = a2_1.y; + float2 m2_1 = make_float2(mv[1], mv[1]); + float2 a2_2 = make_float2(acc[4], acc[5]); + a2_2 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[0], xv[1]), a2_2); + acc[4] = a2_2.x; + acc[5] = a2_2.y; + float2 a2_3 = make_float2(acc[6], acc[7]); + a2_3 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[2], xv[3]), a2_3); + acc[6] = a2_3.x; + acc[7] = a2_3.y; + float2 m2_2 = make_float2(mv[2], mv[2]); + float2 a2_4 = make_float2(acc[8], acc[9]); + a2_4 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[0], xv[1]), a2_4); + acc[8] = a2_4.x; + acc[9] = a2_4.y; + float2 a2_5 = make_float2(acc[10], acc[11]); + a2_5 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[2], xv[3]), a2_5); + acc[10] = a2_5.x; + acc[11] = a2_5.y; + float2 m2_3 = make_float2(mv[3], mv[3]); + float2 a2_6 = make_float2(acc[12], acc[13]); + a2_6 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[0], xv[1]), a2_6); + acc[12] = a2_6.x; + acc[13] = a2_6.y; + float2 a2_7 = make_float2(acc[14], acc[15]); + a2_7 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[2], xv[3]), a2_7); + acc[14] = a2_7.x; + acc[15] = a2_7.y; +#endif + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (6144 - buf * 6144)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (6144 - buf * 6144) + 48), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (6144 - buf * 6144) + 96), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (6144 - buf * 6144) + 144), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (y_pub >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub + g_off) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub + g_off + 128) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub + g_off + 256) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub + g_off + 384) = _v4; + } + } + __syncthreads(); + buf = 1 - buf; + m_cur = m_next; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } else { + int row0_1 = warp_1 * 64 + g; + long long g_off_1 = (long long)col_tile * 16384 + (long long)(row0_1 * 128 + v0 + cg * 4); + unsigned int x_thr_1 = xs_addr + (unsigned int)((row0_1 * 12 + cg * 4) * 4); + unsigned int x_rd_1 = xs_addr + (unsigned int)(cg * 16); + unsigned int m_rd_1 = ring_w + (unsigned int)(g * 80); + int cp_i_1 = lane / 4; + int cp_j_1 = lane % 4; + int cp_src_1 = cp_i_1 * 128 + warp_1 * 8192 + 4 * cp_j_1; + unsigned int cp_dst_1 = ring_w + (unsigned int)((cp_i_1 * 20 + 4 * cp_j_1) * 4); + long long m_cur_1 = 0; + if (e_begin < e_end) { + m_cur_1 = elem_m[e_begin]; + long long src0_4 = m_cur_1 + (long long)cp_src_1; + unsigned int dst0_4 = cp_dst_1 + 0; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4), "l"(m + src0_4)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 640), "l"(m + (src0_4 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1280), "l"(m + (src0_4 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 1920), "l"(m + (src0_4 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 2560), "l"(m + (src0_4 + 4096))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 3200), "l"(m + (src0_4 + 5120))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 3840), "l"(m + (src0_4 + 6144))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_4 + 4480), "l"(m + (src0_4 + 7168))); + asm volatile("cp.async.commit_group;"); + long long src0_2_1 = m_cur_1 + (long long)(cp_src_1 + 16); + unsigned int dst0_2_1 = cp_dst_1 + 5120; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1), "l"(m + src0_2_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 640), "l"(m + (src0_2_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1280), "l"(m + (src0_2_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 1920), "l"(m + (src0_2_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 2560), "l"(m + (src0_2_1 + 4096))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 3200), "l"(m + (src0_2_1 + 5120))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 3840), "l"(m + (src0_2_1 + 6144))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_2_1 + 4480), "l"(m + (src0_2_1 + 7168))); + asm volatile("cp.async.commit_group;"); + long long src0_3_1 = m_cur_1 + (long long)(cp_src_1 + 32); + unsigned int dst0_3_1 = cp_dst_1 + 10240; + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1), "l"(m + src0_3_1)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 640), "l"(m + (src0_3_1 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1280), "l"(m + (src0_3_1 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 1920), "l"(m + (src0_3_1 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 2560), "l"(m + (src0_3_1 + 4096))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 3200), "l"(m + (src0_3_1 + 5120))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 3840), "l"(m + (src0_3_1 + 6144))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_3_1 + 4480), "l"(m + (src0_3_1 + 7168))); + asm volatile("cp.async.commit_group;"); + long long c_first_1 = -1; + if (col_tile == 0) { + c_first_1 = elem_c[e_begin]; + } + if (c_first_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 2048); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 4096); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_first_1 + g_off_1 + 6144); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + long long x0_off_1 = x0_index[chain]; + if (x0_off_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1); + acc[0 + 0] = _v4.x; + acc[0 + 1] = _v4.y; + acc[0 + 2] = _v4.z; + acc[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 2048); + acc[4 + 0] = _v4.x; + acc[4 + 1] = _v4.y; + acc[4 + 2] = _v4.z; + acc[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 4096); + acc[8 + 0] = _v4.x; + acc[8 + 1] = _v4.y; + acc[8 + 2] = _v4.z; + acc[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(x0 + x0_off_1 + g_off_1 + 6144); + acc[12 + 0] = _v4.x; + acc[12 + 1] = _v4.y; + acc[12 + 2] = _v4.z; + acc[12 + 3] = _v4.w; + } + } else { + int diag_1 = row0_1 - v0 - cg * 4; + if (col_tile == 0) { + diag_1 = -4096; + } + acc[0] = ((diag_1 == 0) ? 1.0f : 0.0f); + acc[1] = ((diag_1 == 1) ? 1.0f : 0.0f); + acc[2] = ((diag_1 == 2) ? 1.0f : 0.0f); + acc[3] = ((diag_1 == 3) ? 1.0f : 0.0f); + acc[4] = ((diag_1 + 16 == 0) ? 1.0f : 0.0f); + acc[5] = ((diag_1 + 16 == 1) ? 1.0f : 0.0f); + acc[6] = ((diag_1 + 16 == 2) ? 1.0f : 0.0f); + acc[7] = ((diag_1 + 16 == 3) ? 1.0f : 0.0f); + acc[8] = ((diag_1 + 32 == 0) ? 1.0f : 0.0f); + acc[9] = ((diag_1 + 32 == 1) ? 1.0f : 0.0f); + acc[10] = ((diag_1 + 32 == 2) ? 1.0f : 0.0f); + acc[11] = ((diag_1 + 32 == 3) ? 1.0f : 0.0f); + acc[12] = ((diag_1 + 48 == 0) ? 1.0f : 0.0f); + acc[13] = ((diag_1 + 48 == 1) ? 1.0f : 0.0f); + acc[14] = ((diag_1 + 48 == 2) ? 1.0f : 0.0f); + acc[15] = ((diag_1 + 48 == 3) ? 1.0f : 0.0f); + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 768), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 1536), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + 2304), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (e_begin < e_end) { + if (write_states != 0) { + long long y_first_1 = elem_y[e_begin]; + if (y_first_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 4096) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_first_1 + g_off_1 + 6144) = _v4; + } + } + } + } else { + long long y_last_1 = y_final[chain]; + if (y_last_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 4096) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_last_1 + g_off_1 + 6144) = _v4; + } + } + } + __syncthreads(); + unsigned int buf_1 = 0; + #pragma unroll 1 + for (int e_1 = e_begin; e_1 < e_end; e_1++) { + int has_next_1 = 0; + long long m_next_1 = 0; + long long c_next_1 = -1; + long long y_pub_1 = -1; + if (e_1 + 1 < e_end) { + has_next_1 = 1; + m_next_1 = elem_m[e_1 + 1]; + if (col_tile == 0) { + c_next_1 = elem_c[e_1 + 1]; + } + if (write_states != 0) { + y_pub_1 = elem_y[e_1 + 1]; + } + } else { + y_pub_1 = y_final[chain]; + } + acc[0] = cn[0]; + acc[1] = cn[1]; + acc[2] = cn[2]; + acc[3] = cn[3]; + acc[4] = cn[4]; + acc[5] = cn[5]; + acc[6] = cn[6]; + acc[7] = cn[7]; + acc[8] = cn[8]; + acc[9] = cn[9]; + acc[10] = cn[10]; + acc[11] = cn[11]; + acc[12] = cn[12]; + acc[13] = cn[13]; + acc[14] = cn[14]; + acc[15] = cn[15]; + unsigned int x_rd_e_1 = x_rd_1 + buf_1 * 6144; + #pragma unroll 1 + for (int ch_1 = 0; ch_1 < 8; ch_1++) { + asm volatile("cp.async.wait_group 2;"); + __syncwarp(); + int nc_1 = ch_1 + 3; + if (nc_1 < 8) { + long long src0_5 = m_cur_1 + (long long)(cp_src_1 + nc_1 * 16); + unsigned int dst0_5 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 5120); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5), "l"(m + src0_5)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 640), "l"(m + (src0_5 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1280), "l"(m + (src0_5 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 1920), "l"(m + (src0_5 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 2560), "l"(m + (src0_5 + 4096))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 3200), "l"(m + (src0_5 + 5120))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 3840), "l"(m + (src0_5 + 6144))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_5 + 4480), "l"(m + (src0_5 + 7168))); + } else if (has_next_1 != 0) { + long long src0_6 = m_next_1 + (long long)(cp_src_1 + (nc_1 - 8) * 16); + unsigned int dst0_6 = cp_dst_1 + (unsigned int)(nc_1 % 4 * 5120); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6), "l"(m + src0_6)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 640), "l"(m + (src0_6 + 1024))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1280), "l"(m + (src0_6 + 2048))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 1920), "l"(m + (src0_6 + 3072))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 2560), "l"(m + (src0_6 + 4096))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 3200), "l"(m + (src0_6 + 5120))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 3840), "l"(m + (src0_6 + 6144))); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(dst0_6 + 4480), "l"(m + (src0_6 + 7168))); + } + asm volatile("cp.async.commit_group;"); + if (ch_1 == 1) { + if (c_next_1 >= 0) { + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1); + cn[0 + 0] = _v4.x; + cn[0 + 1] = _v4.y; + cn[0 + 2] = _v4.z; + cn[0 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 2048); + cn[4 + 0] = _v4.x; + cn[4 + 1] = _v4.y; + cn[4 + 2] = _v4.z; + cn[4 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 4096); + cn[8 + 0] = _v4.x; + cn[8 + 1] = _v4.y; + cn[8 + 2] = _v4.z; + cn[8 + 3] = _v4.w; + } + { + float4 _v4 = *reinterpret_cast(c + c_next_1 + g_off_1 + 6144); + cn[12 + 0] = _v4.x; + cn[12 + 1] = _v4.y; + cn[12 + 2] = _v4.z; + cn[12 + 3] = _v4.w; + } + } else { + cn[0] = 0.0f; + cn[1] = 0.0f; + cn[2] = 0.0f; + cn[3] = 0.0f; + cn[4] = 0.0f; + cn[5] = 0.0f; + cn[6] = 0.0f; + cn[7] = 0.0f; + cn[8] = 0.0f; + cn[9] = 0.0f; + cn[10] = 0.0f; + cn[11] = 0.0f; + cn[12] = 0.0f; + cn[13] = 0.0f; + cn[14] = 0.0f; + cn[15] = 0.0f; + } + } + #pragma unroll + for (int q = 0; q < 4; q++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[0])), "=r"(*reinterpret_cast(&mv[(0) + 1])), "=r"(*reinterpret_cast(&mv[(0) + 2])), "=r"(*reinterpret_cast(&mv[(0) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 5120) + (unsigned int)(q * 16))); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[4])), "=r"(*reinterpret_cast(&mv[(4) + 1])), "=r"(*reinterpret_cast(&mv[(4) + 2])), "=r"(*reinterpret_cast(&mv[(4) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 5120) + (unsigned int)(q * 16) + 1280)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[8])), "=r"(*reinterpret_cast(&mv[(8) + 1])), "=r"(*reinterpret_cast(&mv[(8) + 2])), "=r"(*reinterpret_cast(&mv[(8) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 5120) + (unsigned int)(q * 16) + 2560)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&mv[12])), "=r"(*reinterpret_cast(&mv[(12) + 1])), "=r"(*reinterpret_cast(&mv[(12) + 2])), "=r"(*reinterpret_cast(&mv[(12) + 3])) + : "r"(m_rd_1 + (unsigned int)(ch_1 % 4 * 5120) + (unsigned int)(q * 16) + 3840)); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 768) + (unsigned int)(q * 192))); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[4], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[4], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[4], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[4], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[8], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[8], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[8], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[8], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[12], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[12], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[12], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[12], xv[3], acc[15]); +#else + float2 m2_4 = make_float2(mv[0], mv[0]); + float2 a2_8 = make_float2(acc[0], acc[1]); + a2_8 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[0], xv[1]), a2_8); + acc[0] = a2_8.x; + acc[1] = a2_8.y; + float2 a2_9 = make_float2(acc[2], acc[3]); + a2_9 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[2], xv[3]), a2_9); + acc[2] = a2_9.x; + acc[3] = a2_9.y; + float2 m2_5 = make_float2(mv[4], mv[4]); + float2 a2_10 = make_float2(acc[4], acc[5]); + a2_10 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[0], xv[1]), a2_10); + acc[4] = a2_10.x; + acc[5] = a2_10.y; + float2 a2_11 = make_float2(acc[6], acc[7]); + a2_11 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[2], xv[3]), a2_11); + acc[6] = a2_11.x; + acc[7] = a2_11.y; + float2 m2_6 = make_float2(mv[8], mv[8]); + float2 a2_12 = make_float2(acc[8], acc[9]); + a2_12 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[0], xv[1]), a2_12); + acc[8] = a2_12.x; + acc[9] = a2_12.y; + float2 a2_13 = make_float2(acc[10], acc[11]); + a2_13 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[2], xv[3]), a2_13); + acc[10] = a2_13.x; + acc[11] = a2_13.y; + float2 m2_7 = make_float2(mv[12], mv[12]); + float2 a2_14 = make_float2(acc[12], acc[13]); + a2_14 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[0], xv[1]), a2_14); + acc[12] = a2_14.x; + acc[13] = a2_14.y; + float2 a2_15 = make_float2(acc[14], acc[15]); + a2_15 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[2], xv[3]), a2_15); + acc[14] = a2_15.x; + acc[15] = a2_15.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 768) + (unsigned int)(q * 192) + 48)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[1], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[1], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[1], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[1], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[5], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[5], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[5], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[5], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[9], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[9], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[9], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[9], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[13], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[13], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[13], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[13], xv[3], acc[15]); +#else + float2 m2_8 = make_float2(mv[1], mv[1]); + float2 a2_16 = make_float2(acc[0], acc[1]); + a2_16 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[0], xv[1]), a2_16); + acc[0] = a2_16.x; + acc[1] = a2_16.y; + float2 a2_17 = make_float2(acc[2], acc[3]); + a2_17 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[2], xv[3]), a2_17); + acc[2] = a2_17.x; + acc[3] = a2_17.y; + float2 m2_9 = make_float2(mv[5], mv[5]); + float2 a2_18 = make_float2(acc[4], acc[5]); + a2_18 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[0], xv[1]), a2_18); + acc[4] = a2_18.x; + acc[5] = a2_18.y; + float2 a2_19 = make_float2(acc[6], acc[7]); + a2_19 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[2], xv[3]), a2_19); + acc[6] = a2_19.x; + acc[7] = a2_19.y; + float2 m2_10 = make_float2(mv[9], mv[9]); + float2 a2_20 = make_float2(acc[8], acc[9]); + a2_20 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[0], xv[1]), a2_20); + acc[8] = a2_20.x; + acc[9] = a2_20.y; + float2 a2_21 = make_float2(acc[10], acc[11]); + a2_21 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[2], xv[3]), a2_21); + acc[10] = a2_21.x; + acc[11] = a2_21.y; + float2 m2_11 = make_float2(mv[13], mv[13]); + float2 a2_22 = make_float2(acc[12], acc[13]); + a2_22 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[0], xv[1]), a2_22); + acc[12] = a2_22.x; + acc[13] = a2_22.y; + float2 a2_23 = make_float2(acc[14], acc[15]); + a2_23 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[2], xv[3]), a2_23); + acc[14] = a2_23.x; + acc[15] = a2_23.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 768) + (unsigned int)(q * 192) + 96)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[2], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[2], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[2], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[2], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[6], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[6], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[6], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[6], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[10], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[10], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[10], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[10], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[14], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[14], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[14], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[14], xv[3], acc[15]); +#else + float2 m2_12 = make_float2(mv[2], mv[2]); + float2 a2_24 = make_float2(acc[0], acc[1]); + a2_24 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[0], xv[1]), a2_24); + acc[0] = a2_24.x; + acc[1] = a2_24.y; + float2 a2_25 = make_float2(acc[2], acc[3]); + a2_25 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[2], xv[3]), a2_25); + acc[2] = a2_25.x; + acc[3] = a2_25.y; + float2 m2_13 = make_float2(mv[6], mv[6]); + float2 a2_26 = make_float2(acc[4], acc[5]); + a2_26 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[0], xv[1]), a2_26); + acc[4] = a2_26.x; + acc[5] = a2_26.y; + float2 a2_27 = make_float2(acc[6], acc[7]); + a2_27 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[2], xv[3]), a2_27); + acc[6] = a2_27.x; + acc[7] = a2_27.y; + float2 m2_14 = make_float2(mv[10], mv[10]); + float2 a2_28 = make_float2(acc[8], acc[9]); + a2_28 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[0], xv[1]), a2_28); + acc[8] = a2_28.x; + acc[9] = a2_28.y; + float2 a2_29 = make_float2(acc[10], acc[11]); + a2_29 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[2], xv[3]), a2_29); + acc[10] = a2_29.x; + acc[11] = a2_29.y; + float2 m2_15 = make_float2(mv[14], mv[14]); + float2 a2_30 = make_float2(acc[12], acc[13]); + a2_30 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[0], xv[1]), a2_30); + acc[12] = a2_30.x; + acc[13] = a2_30.y; + float2 a2_31 = make_float2(acc[14], acc[15]); + a2_31 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[2], xv[3]), a2_31); + acc[14] = a2_31.x; + acc[15] = a2_31.y; +#endif + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) + : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 768) + (unsigned int)(q * 192) + 144)); +#if __CUDA_ARCH__ == 1070 + acc[0] = __fmaf_rn(mv[3], xv[0], acc[0]); + acc[1] = __fmaf_rn(mv[3], xv[1], acc[1]); + acc[2] = __fmaf_rn(mv[3], xv[2], acc[2]); + acc[3] = __fmaf_rn(mv[3], xv[3], acc[3]); + acc[4] = __fmaf_rn(mv[7], xv[0], acc[4]); + acc[5] = __fmaf_rn(mv[7], xv[1], acc[5]); + acc[6] = __fmaf_rn(mv[7], xv[2], acc[6]); + acc[7] = __fmaf_rn(mv[7], xv[3], acc[7]); + acc[8] = __fmaf_rn(mv[11], xv[0], acc[8]); + acc[9] = __fmaf_rn(mv[11], xv[1], acc[9]); + acc[10] = __fmaf_rn(mv[11], xv[2], acc[10]); + acc[11] = __fmaf_rn(mv[11], xv[3], acc[11]); + acc[12] = __fmaf_rn(mv[15], xv[0], acc[12]); + acc[13] = __fmaf_rn(mv[15], xv[1], acc[13]); + acc[14] = __fmaf_rn(mv[15], xv[2], acc[14]); + acc[15] = __fmaf_rn(mv[15], xv[3], acc[15]); +#else + float2 m2_16 = make_float2(mv[3], mv[3]); + float2 a2_32 = make_float2(acc[0], acc[1]); + a2_32 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[0], xv[1]), a2_32); + acc[0] = a2_32.x; + acc[1] = a2_32.y; + float2 a2_33 = make_float2(acc[2], acc[3]); + a2_33 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[2], xv[3]), a2_33); + acc[2] = a2_33.x; + acc[3] = a2_33.y; + float2 m2_17 = make_float2(mv[7], mv[7]); + float2 a2_34 = make_float2(acc[4], acc[5]); + a2_34 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[0], xv[1]), a2_34); + acc[4] = a2_34.x; + acc[5] = a2_34.y; + float2 a2_35 = make_float2(acc[6], acc[7]); + a2_35 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[2], xv[3]), a2_35); + acc[6] = a2_35.x; + acc[7] = a2_35.y; + float2 m2_18 = make_float2(mv[11], mv[11]); + float2 a2_36 = make_float2(acc[8], acc[9]); + a2_36 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[0], xv[1]), a2_36); + acc[8] = a2_36.x; + acc[9] = a2_36.y; + float2 a2_37 = make_float2(acc[10], acc[11]); + a2_37 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[2], xv[3]), a2_37); + acc[10] = a2_37.x; + acc[11] = a2_37.y; + float2 m2_19 = make_float2(mv[15], mv[15]); + float2 a2_38 = make_float2(acc[12], acc[13]); + a2_38 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[0], xv[1]), a2_38); + acc[12] = a2_38.x; + acc[13] = a2_38.y; + float2 a2_39 = make_float2(acc[14], acc[15]); + a2_39 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[2], xv[3]), a2_39); + acc[14] = a2_39.x; + acc[15] = a2_39.y; +#endif + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (6144 - buf_1 * 6144)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (6144 - buf_1 * 6144) + 768), "r"(__float_as_uint(acc[4])), "r"(__float_as_uint(acc[5])), "r"(__float_as_uint(acc[6])), "r"(__float_as_uint(acc[7])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (6144 - buf_1 * 6144) + 1536), "r"(__float_as_uint(acc[8])), "r"(__float_as_uint(acc[9])), "r"(__float_as_uint(acc[10])), "r"(__float_as_uint(acc[11])) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (6144 - buf_1 * 6144) + 2304), "r"(__float_as_uint(acc[12])), "r"(__float_as_uint(acc[13])), "r"(__float_as_uint(acc[14])), "r"(__float_as_uint(acc[15])) : "memory"); + if (y_pub_1 >= 0) { + { + float4 _v4 = make_float4(acc[0 + 0], acc[0 + 1], acc[0 + 2], acc[0 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1) = _v4; + } + { + float4 _v4 = make_float4(acc[4 + 0], acc[4 + 1], acc[4 + 2], acc[4 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 2048) = _v4; + } + { + float4 _v4 = make_float4(acc[8 + 0], acc[8 + 1], acc[8 + 2], acc[8 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 4096) = _v4; + } + { + float4 _v4 = make_float4(acc[12 + 0], acc[12 + 1], acc[12 + 2], acc[12 + 3]); + *reinterpret_cast(y + y_pub_1 + g_off_1 + 6144) = _v4; + } + } + __syncthreads(); + buf_1 = 1 - buf_1; + m_cur_1 = m_next_1; + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } +} + +} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu index 42c6117094a..f93dc9900de 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7ddde81b137fb82d27d5_kcp_kernel.cu @@ -50,8 +50,6 @@ __device__ __forceinline__ int make_warp_uniform(int x) { #define LAUNCH_MIN_BLOCKS 1 #include -#if __CUDA_ARCH__ == 1070 -#else // ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) @@ -65,7 +63,6 @@ __device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) "l"(*(const unsigned long long*)&c)); return r; } -#endif extern "C" { @@ -1030,72 +1027,6 @@ kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) : "r"(x_rd_e + (unsigned int)(ch * 8704) + (unsigned int)(kk * 272) + 128)); -#if __CUDA_ARCH__ == 1070 - acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); - acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); - acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); - acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); - acc[4] = __fmaf_rn(mv[0], xv[4], acc[4]); - acc[5] = __fmaf_rn(mv[0], xv[5], acc[5]); - acc[6] = __fmaf_rn(mv[0], xv[6], acc[6]); - acc[7] = __fmaf_rn(mv[0], xv[7], acc[7]); - acc[8] = __fmaf_rn(mv[1], xv[0], acc[8]); - acc[9] = __fmaf_rn(mv[1], xv[1], acc[9]); - acc[10] = __fmaf_rn(mv[1], xv[2], acc[10]); - acc[11] = __fmaf_rn(mv[1], xv[3], acc[11]); - acc[12] = __fmaf_rn(mv[1], xv[4], acc[12]); - acc[13] = __fmaf_rn(mv[1], xv[5], acc[13]); - acc[14] = __fmaf_rn(mv[1], xv[6], acc[14]); - acc[15] = __fmaf_rn(mv[1], xv[7], acc[15]); - acc[16] = __fmaf_rn(mv[2], xv[0], acc[16]); - acc[17] = __fmaf_rn(mv[2], xv[1], acc[17]); - acc[18] = __fmaf_rn(mv[2], xv[2], acc[18]); - acc[19] = __fmaf_rn(mv[2], xv[3], acc[19]); - acc[20] = __fmaf_rn(mv[2], xv[4], acc[20]); - acc[21] = __fmaf_rn(mv[2], xv[5], acc[21]); - acc[22] = __fmaf_rn(mv[2], xv[6], acc[22]); - acc[23] = __fmaf_rn(mv[2], xv[7], acc[23]); - acc[24] = __fmaf_rn(mv[3], xv[0], acc[24]); - acc[25] = __fmaf_rn(mv[3], xv[1], acc[25]); - acc[26] = __fmaf_rn(mv[3], xv[2], acc[26]); - acc[27] = __fmaf_rn(mv[3], xv[3], acc[27]); - acc[28] = __fmaf_rn(mv[3], xv[4], acc[28]); - acc[29] = __fmaf_rn(mv[3], xv[5], acc[29]); - acc[30] = __fmaf_rn(mv[3], xv[6], acc[30]); - acc[31] = __fmaf_rn(mv[3], xv[7], acc[31]); - acc[32] = __fmaf_rn(mv[4], xv[0], acc[32]); - acc[33] = __fmaf_rn(mv[4], xv[1], acc[33]); - acc[34] = __fmaf_rn(mv[4], xv[2], acc[34]); - acc[35] = __fmaf_rn(mv[4], xv[3], acc[35]); - acc[36] = __fmaf_rn(mv[4], xv[4], acc[36]); - acc[37] = __fmaf_rn(mv[4], xv[5], acc[37]); - acc[38] = __fmaf_rn(mv[4], xv[6], acc[38]); - acc[39] = __fmaf_rn(mv[4], xv[7], acc[39]); - acc[40] = __fmaf_rn(mv[5], xv[0], acc[40]); - acc[41] = __fmaf_rn(mv[5], xv[1], acc[41]); - acc[42] = __fmaf_rn(mv[5], xv[2], acc[42]); - acc[43] = __fmaf_rn(mv[5], xv[3], acc[43]); - acc[44] = __fmaf_rn(mv[5], xv[4], acc[44]); - acc[45] = __fmaf_rn(mv[5], xv[5], acc[45]); - acc[46] = __fmaf_rn(mv[5], xv[6], acc[46]); - acc[47] = __fmaf_rn(mv[5], xv[7], acc[47]); - acc[48] = __fmaf_rn(mv[6], xv[0], acc[48]); - acc[49] = __fmaf_rn(mv[6], xv[1], acc[49]); - acc[50] = __fmaf_rn(mv[6], xv[2], acc[50]); - acc[51] = __fmaf_rn(mv[6], xv[3], acc[51]); - acc[52] = __fmaf_rn(mv[6], xv[4], acc[52]); - acc[53] = __fmaf_rn(mv[6], xv[5], acc[53]); - acc[54] = __fmaf_rn(mv[6], xv[6], acc[54]); - acc[55] = __fmaf_rn(mv[6], xv[7], acc[55]); - acc[56] = __fmaf_rn(mv[7], xv[0], acc[56]); - acc[57] = __fmaf_rn(mv[7], xv[1], acc[57]); - acc[58] = __fmaf_rn(mv[7], xv[2], acc[58]); - acc[59] = __fmaf_rn(mv[7], xv[3], acc[59]); - acc[60] = __fmaf_rn(mv[7], xv[4], acc[60]); - acc[61] = __fmaf_rn(mv[7], xv[5], acc[61]); - acc[62] = __fmaf_rn(mv[7], xv[6], acc[62]); - acc[63] = __fmaf_rn(mv[7], xv[7], acc[63]); -#else float2 m2_0 = make_float2(mv[0], mv[0]); float2 a2_0 = make_float2(acc[0], acc[1]); a2_0 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[0], xv[1]), a2_0); @@ -1232,7 +1163,6 @@ kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* a2_31 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[6], xv[7]), a2_31); acc[62] = a2_31.x; acc[63] = a2_31.y; -#endif } } asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (34816 - buf * 34816)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); @@ -2259,72 +2189,6 @@ kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 128)); -#if __CUDA_ARCH__ == 1070 - acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); - acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); - acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); - acc[3] = __fmaf_rn(mv[0], xv[3], acc[3]); - acc[4] = __fmaf_rn(mv[0], xv[4], acc[4]); - acc[5] = __fmaf_rn(mv[0], xv[5], acc[5]); - acc[6] = __fmaf_rn(mv[0], xv[6], acc[6]); - acc[7] = __fmaf_rn(mv[0], xv[7], acc[7]); - acc[8] = __fmaf_rn(mv[4], xv[0], acc[8]); - acc[9] = __fmaf_rn(mv[4], xv[1], acc[9]); - acc[10] = __fmaf_rn(mv[4], xv[2], acc[10]); - acc[11] = __fmaf_rn(mv[4], xv[3], acc[11]); - acc[12] = __fmaf_rn(mv[4], xv[4], acc[12]); - acc[13] = __fmaf_rn(mv[4], xv[5], acc[13]); - acc[14] = __fmaf_rn(mv[4], xv[6], acc[14]); - acc[15] = __fmaf_rn(mv[4], xv[7], acc[15]); - acc[16] = __fmaf_rn(mv[8], xv[0], acc[16]); - acc[17] = __fmaf_rn(mv[8], xv[1], acc[17]); - acc[18] = __fmaf_rn(mv[8], xv[2], acc[18]); - acc[19] = __fmaf_rn(mv[8], xv[3], acc[19]); - acc[20] = __fmaf_rn(mv[8], xv[4], acc[20]); - acc[21] = __fmaf_rn(mv[8], xv[5], acc[21]); - acc[22] = __fmaf_rn(mv[8], xv[6], acc[22]); - acc[23] = __fmaf_rn(mv[8], xv[7], acc[23]); - acc[24] = __fmaf_rn(mv[12], xv[0], acc[24]); - acc[25] = __fmaf_rn(mv[12], xv[1], acc[25]); - acc[26] = __fmaf_rn(mv[12], xv[2], acc[26]); - acc[27] = __fmaf_rn(mv[12], xv[3], acc[27]); - acc[28] = __fmaf_rn(mv[12], xv[4], acc[28]); - acc[29] = __fmaf_rn(mv[12], xv[5], acc[29]); - acc[30] = __fmaf_rn(mv[12], xv[6], acc[30]); - acc[31] = __fmaf_rn(mv[12], xv[7], acc[31]); - acc[32] = __fmaf_rn(mv[16], xv[0], acc[32]); - acc[33] = __fmaf_rn(mv[16], xv[1], acc[33]); - acc[34] = __fmaf_rn(mv[16], xv[2], acc[34]); - acc[35] = __fmaf_rn(mv[16], xv[3], acc[35]); - acc[36] = __fmaf_rn(mv[16], xv[4], acc[36]); - acc[37] = __fmaf_rn(mv[16], xv[5], acc[37]); - acc[38] = __fmaf_rn(mv[16], xv[6], acc[38]); - acc[39] = __fmaf_rn(mv[16], xv[7], acc[39]); - acc[40] = __fmaf_rn(mv[20], xv[0], acc[40]); - acc[41] = __fmaf_rn(mv[20], xv[1], acc[41]); - acc[42] = __fmaf_rn(mv[20], xv[2], acc[42]); - acc[43] = __fmaf_rn(mv[20], xv[3], acc[43]); - acc[44] = __fmaf_rn(mv[20], xv[4], acc[44]); - acc[45] = __fmaf_rn(mv[20], xv[5], acc[45]); - acc[46] = __fmaf_rn(mv[20], xv[6], acc[46]); - acc[47] = __fmaf_rn(mv[20], xv[7], acc[47]); - acc[48] = __fmaf_rn(mv[24], xv[0], acc[48]); - acc[49] = __fmaf_rn(mv[24], xv[1], acc[49]); - acc[50] = __fmaf_rn(mv[24], xv[2], acc[50]); - acc[51] = __fmaf_rn(mv[24], xv[3], acc[51]); - acc[52] = __fmaf_rn(mv[24], xv[4], acc[52]); - acc[53] = __fmaf_rn(mv[24], xv[5], acc[53]); - acc[54] = __fmaf_rn(mv[24], xv[6], acc[54]); - acc[55] = __fmaf_rn(mv[24], xv[7], acc[55]); - acc[56] = __fmaf_rn(mv[28], xv[0], acc[56]); - acc[57] = __fmaf_rn(mv[28], xv[1], acc[57]); - acc[58] = __fmaf_rn(mv[28], xv[2], acc[58]); - acc[59] = __fmaf_rn(mv[28], xv[3], acc[59]); - acc[60] = __fmaf_rn(mv[28], xv[4], acc[60]); - acc[61] = __fmaf_rn(mv[28], xv[5], acc[61]); - acc[62] = __fmaf_rn(mv[28], xv[6], acc[62]); - acc[63] = __fmaf_rn(mv[28], xv[7], acc[63]); -#else float2 m2_8 = make_float2(mv[0], mv[0]); float2 a2_32 = make_float2(acc[0], acc[1]); a2_32 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[0], xv[1]), a2_32); @@ -2461,79 +2325,12 @@ kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* a2_63 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[6], xv[7]), a2_63); acc[62] = a2_63.x; acc[63] = a2_63.y; -#endif asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 272)); asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 400)); -#if __CUDA_ARCH__ == 1070 - acc[0] = __fmaf_rn(mv[1], xv[0], acc[0]); - acc[1] = __fmaf_rn(mv[1], xv[1], acc[1]); - acc[2] = __fmaf_rn(mv[1], xv[2], acc[2]); - acc[3] = __fmaf_rn(mv[1], xv[3], acc[3]); - acc[4] = __fmaf_rn(mv[1], xv[4], acc[4]); - acc[5] = __fmaf_rn(mv[1], xv[5], acc[5]); - acc[6] = __fmaf_rn(mv[1], xv[6], acc[6]); - acc[7] = __fmaf_rn(mv[1], xv[7], acc[7]); - acc[8] = __fmaf_rn(mv[5], xv[0], acc[8]); - acc[9] = __fmaf_rn(mv[5], xv[1], acc[9]); - acc[10] = __fmaf_rn(mv[5], xv[2], acc[10]); - acc[11] = __fmaf_rn(mv[5], xv[3], acc[11]); - acc[12] = __fmaf_rn(mv[5], xv[4], acc[12]); - acc[13] = __fmaf_rn(mv[5], xv[5], acc[13]); - acc[14] = __fmaf_rn(mv[5], xv[6], acc[14]); - acc[15] = __fmaf_rn(mv[5], xv[7], acc[15]); - acc[16] = __fmaf_rn(mv[9], xv[0], acc[16]); - acc[17] = __fmaf_rn(mv[9], xv[1], acc[17]); - acc[18] = __fmaf_rn(mv[9], xv[2], acc[18]); - acc[19] = __fmaf_rn(mv[9], xv[3], acc[19]); - acc[20] = __fmaf_rn(mv[9], xv[4], acc[20]); - acc[21] = __fmaf_rn(mv[9], xv[5], acc[21]); - acc[22] = __fmaf_rn(mv[9], xv[6], acc[22]); - acc[23] = __fmaf_rn(mv[9], xv[7], acc[23]); - acc[24] = __fmaf_rn(mv[13], xv[0], acc[24]); - acc[25] = __fmaf_rn(mv[13], xv[1], acc[25]); - acc[26] = __fmaf_rn(mv[13], xv[2], acc[26]); - acc[27] = __fmaf_rn(mv[13], xv[3], acc[27]); - acc[28] = __fmaf_rn(mv[13], xv[4], acc[28]); - acc[29] = __fmaf_rn(mv[13], xv[5], acc[29]); - acc[30] = __fmaf_rn(mv[13], xv[6], acc[30]); - acc[31] = __fmaf_rn(mv[13], xv[7], acc[31]); - acc[32] = __fmaf_rn(mv[17], xv[0], acc[32]); - acc[33] = __fmaf_rn(mv[17], xv[1], acc[33]); - acc[34] = __fmaf_rn(mv[17], xv[2], acc[34]); - acc[35] = __fmaf_rn(mv[17], xv[3], acc[35]); - acc[36] = __fmaf_rn(mv[17], xv[4], acc[36]); - acc[37] = __fmaf_rn(mv[17], xv[5], acc[37]); - acc[38] = __fmaf_rn(mv[17], xv[6], acc[38]); - acc[39] = __fmaf_rn(mv[17], xv[7], acc[39]); - acc[40] = __fmaf_rn(mv[21], xv[0], acc[40]); - acc[41] = __fmaf_rn(mv[21], xv[1], acc[41]); - acc[42] = __fmaf_rn(mv[21], xv[2], acc[42]); - acc[43] = __fmaf_rn(mv[21], xv[3], acc[43]); - acc[44] = __fmaf_rn(mv[21], xv[4], acc[44]); - acc[45] = __fmaf_rn(mv[21], xv[5], acc[45]); - acc[46] = __fmaf_rn(mv[21], xv[6], acc[46]); - acc[47] = __fmaf_rn(mv[21], xv[7], acc[47]); - acc[48] = __fmaf_rn(mv[25], xv[0], acc[48]); - acc[49] = __fmaf_rn(mv[25], xv[1], acc[49]); - acc[50] = __fmaf_rn(mv[25], xv[2], acc[50]); - acc[51] = __fmaf_rn(mv[25], xv[3], acc[51]); - acc[52] = __fmaf_rn(mv[25], xv[4], acc[52]); - acc[53] = __fmaf_rn(mv[25], xv[5], acc[53]); - acc[54] = __fmaf_rn(mv[25], xv[6], acc[54]); - acc[55] = __fmaf_rn(mv[25], xv[7], acc[55]); - acc[56] = __fmaf_rn(mv[29], xv[0], acc[56]); - acc[57] = __fmaf_rn(mv[29], xv[1], acc[57]); - acc[58] = __fmaf_rn(mv[29], xv[2], acc[58]); - acc[59] = __fmaf_rn(mv[29], xv[3], acc[59]); - acc[60] = __fmaf_rn(mv[29], xv[4], acc[60]); - acc[61] = __fmaf_rn(mv[29], xv[5], acc[61]); - acc[62] = __fmaf_rn(mv[29], xv[6], acc[62]); - acc[63] = __fmaf_rn(mv[29], xv[7], acc[63]); -#else float2 m2_16 = make_float2(mv[1], mv[1]); float2 a2_64 = make_float2(acc[0], acc[1]); a2_64 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[0], xv[1]), a2_64); @@ -2670,79 +2467,12 @@ kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* a2_95 = fma_f32x2_rn_ftz(m2_23, make_float2(xv[6], xv[7]), a2_95); acc[62] = a2_95.x; acc[63] = a2_95.y; -#endif asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 544)); asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 672)); -#if __CUDA_ARCH__ == 1070 - acc[0] = __fmaf_rn(mv[2], xv[0], acc[0]); - acc[1] = __fmaf_rn(mv[2], xv[1], acc[1]); - acc[2] = __fmaf_rn(mv[2], xv[2], acc[2]); - acc[3] = __fmaf_rn(mv[2], xv[3], acc[3]); - acc[4] = __fmaf_rn(mv[2], xv[4], acc[4]); - acc[5] = __fmaf_rn(mv[2], xv[5], acc[5]); - acc[6] = __fmaf_rn(mv[2], xv[6], acc[6]); - acc[7] = __fmaf_rn(mv[2], xv[7], acc[7]); - acc[8] = __fmaf_rn(mv[6], xv[0], acc[8]); - acc[9] = __fmaf_rn(mv[6], xv[1], acc[9]); - acc[10] = __fmaf_rn(mv[6], xv[2], acc[10]); - acc[11] = __fmaf_rn(mv[6], xv[3], acc[11]); - acc[12] = __fmaf_rn(mv[6], xv[4], acc[12]); - acc[13] = __fmaf_rn(mv[6], xv[5], acc[13]); - acc[14] = __fmaf_rn(mv[6], xv[6], acc[14]); - acc[15] = __fmaf_rn(mv[6], xv[7], acc[15]); - acc[16] = __fmaf_rn(mv[10], xv[0], acc[16]); - acc[17] = __fmaf_rn(mv[10], xv[1], acc[17]); - acc[18] = __fmaf_rn(mv[10], xv[2], acc[18]); - acc[19] = __fmaf_rn(mv[10], xv[3], acc[19]); - acc[20] = __fmaf_rn(mv[10], xv[4], acc[20]); - acc[21] = __fmaf_rn(mv[10], xv[5], acc[21]); - acc[22] = __fmaf_rn(mv[10], xv[6], acc[22]); - acc[23] = __fmaf_rn(mv[10], xv[7], acc[23]); - acc[24] = __fmaf_rn(mv[14], xv[0], acc[24]); - acc[25] = __fmaf_rn(mv[14], xv[1], acc[25]); - acc[26] = __fmaf_rn(mv[14], xv[2], acc[26]); - acc[27] = __fmaf_rn(mv[14], xv[3], acc[27]); - acc[28] = __fmaf_rn(mv[14], xv[4], acc[28]); - acc[29] = __fmaf_rn(mv[14], xv[5], acc[29]); - acc[30] = __fmaf_rn(mv[14], xv[6], acc[30]); - acc[31] = __fmaf_rn(mv[14], xv[7], acc[31]); - acc[32] = __fmaf_rn(mv[18], xv[0], acc[32]); - acc[33] = __fmaf_rn(mv[18], xv[1], acc[33]); - acc[34] = __fmaf_rn(mv[18], xv[2], acc[34]); - acc[35] = __fmaf_rn(mv[18], xv[3], acc[35]); - acc[36] = __fmaf_rn(mv[18], xv[4], acc[36]); - acc[37] = __fmaf_rn(mv[18], xv[5], acc[37]); - acc[38] = __fmaf_rn(mv[18], xv[6], acc[38]); - acc[39] = __fmaf_rn(mv[18], xv[7], acc[39]); - acc[40] = __fmaf_rn(mv[22], xv[0], acc[40]); - acc[41] = __fmaf_rn(mv[22], xv[1], acc[41]); - acc[42] = __fmaf_rn(mv[22], xv[2], acc[42]); - acc[43] = __fmaf_rn(mv[22], xv[3], acc[43]); - acc[44] = __fmaf_rn(mv[22], xv[4], acc[44]); - acc[45] = __fmaf_rn(mv[22], xv[5], acc[45]); - acc[46] = __fmaf_rn(mv[22], xv[6], acc[46]); - acc[47] = __fmaf_rn(mv[22], xv[7], acc[47]); - acc[48] = __fmaf_rn(mv[26], xv[0], acc[48]); - acc[49] = __fmaf_rn(mv[26], xv[1], acc[49]); - acc[50] = __fmaf_rn(mv[26], xv[2], acc[50]); - acc[51] = __fmaf_rn(mv[26], xv[3], acc[51]); - acc[52] = __fmaf_rn(mv[26], xv[4], acc[52]); - acc[53] = __fmaf_rn(mv[26], xv[5], acc[53]); - acc[54] = __fmaf_rn(mv[26], xv[6], acc[54]); - acc[55] = __fmaf_rn(mv[26], xv[7], acc[55]); - acc[56] = __fmaf_rn(mv[30], xv[0], acc[56]); - acc[57] = __fmaf_rn(mv[30], xv[1], acc[57]); - acc[58] = __fmaf_rn(mv[30], xv[2], acc[58]); - acc[59] = __fmaf_rn(mv[30], xv[3], acc[59]); - acc[60] = __fmaf_rn(mv[30], xv[4], acc[60]); - acc[61] = __fmaf_rn(mv[30], xv[5], acc[61]); - acc[62] = __fmaf_rn(mv[30], xv[6], acc[62]); - acc[63] = __fmaf_rn(mv[30], xv[7], acc[63]); -#else float2 m2_24 = make_float2(mv[2], mv[2]); float2 a2_96 = make_float2(acc[0], acc[1]); a2_96 = fma_f32x2_rn_ftz(m2_24, make_float2(xv[0], xv[1]), a2_96); @@ -2879,79 +2609,12 @@ kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* a2_127 = fma_f32x2_rn_ftz(m2_31, make_float2(xv[6], xv[7]), a2_127); acc[62] = a2_127.x; acc[63] = a2_127.y; -#endif asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 816)); asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[4])), "=r"(*reinterpret_cast(&xv[(4) + 1])), "=r"(*reinterpret_cast(&xv[(4) + 2])), "=r"(*reinterpret_cast(&xv[(4) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 8704) + (unsigned int)(q * 1088) + 944)); -#if __CUDA_ARCH__ == 1070 - acc[0] = __fmaf_rn(mv[3], xv[0], acc[0]); - acc[1] = __fmaf_rn(mv[3], xv[1], acc[1]); - acc[2] = __fmaf_rn(mv[3], xv[2], acc[2]); - acc[3] = __fmaf_rn(mv[3], xv[3], acc[3]); - acc[4] = __fmaf_rn(mv[3], xv[4], acc[4]); - acc[5] = __fmaf_rn(mv[3], xv[5], acc[5]); - acc[6] = __fmaf_rn(mv[3], xv[6], acc[6]); - acc[7] = __fmaf_rn(mv[3], xv[7], acc[7]); - acc[8] = __fmaf_rn(mv[7], xv[0], acc[8]); - acc[9] = __fmaf_rn(mv[7], xv[1], acc[9]); - acc[10] = __fmaf_rn(mv[7], xv[2], acc[10]); - acc[11] = __fmaf_rn(mv[7], xv[3], acc[11]); - acc[12] = __fmaf_rn(mv[7], xv[4], acc[12]); - acc[13] = __fmaf_rn(mv[7], xv[5], acc[13]); - acc[14] = __fmaf_rn(mv[7], xv[6], acc[14]); - acc[15] = __fmaf_rn(mv[7], xv[7], acc[15]); - acc[16] = __fmaf_rn(mv[11], xv[0], acc[16]); - acc[17] = __fmaf_rn(mv[11], xv[1], acc[17]); - acc[18] = __fmaf_rn(mv[11], xv[2], acc[18]); - acc[19] = __fmaf_rn(mv[11], xv[3], acc[19]); - acc[20] = __fmaf_rn(mv[11], xv[4], acc[20]); - acc[21] = __fmaf_rn(mv[11], xv[5], acc[21]); - acc[22] = __fmaf_rn(mv[11], xv[6], acc[22]); - acc[23] = __fmaf_rn(mv[11], xv[7], acc[23]); - acc[24] = __fmaf_rn(mv[15], xv[0], acc[24]); - acc[25] = __fmaf_rn(mv[15], xv[1], acc[25]); - acc[26] = __fmaf_rn(mv[15], xv[2], acc[26]); - acc[27] = __fmaf_rn(mv[15], xv[3], acc[27]); - acc[28] = __fmaf_rn(mv[15], xv[4], acc[28]); - acc[29] = __fmaf_rn(mv[15], xv[5], acc[29]); - acc[30] = __fmaf_rn(mv[15], xv[6], acc[30]); - acc[31] = __fmaf_rn(mv[15], xv[7], acc[31]); - acc[32] = __fmaf_rn(mv[19], xv[0], acc[32]); - acc[33] = __fmaf_rn(mv[19], xv[1], acc[33]); - acc[34] = __fmaf_rn(mv[19], xv[2], acc[34]); - acc[35] = __fmaf_rn(mv[19], xv[3], acc[35]); - acc[36] = __fmaf_rn(mv[19], xv[4], acc[36]); - acc[37] = __fmaf_rn(mv[19], xv[5], acc[37]); - acc[38] = __fmaf_rn(mv[19], xv[6], acc[38]); - acc[39] = __fmaf_rn(mv[19], xv[7], acc[39]); - acc[40] = __fmaf_rn(mv[23], xv[0], acc[40]); - acc[41] = __fmaf_rn(mv[23], xv[1], acc[41]); - acc[42] = __fmaf_rn(mv[23], xv[2], acc[42]); - acc[43] = __fmaf_rn(mv[23], xv[3], acc[43]); - acc[44] = __fmaf_rn(mv[23], xv[4], acc[44]); - acc[45] = __fmaf_rn(mv[23], xv[5], acc[45]); - acc[46] = __fmaf_rn(mv[23], xv[6], acc[46]); - acc[47] = __fmaf_rn(mv[23], xv[7], acc[47]); - acc[48] = __fmaf_rn(mv[27], xv[0], acc[48]); - acc[49] = __fmaf_rn(mv[27], xv[1], acc[49]); - acc[50] = __fmaf_rn(mv[27], xv[2], acc[50]); - acc[51] = __fmaf_rn(mv[27], xv[3], acc[51]); - acc[52] = __fmaf_rn(mv[27], xv[4], acc[52]); - acc[53] = __fmaf_rn(mv[27], xv[5], acc[53]); - acc[54] = __fmaf_rn(mv[27], xv[6], acc[54]); - acc[55] = __fmaf_rn(mv[27], xv[7], acc[55]); - acc[56] = __fmaf_rn(mv[31], xv[0], acc[56]); - acc[57] = __fmaf_rn(mv[31], xv[1], acc[57]); - acc[58] = __fmaf_rn(mv[31], xv[2], acc[58]); - acc[59] = __fmaf_rn(mv[31], xv[3], acc[59]); - acc[60] = __fmaf_rn(mv[31], xv[4], acc[60]); - acc[61] = __fmaf_rn(mv[31], xv[5], acc[61]); - acc[62] = __fmaf_rn(mv[31], xv[6], acc[62]); - acc[63] = __fmaf_rn(mv[31], xv[7], acc[63]); -#else float2 m2_32 = make_float2(mv[3], mv[3]); float2 a2_128 = make_float2(acc[0], acc[1]); a2_128 = fma_f32x2_rn_ftz(m2_32, make_float2(xv[0], xv[1]), a2_128); @@ -3088,7 +2751,6 @@ kernel_cake_kda_det_train_7ddde81b137fb82d27d5_kcp(float* __restrict__ m, float* a2_159 = fma_f32x2_rn_ftz(m2_39, make_float2(xv[6], xv[7]), a2_159); acc[62] = a2_159.x; acc[63] = a2_159.y; -#endif } } asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (34816 - buf_1 * 34816)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu similarity index 92% rename from flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu rename to flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu index c065412316b..81812df5241 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_binding.cu +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_15e81414b27b4b831808_kcp, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); +extern "C" __global__ void kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta); -namespace cake_binding_kernel_cake_kda_det_train_15e81414b27b4b831808_kcp { +namespace cake_binding_kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_15e81414b27b4b831808_kcp"; +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -122,7 +122,7 @@ void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, Tensor static thread_local uint64_t prepared_devices = 0; const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_15e81414b27b4b831808_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 61440), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 61440), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -146,11 +146,11 @@ void Run(TensorView a_m, TensorView a_c, TensorView a_x0, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 61440; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_15e81414b27b4b831808_kcp), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_kda_det_train_15e81414b27b4b831808_kcp +} // namespace cake_binding_kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_15e81414b27b4b831808_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu similarity index 80% rename from flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu rename to flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu index 72146109f05..40682eea47f 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_15e81414b27b4b831808_kcp_kernel.cu +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9ab870ac43ee331f2dd3_kcp_kernel.cu @@ -50,11 +50,27 @@ __device__ __forceinline__ int make_warp_uniform(int x) { #define LAUNCH_MIN_BLOCKS 3 #include +#if __CUDA_ARCH__ == 1070 +#else + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} +#endif extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) +kernel_cake_kda_det_train_9ab870ac43ee331f2dd3_kcp(float* __restrict__ m, float* __restrict__ c, float* __restrict__ x0, float* __restrict__ y, int* __restrict__ chain_start, long long* __restrict__ elem_m, long long* __restrict__ elem_c, long long* __restrict__ elem_y, long long* __restrict__ x0_index, long long* __restrict__ y_final, int ncols, int transpose_m, int write_states, int cols_per_cta) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -416,6 +432,7 @@ kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e + (unsigned int)(ch * 1280) + (unsigned int)(kk * 80))); +#if __CUDA_ARCH__ == 1070 acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); @@ -432,6 +449,44 @@ kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* acc[13] = __fmaf_rn(mv[3], xv[1], acc[13]); acc[14] = __fmaf_rn(mv[3], xv[2], acc[14]); acc[15] = __fmaf_rn(mv[3], xv[3], acc[15]); +#else + float2 m2_0 = make_float2(mv[0], mv[0]); + float2 a2_0 = make_float2(acc[0], acc[1]); + a2_0 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[0], xv[1]), a2_0); + acc[0] = a2_0.x; + acc[1] = a2_0.y; + float2 a2_1 = make_float2(acc[2], acc[3]); + a2_1 = fma_f32x2_rn_ftz(m2_0, make_float2(xv[2], xv[3]), a2_1); + acc[2] = a2_1.x; + acc[3] = a2_1.y; + float2 m2_1 = make_float2(mv[1], mv[1]); + float2 a2_2 = make_float2(acc[4], acc[5]); + a2_2 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[0], xv[1]), a2_2); + acc[4] = a2_2.x; + acc[5] = a2_2.y; + float2 a2_3 = make_float2(acc[6], acc[7]); + a2_3 = fma_f32x2_rn_ftz(m2_1, make_float2(xv[2], xv[3]), a2_3); + acc[6] = a2_3.x; + acc[7] = a2_3.y; + float2 m2_2 = make_float2(mv[2], mv[2]); + float2 a2_4 = make_float2(acc[8], acc[9]); + a2_4 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[0], xv[1]), a2_4); + acc[8] = a2_4.x; + acc[9] = a2_4.y; + float2 a2_5 = make_float2(acc[10], acc[11]); + a2_5 = fma_f32x2_rn_ftz(m2_2, make_float2(xv[2], xv[3]), a2_5); + acc[10] = a2_5.x; + acc[11] = a2_5.y; + float2 m2_3 = make_float2(mv[3], mv[3]); + float2 a2_6 = make_float2(acc[12], acc[13]); + a2_6 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[0], xv[1]), a2_6); + acc[12] = a2_6.x; + acc[13] = a2_6.y; + float2 a2_7 = make_float2(acc[14], acc[15]); + a2_7 = fma_f32x2_rn_ftz(m2_3, make_float2(xv[2], xv[3]), a2_7); + acc[14] = a2_7.x; + acc[15] = a2_7.y; +#endif } } asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr + (10240 - buf * 10240)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); @@ -791,6 +846,7 @@ kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320))); +#if __CUDA_ARCH__ == 1070 acc[0] = __fmaf_rn(mv[0], xv[0], acc[0]); acc[1] = __fmaf_rn(mv[0], xv[1], acc[1]); acc[2] = __fmaf_rn(mv[0], xv[2], acc[2]); @@ -807,9 +863,48 @@ kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* acc[13] = __fmaf_rn(mv[12], xv[1], acc[13]); acc[14] = __fmaf_rn(mv[12], xv[2], acc[14]); acc[15] = __fmaf_rn(mv[12], xv[3], acc[15]); +#else + float2 m2_4 = make_float2(mv[0], mv[0]); + float2 a2_8 = make_float2(acc[0], acc[1]); + a2_8 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[0], xv[1]), a2_8); + acc[0] = a2_8.x; + acc[1] = a2_8.y; + float2 a2_9 = make_float2(acc[2], acc[3]); + a2_9 = fma_f32x2_rn_ftz(m2_4, make_float2(xv[2], xv[3]), a2_9); + acc[2] = a2_9.x; + acc[3] = a2_9.y; + float2 m2_5 = make_float2(mv[4], mv[4]); + float2 a2_10 = make_float2(acc[4], acc[5]); + a2_10 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[0], xv[1]), a2_10); + acc[4] = a2_10.x; + acc[5] = a2_10.y; + float2 a2_11 = make_float2(acc[6], acc[7]); + a2_11 = fma_f32x2_rn_ftz(m2_5, make_float2(xv[2], xv[3]), a2_11); + acc[6] = a2_11.x; + acc[7] = a2_11.y; + float2 m2_6 = make_float2(mv[8], mv[8]); + float2 a2_12 = make_float2(acc[8], acc[9]); + a2_12 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[0], xv[1]), a2_12); + acc[8] = a2_12.x; + acc[9] = a2_12.y; + float2 a2_13 = make_float2(acc[10], acc[11]); + a2_13 = fma_f32x2_rn_ftz(m2_6, make_float2(xv[2], xv[3]), a2_13); + acc[10] = a2_13.x; + acc[11] = a2_13.y; + float2 m2_7 = make_float2(mv[12], mv[12]); + float2 a2_14 = make_float2(acc[12], acc[13]); + a2_14 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[0], xv[1]), a2_14); + acc[12] = a2_14.x; + acc[13] = a2_14.y; + float2 a2_15 = make_float2(acc[14], acc[15]); + a2_15 = fma_f32x2_rn_ftz(m2_7, make_float2(xv[2], xv[3]), a2_15); + acc[14] = a2_15.x; + acc[15] = a2_15.y; +#endif asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320) + 80)); +#if __CUDA_ARCH__ == 1070 acc[0] = __fmaf_rn(mv[1], xv[0], acc[0]); acc[1] = __fmaf_rn(mv[1], xv[1], acc[1]); acc[2] = __fmaf_rn(mv[1], xv[2], acc[2]); @@ -826,9 +921,48 @@ kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* acc[13] = __fmaf_rn(mv[13], xv[1], acc[13]); acc[14] = __fmaf_rn(mv[13], xv[2], acc[14]); acc[15] = __fmaf_rn(mv[13], xv[3], acc[15]); +#else + float2 m2_8 = make_float2(mv[1], mv[1]); + float2 a2_16 = make_float2(acc[0], acc[1]); + a2_16 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[0], xv[1]), a2_16); + acc[0] = a2_16.x; + acc[1] = a2_16.y; + float2 a2_17 = make_float2(acc[2], acc[3]); + a2_17 = fma_f32x2_rn_ftz(m2_8, make_float2(xv[2], xv[3]), a2_17); + acc[2] = a2_17.x; + acc[3] = a2_17.y; + float2 m2_9 = make_float2(mv[5], mv[5]); + float2 a2_18 = make_float2(acc[4], acc[5]); + a2_18 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[0], xv[1]), a2_18); + acc[4] = a2_18.x; + acc[5] = a2_18.y; + float2 a2_19 = make_float2(acc[6], acc[7]); + a2_19 = fma_f32x2_rn_ftz(m2_9, make_float2(xv[2], xv[3]), a2_19); + acc[6] = a2_19.x; + acc[7] = a2_19.y; + float2 m2_10 = make_float2(mv[9], mv[9]); + float2 a2_20 = make_float2(acc[8], acc[9]); + a2_20 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[0], xv[1]), a2_20); + acc[8] = a2_20.x; + acc[9] = a2_20.y; + float2 a2_21 = make_float2(acc[10], acc[11]); + a2_21 = fma_f32x2_rn_ftz(m2_10, make_float2(xv[2], xv[3]), a2_21); + acc[10] = a2_21.x; + acc[11] = a2_21.y; + float2 m2_11 = make_float2(mv[13], mv[13]); + float2 a2_22 = make_float2(acc[12], acc[13]); + a2_22 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[0], xv[1]), a2_22); + acc[12] = a2_22.x; + acc[13] = a2_22.y; + float2 a2_23 = make_float2(acc[14], acc[15]); + a2_23 = fma_f32x2_rn_ftz(m2_11, make_float2(xv[2], xv[3]), a2_23); + acc[14] = a2_23.x; + acc[15] = a2_23.y; +#endif asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320) + 160)); +#if __CUDA_ARCH__ == 1070 acc[0] = __fmaf_rn(mv[2], xv[0], acc[0]); acc[1] = __fmaf_rn(mv[2], xv[1], acc[1]); acc[2] = __fmaf_rn(mv[2], xv[2], acc[2]); @@ -845,9 +979,48 @@ kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* acc[13] = __fmaf_rn(mv[14], xv[1], acc[13]); acc[14] = __fmaf_rn(mv[14], xv[2], acc[14]); acc[15] = __fmaf_rn(mv[14], xv[3], acc[15]); +#else + float2 m2_12 = make_float2(mv[2], mv[2]); + float2 a2_24 = make_float2(acc[0], acc[1]); + a2_24 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[0], xv[1]), a2_24); + acc[0] = a2_24.x; + acc[1] = a2_24.y; + float2 a2_25 = make_float2(acc[2], acc[3]); + a2_25 = fma_f32x2_rn_ftz(m2_12, make_float2(xv[2], xv[3]), a2_25); + acc[2] = a2_25.x; + acc[3] = a2_25.y; + float2 m2_13 = make_float2(mv[6], mv[6]); + float2 a2_26 = make_float2(acc[4], acc[5]); + a2_26 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[0], xv[1]), a2_26); + acc[4] = a2_26.x; + acc[5] = a2_26.y; + float2 a2_27 = make_float2(acc[6], acc[7]); + a2_27 = fma_f32x2_rn_ftz(m2_13, make_float2(xv[2], xv[3]), a2_27); + acc[6] = a2_27.x; + acc[7] = a2_27.y; + float2 m2_14 = make_float2(mv[10], mv[10]); + float2 a2_28 = make_float2(acc[8], acc[9]); + a2_28 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[0], xv[1]), a2_28); + acc[8] = a2_28.x; + acc[9] = a2_28.y; + float2 a2_29 = make_float2(acc[10], acc[11]); + a2_29 = fma_f32x2_rn_ftz(m2_14, make_float2(xv[2], xv[3]), a2_29); + acc[10] = a2_29.x; + acc[11] = a2_29.y; + float2 m2_15 = make_float2(mv[14], mv[14]); + float2 a2_30 = make_float2(acc[12], acc[13]); + a2_30 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[0], xv[1]), a2_30); + acc[12] = a2_30.x; + acc[13] = a2_30.y; + float2 a2_31 = make_float2(acc[14], acc[15]); + a2_31 = fma_f32x2_rn_ftz(m2_15, make_float2(xv[2], xv[3]), a2_31); + acc[14] = a2_31.x; + acc[15] = a2_31.y; +#endif asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" : "=r"(*reinterpret_cast(&xv[0])), "=r"(*reinterpret_cast(&xv[(0) + 1])), "=r"(*reinterpret_cast(&xv[(0) + 2])), "=r"(*reinterpret_cast(&xv[(0) + 3])) : "r"(x_rd_e_1 + (unsigned int)(ch_1 * 1280) + (unsigned int)(q * 320) + 240)); +#if __CUDA_ARCH__ == 1070 acc[0] = __fmaf_rn(mv[3], xv[0], acc[0]); acc[1] = __fmaf_rn(mv[3], xv[1], acc[1]); acc[2] = __fmaf_rn(mv[3], xv[2], acc[2]); @@ -864,6 +1037,44 @@ kernel_cake_kda_det_train_15e81414b27b4b831808_kcp(float* __restrict__ m, float* acc[13] = __fmaf_rn(mv[15], xv[1], acc[13]); acc[14] = __fmaf_rn(mv[15], xv[2], acc[14]); acc[15] = __fmaf_rn(mv[15], xv[3], acc[15]); +#else + float2 m2_16 = make_float2(mv[3], mv[3]); + float2 a2_32 = make_float2(acc[0], acc[1]); + a2_32 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[0], xv[1]), a2_32); + acc[0] = a2_32.x; + acc[1] = a2_32.y; + float2 a2_33 = make_float2(acc[2], acc[3]); + a2_33 = fma_f32x2_rn_ftz(m2_16, make_float2(xv[2], xv[3]), a2_33); + acc[2] = a2_33.x; + acc[3] = a2_33.y; + float2 m2_17 = make_float2(mv[7], mv[7]); + float2 a2_34 = make_float2(acc[4], acc[5]); + a2_34 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[0], xv[1]), a2_34); + acc[4] = a2_34.x; + acc[5] = a2_34.y; + float2 a2_35 = make_float2(acc[6], acc[7]); + a2_35 = fma_f32x2_rn_ftz(m2_17, make_float2(xv[2], xv[3]), a2_35); + acc[6] = a2_35.x; + acc[7] = a2_35.y; + float2 m2_18 = make_float2(mv[11], mv[11]); + float2 a2_36 = make_float2(acc[8], acc[9]); + a2_36 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[0], xv[1]), a2_36); + acc[8] = a2_36.x; + acc[9] = a2_36.y; + float2 a2_37 = make_float2(acc[10], acc[11]); + a2_37 = fma_f32x2_rn_ftz(m2_18, make_float2(xv[2], xv[3]), a2_37); + acc[10] = a2_37.x; + acc[11] = a2_37.y; + float2 m2_19 = make_float2(mv[15], mv[15]); + float2 a2_38 = make_float2(acc[12], acc[13]); + a2_38 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[0], xv[1]), a2_38); + acc[12] = a2_38.x; + acc[13] = a2_38.y; + float2 a2_39 = make_float2(acc[14], acc[15]); + a2_39 = fma_f32x2_rn_ftz(m2_19, make_float2(xv[2], xv[3]), a2_39); + acc[14] = a2_39.x; + acc[15] = a2_39.y; +#endif } } asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_thr_1 + (10240 - buf_1 * 10240)), "r"(__float_as_uint(acc[0])), "r"(__float_as_uint(acc[1])), "r"(__float_as_uint(acc[2])), "r"(__float_as_uint(acc[3])) : "memory"); From 1acdc6d37c2232a98abd2d3fcd75eef30dfc69fc Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 17:50:06 -0800 Subject: [PATCH 05/10] cake_kda_train_kcp: the mirror tests leave the affine merge builds out The deterministic package now ships its own register-blocked merge builds; this package keeps its column-block c8/c16 merge builds. The mirror tests compare every other record and unit, and pin this package's merge records. --- .../test_cake_kda_train_kcp_build.py | 47 ++++++++++++++++--- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/tests/experimental/test_cake_kda_train_kcp_build.py b/tests/experimental/test_cake_kda_train_kcp_build.py index 22880f7601d..aef261be949 100644 --- a/tests/experimental/test_cake_kda_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_train_kcp_build.py @@ -490,6 +490,22 @@ def test_test_manifest_is_the_test_only_programs_under_the_opt_in( DET_PREFIX, PREFIX = "cake_kda_det_train_", "cake_kda_train_" +# The affine merge is the one program whose builds differ between the two packages: the deterministic package ships +# register-blocked merge builds selected by its MERGE_RULE; this package keeps its two column-block builds (c8, c16). +# Every other record and unit is the deterministic package's under this package's names. +MERGE = "kcp_affine_merge" +OWN_MERGE = { + "kcp_affine_merge_c8": PREFIX + "c6958a0a4ed5cb55a9a1_kcp", + "kcp_affine_merge_c16": PREFIX + "89f5ff59f03eb170e3bf_kcp", +} + + +def _merge_records(table: dict) -> set[str]: + return {name for name, record in table.items() if record["program"] == MERGE} + + +def _library_stem(name: str) -> str: + return name.rsplit("_sm_", 1)[0] def _renamed(value): @@ -506,13 +522,21 @@ def _renamed(value): def test_the_records_are_the_deterministic_records_under_this_packages_names() -> None: """The KCP programs have no non-deterministic form: both tables hold the deterministic package's records, every field equal but the library stem and the kernel symbol, which carry the same twenty hex digits under this - package's prefix; no library or kernel symbol of this package is one of the deterministic package's.""" + package's prefix; no library or kernel symbol of this package is one of the deterministic package's. The affine + merge records are this package's own column-block builds.""" for table in ("MODULES", "TEST_MODULES"): ours, theirs = getattr(kcp_jit, table), getattr(det_kcp_jit, table) - assert set(ours) == set(theirs), table + assert set(ours) - _merge_records(ours) == set(theirs) - _merge_records( + theirs + ), table for name, record in ours.items(): assert record["library"].startswith(PREFIX), name - assert record == _renamed(theirs[name]), (table, name) + if record["program"] != MERGE: + assert record == _renamed(theirs[name]), (table, name) + assert { + name: kcp_jit.MODULES[name]["library"] + for name in _merge_records(kcp_jit.MODULES) + } == OWN_MERGE ours = {r["library"] for r in _records().values()} theirs = { r["library"] @@ -530,14 +554,25 @@ def test_the_units_are_the_deterministic_units_under_this_packages_names( ) -> None: """Every library of this package's manifest is the deterministic package's library of the same digits: the same architecture, flags and entry, and sources whose text is the deterministic unit's with the deterministic prefix - renamed (file names, kernel symbols, binding namespaces); nothing else differs.""" + renamed (file names, kernel symbols, binding namespaces); nothing else differs. The affine merge libraries, which + differ between the two packages, are left out on both sides.""" det_dir = Path(det_kcp_jit.__file__).resolve().parent path, det_path = PACKAGE_DIR / manifest, det_dir / manifest if not path.is_file() or not det_path.is_file(): pytest.skip(f"{manifest} is not rendered in both packages here") - ours = {e["name"]: e for e in json.loads(path.read_text())["libraries"]} + own_merge = {kcp_jit.MODULES[n]["library"] for n in _merge_records(kcp_jit.MODULES)} + det_merge = { + det_kcp_jit.MODULES[n]["library"] for n in _merge_records(det_kcp_jit.MODULES) + } + ours = { + e["name"]: e + for e in json.loads(path.read_text())["libraries"] + if _library_stem(e["name"]) not in own_merge + } theirs = { - _renamed(e["name"]): e for e in json.loads(det_path.read_text())["libraries"] + _renamed(e["name"]): e + for e in json.loads(det_path.read_text())["libraries"] + if _library_stem(e["name"]) not in det_merge } assert set(ours) == set(theirs), manifest for name, entry in ours.items(): From 05b52f5dba97e662c1a244d31f2211ac5a97dece Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 18:00:32 -0800 Subject: [PATCH 06/10] cake_kda_det_train_kcp: the full-width offset walks take the split-role and overlap builds The forward offset walk's full-width record is the split-role build (320 threads). The adjoint offset walk ships the overlap-schedule build beside the plain one; the walk_heads rule (WALK_HEADS_RULE) picks the overlap build from 64 value heads. Every build gives the same bits. --- .../cake_kda_det_train_kcp/README.md | 12 +- .../cake_kda_det_train_kcp/cake_kcp.py | 5 + .../cake_kda_det_train_kcp/cake_kcp_jit.py | 130 ++++++++++++++++-- .../test_cake_kda_det_train_kcp_build.py | 59 +++++++- .../test_cake_kda_train_kcp_build.py | 53 ++++--- 5 files changed, 213 insertions(+), 46 deletions(-) diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/README.md b/flashinfer/experimental/cake_kda_det_train_kcp/README.md index 6e9e05630e1..299fbe703c9 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/README.md +++ b/flashinfer/experimental/cake_kda_det_train_kcp/README.md @@ -472,14 +472,16 @@ families of a phase in one launch; off by default). The files of this package: -* `cake_kcp_jit.py` -- the `MODULES` registry of the thirteen shipped programs (the +* `cake_kcp_jit.py` -- the `MODULES` registry of the fourteen shipped programs (the fragment walks of the forward and backward state recurrences in their `offset` and `tape` forms, the BF16x3 transfer walk (two CTAs per segment and head, each over one 64-column half of the state), the affine merge in one record per register-blocked build, and the two chunk stages of the context-parallel route, `kcp_fwd_o` (the output from the K1 tapes) and `kcp_dav` (`dA_qk` and the `dv1` K2 reads); selection by `always`, by - `walk_form`, or by `merge_rule` (the `MERGE_RULE` bands of the merge + `walk_form`, by `walk_heads` (the `WALK_HEADS_RULE` bands of the launch's + head count: the plain or the overlap-schedule full-width adjoint offset + walk), or by `merge_rule` (the `MERGE_RULE` bands of the merge launch's 128-column chain blocks per architecture), grids by `kcp_segments_heads`, `kcp_chains`, `kcp_tile_items` or `kcp_pairs`), the logical vocabulary of the KCP arguments, `run_program(name, args, device_index=)`, the lowering @@ -500,18 +502,18 @@ The files of this package: tests; without the file those tests skip. The package reads no environment variable. * `cake_build.py` + `build.json` -- the loader shared by every Cake-exported - package (byte-identical copies) and the manifest of the thirteen shipped + package (byte-identical copies) and the manifest of the fourteen shipped programs (schema `cake-build/1`, one entry per `_`); `aot/` is never tracked (`.gitignore`). * `csrc/cake_ffi_checks.h` -- the TVM-FFI helper header the bindings include. -* `csrc/cake_kda_det_train__kcp_{kernel,binding}.cu` of the seventeen +* `csrc/cake_kda_det_train__kcp_{kernel,binding}.cu` of the eighteen context-parallel-only programs and `csrc/cake_kda_det_train__{kernel,binding}.cu` of the two chunk stages (`kcp_fwd_o`, `kcp_dav`), rendered like the single-rank ones (digests in `build.json` and `build.test.json`). ## Status -The registry holds thirteen shipped programs (`cake_kcp_jit.MODULES`, in +The registry holds fourteen shipped programs (`cake_kcp_jit.MODULES`, in `build.json`) and six test-only programs (`cake_kcp_jit.TEST_MODULES`, in `build.test.json`) for `sm_100a`, `sm_103a` and `sm_107a`. diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py index 3cb6e756b9e..1f2fb6fe950 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py @@ -2401,6 +2401,11 @@ def _launch( if rule == "fit": if walk_form == "split": record += "_vh64" + elif record in getattr(jit, "WALK_HEADS_RULE", {}): + # the full-width record the head-count rule names (the ABI carries the head count) + record = jit.walk_heads_record( + record, int(abi[ABI_PARAMS[program].index("num_heads")]) + ) elif walk_form in ("full", "split") and walk_form != rule: record += f"_{walk_form}" if walk_form == "split": # two CTAs per segment and head diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index 5973ffa40d8..ffd43c57163 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -46,7 +46,10 @@ launch from its walk list; * ``kcp_dhu_offset`` / ``kcp_dhu_tape`` -- the adjoint recurrence likewise (reverse walk; the tape program writes the ``dh`` / ``dv2`` tapes), in the - same two forms; + same two forms; the offset program's full-width form ships in two records, + bitwise equal (``select_rule`` ``walk_heads``): the plain build and the + overlap-schedule build ``_overlap``, picked per launch from the head count + (:data:`WALK_HEADS_RULE`, :func:`walk_heads_record`); * ``kcp_transfer_walk`` -- the FP32-class transfer ``M_f`` (``^T``) of a segment (BF16x3 on the tensor cores: two CTAs per segment and head, each over one 64-column half of the state) or, with ``init_mode = 1``, its @@ -254,7 +257,7 @@ # The ``kcp_affine_merge`` record of a launch (:func:`merge_record`), by architecture: ``T = n_chains * ncols / 128`` # (the launch's 128-column chain blocks; one rule for both operator layouts) -> the record of the first band whose # upper bound holds ``T`` (``None``: no bound). Every record gives the same bits. The bands are the fastest build per -# ``T`` in the r2611 K4 A/B on GB300 (sm_103a, 152 SMs: 17 launch rows x 11 builds, every row bitwise equal to the +# ``T`` in an A/B on GB300 (sm_103a, 152 SMs: 17 launch rows x 11 builds, every row bitwise equal to the # previous build); sm_100a takes the GB300 bands (not measured on sm_100a). sm_107a has no packed FMA pairs and ships # no 8 x 8 build: its bands are the GB300 bands with the scalar builds and the 4 x 8 build above ``T`` 56 (a proxy, not # measured on sm_107a). @@ -281,6 +284,16 @@ (None, "kcp_affine_merge_c32"), ), } +# The full-width record of a K1 / K2 offset walk under the ``walk_heads`` rule: walk program -> bands ``(upper bound +# of the launch's value-head count H or None, record)``, the record of the first band whose bound holds ``H``. Both +# records give the same bits. kcp_dhu_offset: the overlap-schedule build against the plain one in an offset-walk +# A/B on GB300 (sm_103a, hot medians): faster on every H = 64 row (W4 rank 1 247.94 -> 241.31 us, W4 rank 0 147.62 -> +# 144.03, z2 548.40 -> 479.36, c14 373.63 -> 369.89, d4 197.12 -> 195.20), slower on both H = 4 rows (default 248.42 -> +# 252.51, p8 224.58 -> 225.76); H from 5 to 63 is not measured and keeps the plain build. sm_100a and sm_107a take +# the GB300 bands (not measured there). +WALK_HEADS_RULE: dict[str, tuple[tuple[int | None, str], ...]] = { + "kcp_dhu_offset": ((63, "kcp_dhu_offset"), (None, "kcp_dhu_offset_overlap")), +} # Filled mechanically by the exporter (the derived leaves ``kernel_symbol``, ``launch``, ``arg_plan``, # ``logical``); the host plan (``program``, ``variant``, ``select_rule``, ``grid_rule``, ``library``, ``arches``) @@ -316,15 +329,15 @@ ], "ffi_entry": "run", "grid_rule": {"kind": "kcp_segments_heads"}, - "kernel_symbol": "kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp", + "kernel_symbol": "kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp", "launch": { - "block": (192, 1, 1), + "block": (320, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 115712, + "dynamic_smem_bytes": 0, "use_pdl": False, }, - "library": "cake_kda_det_train_7658fcc297c6fd906d53_kcp", + "library": "cake_kda_det_train_7d43005ed0adb9ffadf2_kcp", "logical": { "gk": "gk", "h0": "h0", @@ -539,9 +552,77 @@ "walk_seg": "walk_seg", }, "program": "kcp_dhu_offset", - "select_rule": {"form": "full", "kind": "walk_form"}, + "select_rule": {"form": "full", "kind": "walk_heads"}, "variant": None, }, + "kcp_dhu_offset_overlap": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "kg_tma"], + ["buffer", "qg_tma"], + ["buffer", "w_tma"], + ["buffer", "do_tma"], + ["buffer", "dv1"], + ["buffer", "gk"], + ["parameter", "num_heads"], + ["buffer", "seq_chunk_start"], + ["parameter", "num_seqs"], + ["parameter", "scale"], + ["buffer", "walk_seg"], + ["buffer", "dht"], + ["buffer", "dht_index"], + ["parameter", "dht_hg"], + ["parameter", "dht_n"], + ["buffer", "out"], + ["buffer", "out_index"], + ["parameter", "out_hg"], + ["parameter", "out_n"], + ["parameter", "tile_stride"], + ["buffer", "dh_tape"], + ["buffer", "dv2"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": {"kind": "kcp_segments_heads"}, + "kernel_symbol": "kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp", + "launch": { + "block": (192, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 0, + "use_pdl": False, + }, + "library": "cake_kda_det_train_4a02b970484a9ef1f3c1_kcp", + "logical": { + "dh_tape": "dh", + "dht": "dht", + "dht_hg": "dht_hg", + "dht_index": "dht_index", + "dht_n": "dht_n", + "do_tma": "do_int", + "dv1": "dv1", + "dv2": "dv2", + "gk": "gk", + "kg_tma": "kg", + "num_heads": "num_heads", + "num_seqs": "num_seqs", + "out": "out", + "out_hg": "out_hg", + "out_index": "out_index", + "out_n": "out_n", + "qg_tma": "qg", + "scale": "scale", + "seq_chunk_start": "seg_chunk_start", + "tile_stride": "tile_stride", + "w_tma": "w", + "walk_seg": "walk_seg", + }, + "program": "kcp_dhu_offset", + "select_rule": {"form": "full", "kind": "walk_heads"}, + "variant": "overlap", + }, "kcp_dhu_offset_vh64": { "arches": ["sm_100a", "sm_103a", "sm_107a"], "arg_plan": [ @@ -1435,6 +1516,7 @@ def _admits( form: str | None = None, n_chains: int | None = None, arch: str | None = None, + num_heads: int | None = None, ) -> bool: kind = rule["kind"] if kind == "always": @@ -1443,6 +1525,10 @@ def _admits( return ncols is not None and int(ncols) == int(rule["ncols"]) if kind == "walk_form": return form is not None and str(form) == str(rule["form"]) + if kind == "walk_heads": + if form is None or str(form) != str(rule["form"]) or num_heads is None: + return False + return walk_heads_record(MODULES[name]["program"], int(num_heads)) == name if kind == "merge_rule": if ncols is None or n_chains is None: return False @@ -1457,6 +1543,7 @@ def select_record( form: str | None = None, n_chains: int | None = None, arch: str | None = None, + num_heads: int | None = None, ) -> str: """The one record of ``program`` the launch's scalars select. @@ -1468,8 +1555,11 @@ def select_record( (``{"kind": "walk_form", "form": F}``) picks the K1 / K2 walk record of form ``F`` -- ``"full"`` (one CTA per segment and head) or ``"split"`` (the two-CTA record over 64 value columns each), the launch's form from - :func:`walk_form`; a program under a fixed rule ships its one ``always`` - record, the other form being test-only (:func:`test_walk_record`). + :func:`walk_form`; ``walk_heads`` (``{"kind": "walk_heads", "form": F}``) + picks, among the records of form ``F``, the one :data:`WALK_HEADS_RULE` + names for the launch's ``num_heads`` (:func:`walk_heads_record`); a + program under a fixed rule ships its one ``always`` record, the other form + being test-only (:func:`test_walk_record`). Exactly one record admits every launch, or the registry is inconsistent. """ admitted = [ @@ -1482,12 +1572,13 @@ def select_record( form=form, n_chains=n_chains, arch=arch, + num_heads=num_heads, ) ] if len(admitted) != 1: raise ValueError( f"program {program!r}: {len(admitted)} records admit ncols={ncols} form={form} n_chains={n_chains} " - f"arch={arch} ({admitted}); expected one" + f"arch={arch} num_heads={num_heads} ({admitted}); expected one" ) return admitted[0] @@ -1526,6 +1617,21 @@ def merge_record(ncols: int, n_chains: int, arch: str | None = None) -> str: return record +def walk_heads_record(program: str, num_heads: int) -> str: + """The full-width record of walk ``program`` for a launch over ``num_heads`` value heads + (:data:`WALK_HEADS_RULE`; every record of the program gives the same bits).""" + bands = WALK_HEADS_RULE.get(program) + if bands is None: + raise ValueError( + f"walk {program!r} has no head-count rule; the rules: {tuple(WALK_HEADS_RULE)}" + ) + if int(num_heads) < 1: + raise ValueError(f"num_heads must be >= 1, got {num_heads!r}") + return next( + record for bound, record in bands if bound is None or int(num_heads) <= bound + ) + + def walk_form( program: str, chains: Sequence[int], @@ -2125,7 +2231,7 @@ def _walk(name: str, args: Mapping[str, Any], call: Any, device_index: int) -> N record = test_walk_record(name, form)[0] module = load_test(record, device_arch(device_index)) else: - record = select_record(name, form=form) + record = select_record(name, form=form, num_heads=H) WALK_FORM_STATS[form] += 1 log = getattr(call, "walk_form_log", None) if log is not None: @@ -2278,6 +2384,7 @@ def run_program(name: str, args: Mapping[str, Any], *, device_index: int) -> Non "WALK_FORM_FIT_STEP_RATIO", "WALK_FORM_RULES", "WALK_FORM_STATS", + "WALK_HEADS_RULE", "WALK_PROGRAMS", "arch_for_capability", "device_arch", @@ -2300,4 +2407,5 @@ def run_program(name: str, args: Mapping[str, Any], *, device_index: int) -> Non "test_specs_for_arch", "test_walk_record", "walk_form", + "walk_heads_record", ] diff --git a/tests/experimental/test_cake_kda_det_train_kcp_build.py b/tests/experimental/test_cake_kda_det_train_kcp_build.py index 9f49a6dbf51..ef8858188fa 100644 --- a/tests/experimental/test_cake_kda_det_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_det_train_kcp_build.py @@ -121,6 +121,7 @@ "kcp_fwdh_offset_vh64": FWDH_PLAN, "kcp_fwdh_tape": FWDH_PLAN, "kcp_dhu_offset": DHU_PLAN, + "kcp_dhu_offset_overlap": DHU_PLAN, "kcp_dhu_offset_vh64": DHU_PLAN, "kcp_dhu_tape": DHU_PLAN, "kcp_transfer_walk": TRANSFER_PLAN, @@ -134,6 +135,9 @@ "kcp_fwdh_tape_full": FWDH_PLAN, "kcp_dhu_tape_full": DHU_PLAN, } +# Launch block of the records that are not 192 threads wide (the merge records: MERGE_RECORDS); the full-width fwdh +# offset walk runs the split-role schedule (four state, four decay, one TMA and one MMA warp). +BLOCKS = {"kcp_fwdh_offset": 320} # The identity rule: ``cake_kda_det_train_<20 hex of sha256("\\n")>``, # then ``_kcp`` for a program only the context-parallel route runs (every one but the single-rank ``dav`` stage). # The digest hashes the source module path the program shipped under, so the digests are those of the previous names @@ -141,8 +145,8 @@ _W = "sm100_sm103_sm107" IDENTITIES = { "kcp_fwdh_offset": ( - f"tart.kda_chunk_kcp_fwd_h.{_W}.kda_chunk_kcp_fwdh_v1", - {"program": "offset"}, + f"tart.kda_chunk_kcp_fwd_h.{_W}.kda_chunk_det_fwdh_v2_kcp", + {"program": "offset", "tile_prefetch": 0}, ), "kcp_fwdh_offset_vh64": ( f"tart.kda_chunk_kcp_fwd_h.{_W}.kda_chunk_kcp_fwdh_v1", @@ -156,6 +160,10 @@ f"tart.kda_chunk_kcp_dhu.{_W}.kda_chunk_kcp_dhu_v1", {"program": "offset"}, ), + "kcp_dhu_offset_overlap": ( + f"tart.kda_chunk_kcp_dhu.{_W}.kda_chunk_det_dhu_v2_kcp", + {"program": "offset", "tile_prefetch": 0}, + ), "kcp_dhu_offset_vh64": ( f"tart.kda_chunk_kcp_dhu.{_W}.kda_chunk_kcp_dhu_v1", {"program": "offset", "vh": 64}, @@ -249,7 +257,7 @@ def test_record_matches_the_frozen_abi(name: str) -> None: assert record["launch"]["block"] == ( (MERGE_RECORDS[name][1], 1, 1) if record["program"] == "kcp_affine_merge" - else (192, 1, 1) + else (BLOCKS.get(name, 192), 1, 1) ) if record["program"] == "kcp_affine_merge": assert record["grid_rule"] == { @@ -461,6 +469,51 @@ def test_walk_form_rules_and_records() -> None: assert kcp_jit.MODULES["kcp_fwdh_tape"]["select_rule"] == {"kind": "always"} +def test_walk_heads_selection() -> None: + # the full-width adjoint offset walk: the plain build up to 63 heads, the overlap-schedule build from 64 (the W4 + # rows); its split form and the forward offset walk keep the walk_form rule + assert kcp_jit.WALK_HEADS_RULE == { + "kcp_dhu_offset": ((63, "kcp_dhu_offset"), (None, "kcp_dhu_offset_overlap")) + } + for name in ("kcp_dhu_offset", "kcp_dhu_offset_overlap"): + assert kcp_jit.MODULES[name]["select_rule"] == { + "form": "full", + "kind": "walk_heads", + } + assert kcp_jit.MODULES[name]["grid_rule"] == {"kind": "kcp_segments_heads"} + for heads, record in ( + (1, "kcp_dhu_offset"), + (4, "kcp_dhu_offset"), + (63, "kcp_dhu_offset"), + (64, "kcp_dhu_offset_overlap"), + (128, "kcp_dhu_offset_overlap"), + ): + assert kcp_jit.walk_heads_record("kcp_dhu_offset", heads) == record + assert ( + kcp_jit.select_record("kcp_dhu_offset", form="full", num_heads=heads) + == record + ) + assert ( + kcp_jit.select_record("kcp_dhu_offset", form="split", num_heads=heads) + == "kcp_dhu_offset_vh64" + ) + assert ( + kcp_jit.select_record("kcp_fwdh_offset", form="full", num_heads=heads) + == "kcp_fwdh_offset" + ) + assert kcp_jit.launch_grid("kcp_dhu_offset_overlap", n_walk=3, num_heads=64) == ( + 192, + 1, + 1, + ) + with pytest.raises(ValueError): + kcp_jit.select_record("kcp_dhu_offset", form="full") # the head count is needed + with pytest.raises(ValueError): + kcp_jit.walk_heads_record("kcp_dhu_offset", 0) + with pytest.raises(ValueError): + kcp_jit.walk_heads_record("kcp_fwdh_offset", 64) + + def test_the_shipped_transfer_walk_is_the_half_width_bf16x3_program() -> None: record = kcp_jit.MODULES["kcp_transfer_walk"] assert record["grid_rule"] == {"ctas_per_head": 2, "kind": "kcp_segments_heads"} diff --git a/tests/experimental/test_cake_kda_train_kcp_build.py b/tests/experimental/test_cake_kda_train_kcp_build.py index aef261be949..7bb08a96760 100644 --- a/tests/experimental/test_cake_kda_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_train_kcp_build.py @@ -490,20 +490,19 @@ def test_test_manifest_is_the_test_only_programs_under_the_opt_in( DET_PREFIX, PREFIX = "cake_kda_det_train_", "cake_kda_train_" -# The affine merge is the one program whose builds differ between the two packages: the deterministic package ships -# register-blocked merge builds selected by its MERGE_RULE; this package keeps its two column-block builds (c8, c16). -# Every other record and unit is the deterministic package's under this package's names. -MERGE = "kcp_affine_merge" -OWN_MERGE = { +# The records whose builds differ between the two packages: the deterministic package ships newer builds of the +# affine merge (selected by its MERGE_RULE) and of the full-width offset walks (the adjoint one beside the plain build, +# selected by its WALK_HEADS_RULE); this package keeps its own builds of these records. The deterministic package's +# extra records belong to those programs. Every other record and unit is the deterministic package's under this +# package's names. +OWN = { + "kcp_fwdh_offset": PREFIX + "7658fcc297c6fd906d53_kcp", + "kcp_dhu_offset": PREFIX + "1c66573e87150245374b_kcp", "kcp_affine_merge_c8": PREFIX + "c6958a0a4ed5cb55a9a1_kcp", "kcp_affine_merge_c16": PREFIX + "89f5ff59f03eb170e3bf_kcp", } -def _merge_records(table: dict) -> set[str]: - return {name for name, record in table.items() if record["program"] == MERGE} - - def _library_stem(name: str) -> str: return name.rsplit("_sm_", 1)[0] @@ -522,21 +521,20 @@ def _renamed(value): def test_the_records_are_the_deterministic_records_under_this_packages_names() -> None: """The KCP programs have no non-deterministic form: both tables hold the deterministic package's records, every field equal but the library stem and the kernel symbol, which carry the same twenty hex digits under this - package's prefix; no library or kernel symbol of this package is one of the deterministic package's. The affine - merge records are this package's own column-block builds.""" + package's prefix; no library or kernel symbol of this package is one of the deterministic package's. The records + of ``OWN`` are this package's own builds.""" + own_programs = {kcp_jit.MODULES[name]["program"] for name in OWN} for table in ("MODULES", "TEST_MODULES"): ours, theirs = getattr(kcp_jit, table), getattr(det_kcp_jit, table) - assert set(ours) - _merge_records(ours) == set(theirs) - _merge_records( - theirs - ), table + assert set(ours) <= set(theirs), table + assert {theirs[n]["program"] for n in set(theirs) - set(ours)} <= own_programs for name, record in ours.items(): assert record["library"].startswith(PREFIX), name - if record["program"] != MERGE: + if name in OWN: + assert record["library"] == OWN[name], name + else: assert record == _renamed(theirs[name]), (table, name) - assert { - name: kcp_jit.MODULES[name]["library"] - for name in _merge_records(kcp_jit.MODULES) - } == OWN_MERGE + assert set(OWN) <= set(kcp_jit.MODULES) ours = {r["library"] for r in _records().values()} theirs = { r["library"] @@ -554,25 +552,26 @@ def test_the_units_are_the_deterministic_units_under_this_packages_names( ) -> None: """Every library of this package's manifest is the deterministic package's library of the same digits: the same architecture, flags and entry, and sources whose text is the deterministic unit's with the deterministic prefix - renamed (file names, kernel symbols, binding namespaces); nothing else differs. The affine merge libraries, which - differ between the two packages, are left out on both sides.""" + renamed (file names, kernel symbols, binding namespaces); nothing else differs. The libraries of the records whose + builds differ between the two packages (``OWN``, and the deterministic package's extra records) are left out on + both sides.""" det_dir = Path(det_kcp_jit.__file__).resolve().parent path, det_path = PACKAGE_DIR / manifest, det_dir / manifest if not path.is_file() or not det_path.is_file(): pytest.skip(f"{manifest} is not rendered in both packages here") - own_merge = {kcp_jit.MODULES[n]["library"] for n in _merge_records(kcp_jit.MODULES)} - det_merge = { - det_kcp_jit.MODULES[n]["library"] for n in _merge_records(det_kcp_jit.MODULES) - } + own = set(OWN.values()) + det_records = {**det_kcp_jit.MODULES, **det_kcp_jit.TEST_MODULES} + left_out = set(OWN) | (set(det_records) - set(_records())) + det_own = {det_records[n]["library"] for n in left_out if n in det_records} ours = { e["name"]: e for e in json.loads(path.read_text())["libraries"] - if _library_stem(e["name"]) not in own_merge + if _library_stem(e["name"]) not in own } theirs = { _renamed(e["name"]): e for e in json.loads(det_path.read_text())["libraries"] - if _library_stem(e["name"]) not in det_merge + if _library_stem(e["name"]) not in det_own } assert set(ours) == set(theirs), manifest for name, entry in ours.items(): From 544a72ff7cd59611003fa209bda24d98fcb6bce0 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 18:08:01 -0800 Subject: [PATCH 07/10] cake_kda_det_train_kcp: render the split-role and overlap offset walk builds --- .../cake_kda_det_train_kcp/build.json | 147 +- .../cake_kda_det_train_kcp/cake_kcp_jit.py | 4 +- ..._train_4a02b970484a9ef1f3c1_kcp_binding.cu | 347 +++++ ...t_train_4a02b970484a9ef1f3c1_kcp_kernel.cu | 1376 +++++++++++++++++ ...t_train_7658fcc297c6fd906d53_kcp_kernel.cu | 636 -------- ...train_7d43005ed0adb9ffadf2_kcp_binding.cu} | 36 +- ...t_train_7d43005ed0adb9ffadf2_kcp_kernel.cu | 926 +++++++++++ 7 files changed, 2779 insertions(+), 693 deletions(-) create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu delete mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu rename flashinfer/experimental/cake_kda_det_train_kcp/csrc/{cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu => cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu} (84%) create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/build.json b/flashinfer/experimental/cake_kda_det_train_kcp/build.json index 5a5e5b4e801..ac691d1a85a 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/build.json +++ b/flashinfer/experimental/cake_kda_det_train_kcp/build.json @@ -186,6 +186,69 @@ "csrc/cake_kda_det_train_32b1aa657b082f2f4cb7_kcp_binding.cu" ] }, + { + "arch": "sm_100a", + "closure_sha256": "b2d7abadf874f420d6132597e7195f739168e51596f03a51ef38a49a46fa92e0", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_sm_100a.so", + "sha256": { + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu": "9e01aaad0429b57644b52c7c3694a4965b07e84b524a7bebab5e111fc03e838e", + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu": "06b9dfde6af6ce63b86f774a60b7ee8648817b0388215f50dc45689ec64a3b1e" + }, + "sources": [ + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu", + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "663fde510cb99a60271c31a343549615ceaa53acbbc702b4070da667d82139d8", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_sm_103a.so", + "sha256": { + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu": "9e01aaad0429b57644b52c7c3694a4965b07e84b524a7bebab5e111fc03e838e", + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu": "06b9dfde6af6ce63b86f774a60b7ee8648817b0388215f50dc45689ec64a3b1e" + }, + "sources": [ + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu", + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "c9ae1ccae892adc89f9ea67fae8c68e6d40c3832f7920db440f7abf931ea0cc4", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_sm_107a.so", + "sha256": { + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu": "9e01aaad0429b57644b52c7c3694a4965b07e84b524a7bebab5e111fc03e838e", + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu": "06b9dfde6af6ce63b86f774a60b7ee8648817b0388215f50dc45689ec64a3b1e" + }, + "sources": [ + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu", + "csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "01db09151728728a86485f98bc4863e8f12ca4a89a738cb0969bc326c8fb1041", @@ -314,128 +377,128 @@ }, { "arch": "sm_100a", - "closure_sha256": "1bb256121e765dc4c41de2673e1d433b0b9e8387e605c501a3e27ca2e726d9ca", + "closure_sha256": "1abe6d29eeab5be3d96e0dc0a4e6b7a646a89817892fbb3b7704ab1ad5883ea0", "entry": "run", "gencode": "-gencode=arch=compute_100a,code=sm_100a", "min_cuda": "12.8", - "name": "cake_kda_det_train_7658fcc297c6fd906d53_kcp_sm_100a", + "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_100a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7658fcc297c6fd906d53_kcp_sm_100a.so", + "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_100a.so", "sha256": { - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu": "410e1ae214daa249128825f0f35562df6646010fdc6879f8c10cb646fe02902b", - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu": "ed6531784ddb459dd3e4c36103eb954290e72998b9b347415503ae61afad3d9a" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" }, "sources": [ - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu", - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" ] }, { "arch": "sm_103a", - "closure_sha256": "7f3fc8f4a66148b90a9d639ef1ea290765d15d97dab4241d0970c4c7cc7b4eb5", + "closure_sha256": "c523a00b55ca9add18bd82866a5bd53223caf3a15ccc9a98fc2a6d4932e53125", "entry": "run", "gencode": "-gencode=arch=compute_103a,code=sm_103a", "min_cuda": "12.9", - "name": "cake_kda_det_train_7658fcc297c6fd906d53_kcp_sm_103a", + "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_103a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7658fcc297c6fd906d53_kcp_sm_103a.so", + "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_103a.so", "sha256": { - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu": "410e1ae214daa249128825f0f35562df6646010fdc6879f8c10cb646fe02902b", - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu": "ed6531784ddb459dd3e4c36103eb954290e72998b9b347415503ae61afad3d9a" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" }, "sources": [ - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu", - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" ] }, { "arch": "sm_107a", - "closure_sha256": "d043e9ca5bc5587eb0ed66ac3b3e0a257396ec42730d55988d617b04b38562bf", + "closure_sha256": "b6b3daf2d4fccef3711a403214d28c2c6386100a7dc5fd3fe0942ba7fb7fd89a", "entry": "run", "gencode": "-gencode=arch=compute_107a,code=sm_107a", "min_cuda": "13.4", - "name": "cake_kda_det_train_7658fcc297c6fd906d53_kcp_sm_107a", + "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_107a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7658fcc297c6fd906d53_kcp_sm_107a.so", + "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_107a.so", "sha256": { - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu": "410e1ae214daa249128825f0f35562df6646010fdc6879f8c10cb646fe02902b", - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu": "ed6531784ddb459dd3e4c36103eb954290e72998b9b347415503ae61afad3d9a" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" }, "sources": [ - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu", - "csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu" + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", + "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" ] }, { "arch": "sm_100a", - "closure_sha256": "1abe6d29eeab5be3d96e0dc0a4e6b7a646a89817892fbb3b7704ab1ad5883ea0", + "closure_sha256": "40a510414f130f7669a45921dd3479cd5b5fdcb792b2ef7884b6e2a77dc76ec9", "entry": "run", "gencode": "-gencode=arch=compute_100a,code=sm_100a", "min_cuda": "12.8", - "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_100a", + "name": "cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_sm_100a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_100a.so", + "path": "aot/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_sm_100a.so", "sha256": { - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu": "02c5687d40bae15f43f603b8ee4959f55080d7e7e9aa1e10dd9acb48f4c1a417", + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu": "0f4e660c0b96d69a574dba57b2eb9befa7407e456ceac4cfe72e50f3b7aa18f7" }, "sources": [ - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu", + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu" ] }, { "arch": "sm_103a", - "closure_sha256": "c523a00b55ca9add18bd82866a5bd53223caf3a15ccc9a98fc2a6d4932e53125", + "closure_sha256": "40968742244a880aabd9fe68d01b51c85864f1aedae382aea9a8d4eee6afac24", "entry": "run", "gencode": "-gencode=arch=compute_103a,code=sm_103a", "min_cuda": "12.9", - "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_103a", + "name": "cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_sm_103a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_103a.so", + "path": "aot/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_sm_103a.so", "sha256": { - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu": "02c5687d40bae15f43f603b8ee4959f55080d7e7e9aa1e10dd9acb48f4c1a417", + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu": "0f4e660c0b96d69a574dba57b2eb9befa7407e456ceac4cfe72e50f3b7aa18f7" }, "sources": [ - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu", + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu" ] }, { "arch": "sm_107a", - "closure_sha256": "b6b3daf2d4fccef3711a403214d28c2c6386100a7dc5fd3fe0942ba7fb7fd89a", + "closure_sha256": "9db6b492ed8a3c7cd5e42cbe96bf4c0732c25d80dd9f0bf59f6bdca88150291f", "entry": "run", "gencode": "-gencode=arch=compute_107a,code=sm_107a", "min_cuda": "13.4", - "name": "cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_107a", + "name": "cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_sm_107a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_sm_107a.so", + "path": "aot/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_sm_107a.so", "sha256": { - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu": "9e0295ae3bc1bf04ab5dc7bbc2753001fb1b63fbafa48368dee99b789c244fb5", - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu": "f72b01b7521e9fb6e8bb390494972540c46f5792a3e59950fff2e66af808ed44" + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu": "02c5687d40bae15f43f603b8ee4959f55080d7e7e9aa1e10dd9acb48f4c1a417", + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu": "0f4e660c0b96d69a574dba57b2eb9befa7407e456ceac4cfe72e50f3b7aa18f7" }, "sources": [ - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_kernel.cu", - "csrc/cake_kda_det_train_7b6b22d3423f3138baf7_kcp_binding.cu" + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu", + "csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu" ] }, { diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index ffd43c57163..223e636eb1f 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -334,7 +334,7 @@ "block": (320, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 115712, "use_pdl": False, }, "library": "cake_kda_det_train_7d43005ed0adb9ffadf2_kcp", @@ -591,7 +591,7 @@ "block": (192, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 115712, "use_pdl": False, }, "library": "cake_kda_det_train_4a02b970484a9ef1f3c1_kcp", diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu new file mode 100644 index 00000000000..3b4a740c8ca --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_binding.cu @@ -0,0 +1,347 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp, for the device source that defines it. Do not edit. +#include +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp(const __grid_constant__ CUtensorMap kg_tma, const __grid_constant__ CUtensorMap qg_tma, const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap do_tma, __nv_bfloat16* __restrict__ dv1, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, float scale, int* __restrict__ walk_seg, float* __restrict__ dht, int* __restrict__ dht_index, int dht_hg, int dht_n, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, __nv_bfloat16* __restrict__ dh_tape, __nv_bfloat16* __restrict__ dv2); + +namespace cake_binding_kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void CheckCu(CUresult result, const char* what) { + if (result != CUDA_SUCCESS) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + TVM_FFI_THROW(RuntimeError) << KERNEL_NAME << ": " << what << " failed with " + << (name ? name : "an unknown CUresult") << " (" << static_cast(result) + << "): " << (message ? message : ""); + } +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +// The product of the extents above the trailing `trailing` dims. +int64_t NumelAbove(const TensorView& t, int trailing) { + int64_t product = 1; + for (int axis = 0; axis < t.ndim() - trailing; ++axis) product *= t.size(axis); + return product; +} + +[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) { + TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b + << " has no int64 value in a host check"; +} + +int64_t CheckedMul(int64_t a, int64_t b) { + int64_t product; + if (__builtin_mul_overflow(a, b, &product)) NoInt64(a, "*", b); + return product; +} + +cuuint64_t Extent(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value > 0, ValueError) + << KERNEL_NAME << ": an extent of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +cuuint64_t Stride(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value >= 0, ValueError) + << KERNEL_NAME << ": a stride of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +CUtensorMap EncodeTma_kg_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'kg_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "kg_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "kg_tma"), Extent(d2, "kg_tma"), Extent(outer2, "kg_tma")}; + cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "kg_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "kg_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'kg_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'kg_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'kg_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of kg_tma"); + return map; +} + +CUtensorMap EncodeTma_qg_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'qg_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "qg_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "qg_tma"), Extent(d2, "qg_tma"), Extent(outer2, "qg_tma")}; + cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "qg_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "qg_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'qg_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'qg_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'qg_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of qg_tma"); + return map; +} + +CUtensorMap EncodeTma_w_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'w_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "w_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "w_tma"), Extent(d2, "w_tma"), Extent(outer2, "w_tma")}; + cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "w_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "w_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'w_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'w_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'w_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of w_tma"); + return map; +} + +CUtensorMap EncodeTma_do_tma(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << KERNEL_NAME << ": descriptor 'do_tma' takes a source of at least 2 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "do_tma"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t outer2 = NumelAbove(t, 2); + cuuint64_t global_dim[3] = {Extent(d1, "do_tma"), Extent(d2, "do_tma"), Extent(outer2, "do_tma")}; + cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "do_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "do_tma")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'do_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'do_tma' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'do_tma' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of do_tma"); + return map; +} + +void Run(TensorView a_kg_tma, TensorView a_qg_tma, TensorView a_w_tma, TensorView a_do_tma, TensorView a_dv1, TensorView a_gk, int64_t a_num_heads, TensorView a_seq_chunk_start, int64_t a_num_seqs, double a_scale, TensorView a_walk_seg, TensorView a_dht, TensorView a_dht_index, int64_t a_dht_hg, int64_t a_dht_n, TensorView a_out, TensorView a_out_index, int64_t a_out_hg, int64_t a_out_n, int64_t a_tile_stride, TensorView a_dh_tape, TensorView a_dv2, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_kg_tma.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_kg_tma, "kg_tma"); + check_dtype(a_kg_tma, DLDataType{kDLBfloat, 16, 1}, "kg_tma"); + check_cuda_tensor(a_qg_tma, "qg_tma"); + check_dtype(a_qg_tma, DLDataType{kDLBfloat, 16, 1}, "qg_tma"); + check_cuda_tensor(a_w_tma, "w_tma"); + check_dtype(a_w_tma, DLDataType{kDLBfloat, 16, 1}, "w_tma"); + check_cuda_tensor(a_do_tma, "do_tma"); + check_dtype(a_do_tma, DLDataType{kDLBfloat, 16, 1}, "do_tma"); + check_cuda_tensor(a_dv1, "dv1"); + check_dtype(a_dv1, DLDataType{kDLBfloat, 16, 1}, "dv1"); + check_cuda_tensor(a_gk, "gk"); + check_dtype(a_gk, DLDataType{kDLFloat, 32, 1}, "gk"); + TVM_FFI_CHECK(a_num_heads >= -2147483648LL && a_num_heads <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_heads' = " << a_num_heads + << " is outside the i32 range [-2147483648, 2147483647]"; + check_cuda_tensor(a_seq_chunk_start, "seq_chunk_start"); + check_dtype(a_seq_chunk_start, DLDataType{kDLInt, 32, 1}, "seq_chunk_start"); + TVM_FFI_CHECK(a_num_seqs >= -2147483648LL && a_num_seqs <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_seqs' = " << a_num_seqs + << " is outside the i32 range [-2147483648, 2147483647]"; + check_cuda_tensor(a_walk_seg, "walk_seg"); + check_dtype(a_walk_seg, DLDataType{kDLInt, 32, 1}, "walk_seg"); + check_cuda_tensor(a_dht, "dht"); + check_dtype(a_dht, DLDataType{kDLFloat, 32, 1}, "dht"); + check_cuda_tensor(a_dht_index, "dht_index"); + check_dtype(a_dht_index, DLDataType{kDLInt, 32, 1}, "dht_index"); + TVM_FFI_CHECK(a_dht_hg >= -2147483648LL && a_dht_hg <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'dht_hg' = " << a_dht_hg + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_dht_n >= -2147483648LL && a_dht_n <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'dht_n' = " << a_dht_n + << " is outside the i32 range [-2147483648, 2147483647]"; + check_cuda_tensor(a_out, "out"); + check_dtype(a_out, DLDataType{kDLFloat, 32, 1}, "out"); + check_cuda_tensor(a_out_index, "out_index"); + check_dtype(a_out_index, DLDataType{kDLInt, 32, 1}, "out_index"); + TVM_FFI_CHECK(a_out_hg >= -2147483648LL && a_out_hg <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'out_hg' = " << a_out_hg + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_out_n >= -2147483648LL && a_out_n <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'out_n' = " << a_out_n + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_tile_stride >= -2147483648LL && a_tile_stride <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'tile_stride' = " << a_tile_stride + << " is outside the i32 range [-2147483648, 2147483647]"; + check_cuda_tensor(a_dh_tape, "dh_tape"); + check_dtype(a_dh_tape, DLDataType{kDLBfloat, 16, 1}, "dh_tape"); + check_cuda_tensor(a_dv2, "dv2"); + check_dtype(a_dv2, DLDataType{kDLBfloat, 16, 1}, "dv2"); + check_contiguous(a_kg_tma, "kg_tma"); + check_contiguous(a_qg_tma, "qg_tma"); + check_same_device(a_qg_tma, a_kg_tma, "qg_tma", "kg_tma"); + check_contiguous(a_w_tma, "w_tma"); + check_same_device(a_w_tma, a_kg_tma, "w_tma", "kg_tma"); + check_contiguous(a_do_tma, "do_tma"); + check_same_device(a_do_tma, a_kg_tma, "do_tma", "kg_tma"); + check_contiguous(a_dv1, "dv1"); + check_same_device(a_dv1, a_kg_tma, "dv1", "kg_tma"); + check_contiguous(a_gk, "gk"); + check_same_device(a_gk, a_kg_tma, "gk", "kg_tma"); + check_contiguous(a_seq_chunk_start, "seq_chunk_start"); + check_same_device(a_seq_chunk_start, a_kg_tma, "seq_chunk_start", "kg_tma"); + check_contiguous(a_walk_seg, "walk_seg"); + check_same_device(a_walk_seg, a_kg_tma, "walk_seg", "kg_tma"); + check_contiguous(a_dht, "dht"); + check_same_device(a_dht, a_kg_tma, "dht", "kg_tma"); + check_contiguous(a_dht_index, "dht_index"); + check_same_device(a_dht_index, a_kg_tma, "dht_index", "kg_tma"); + check_contiguous(a_out, "out"); + check_same_device(a_out, a_kg_tma, "out", "kg_tma"); + check_contiguous(a_out_index, "out_index"); + check_same_device(a_out_index, a_kg_tma, "out_index", "kg_tma"); + check_contiguous(a_dh_tape, "dh_tape"); + check_same_device(a_dh_tape, a_kg_tma, "dh_tape", "kg_tma"); + check_contiguous(a_dv2, "dv2"); + check_same_device(a_dv2, a_kg_tma, "dv2", "kg_tma"); + TVM_FFI_CHECK((a_num_heads >= 1LL), ValueError) + << KERNEL_NAME << ": requires (num_heads >= 1)"; + TVM_FFI_CHECK((a_num_seqs >= 0LL), ValueError) + << KERNEL_NAME << ": requires (num_seqs >= 0)"; + TVM_FFI_CHECK((a_tile_stride >= 1LL), ValueError) + << KERNEL_NAME << ": requires (tile_stride >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device; the tensor-map encoders need this thread's context of the device bound. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 115712), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + CUtensorMap v_kg_tma = EncodeTma_kg_tma(a_kg_tma); + CUtensorMap v_qg_tma = EncodeTma_qg_tma(a_qg_tma); + CUtensorMap v_w_tma = EncodeTma_w_tma(a_w_tma); + CUtensorMap v_do_tma = EncodeTma_do_tma(a_do_tma); + void* v_dv1 = DataPtr(a_dv1); + void* v_gk = DataPtr(a_gk); + int32_t v_num_heads = static_cast(a_num_heads); + void* v_seq_chunk_start = DataPtr(a_seq_chunk_start); + int32_t v_num_seqs = static_cast(a_num_seqs); + float v_scale = static_cast(a_scale); + void* v_walk_seg = DataPtr(a_walk_seg); + void* v_dht = DataPtr(a_dht); + void* v_dht_index = DataPtr(a_dht_index); + int32_t v_dht_hg = static_cast(a_dht_hg); + int32_t v_dht_n = static_cast(a_dht_n); + void* v_out = DataPtr(a_out); + void* v_out_index = DataPtr(a_out_index); + int32_t v_out_hg = static_cast(a_out_hg); + int32_t v_out_n = static_cast(a_out_n); + int32_t v_tile_stride = static_cast(a_tile_stride); + void* v_dh_tape = DataPtr(a_dh_tape); + void* v_dv2 = DataPtr(a_dv2); + void* kernel_args[] = {&v_kg_tma, &v_qg_tma, &v_w_tma, &v_do_tma, &v_dv1, &v_gk, &v_num_heads, &v_seq_chunk_start, &v_num_seqs, &v_scale, &v_walk_seg, &v_dht, &v_dht_index, &v_dht_hg, &v_dht_n, &v_out, &v_out_index, &v_out_hg, &v_out_n, &v_tile_stride, &v_dh_tape, &v_dv2}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(192, 1, 1); + config.dynamicSmemBytes = 115712; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu new file mode 100644 index 00000000000..8d3fc09f434 --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_4a02b970484a9ef1f3c1_kcp_kernel.cu @@ -0,0 +1,1376 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LAUNCH_MIN_BLOCKS 2 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +extern "C" { + +__global__ __launch_bounds__(192, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_det_train_4a02b970484a9ef1f3c1_kcp(const __grid_constant__ CUtensorMap kg_tma, const __grid_constant__ CUtensorMap qg_tma, const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap do_tma, __nv_bfloat16* __restrict__ dv1, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, float scale, int* __restrict__ walk_seg, float* __restrict__ dht, int* __restrict__ dht_index, int dht_hg, int dht_n, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, __nv_bfloat16* __restrict__ dh_tape, __nv_bfloat16* __restrict__ dv2) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define kg_full_addr (mbar_base + 0) + #define kg_empty_addr (mbar_base + 8) + #define qd_full_addr (mbar_base + 16) + #define qd_empty_addr (mbar_base + 24) + #define w_full_addr (mbar_base + 32) + #define w_empty_addr (mbar_base + 40) + #define dhs_ready_addr (mbar_base + 48) + #define dv_ready_addr (mbar_base + 56) + #define dv_read_addr (mbar_base + 64) + #define bdv_ready_addr (mbar_base + 72) + #define dh_ready_addr (mbar_base + 80) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Prologue + __nv_bfloat16* kg_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int kg_panel_addr = smem + 1024; + __nv_bfloat16* qg_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int qg_panel_addr = smem + 17408; + __nv_bfloat16* w_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int w_panel_addr = smem + 33792; + __nv_bfloat16* do_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int do_panel_addr = smem + 50176; + __nv_bfloat16* kg_op = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int kg_op_addr = smem + 1024; + __nv_bfloat16* qg_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int qg_mn_addr = smem + 17408; + __nv_bfloat16* w_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int w_mn_addr = smem + 33792; + __nv_bfloat16* do_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int do_mn_addr = smem + 50176; + __nv_bfloat16* dhs = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int dhs_addr = smem + 66560; + __nv_bfloat16* dhs_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int dhs_mn_addr = smem + 66560; + __nv_bfloat16* bdv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int bdv_addr = smem + 99328; + __nv_bfloat16* bdv_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int bdv_mn_addr = smem + 99328; + + // Mbarrier init (11 barriers at smem_raw[0..88)) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'kg_pipe' --- + // kg_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // kg_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'qd_pipe' --- + // qd_full: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + // qd_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 24, 1); + // --- pipeline 'w_pipe' --- + // w_full: 1 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + // w_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // --- pipeline 'one' --- + // dhs_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 48, 128); + // dv_ready: 1 barriers, init_count=1 + mbarrier_init(smem + 56, 1); + // dv_read: 1 barriers, init_count=128 + mbarrier_init(smem + 64, 128); + // bdv_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 72, 128); + // dh_ready: 1 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 88); + if (warp == 0) { + int _tmem_hold = smem + 88; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const unsigned int taddr = tmem_addr_storage[0]; + const unsigned int tmem_acc_dh = taddr; + const unsigned int tmem_acc_dv = taddr + 128; + const unsigned int tmem_acc_p = taddr + 128; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int head = blockIdx.x % num_heads; + int walk = blockIdx.x / num_heads; + int seg = walk_seg[walk]; + int c0 = seq_chunk_start[seg]; + int nc = seq_chunk_start[seg + 1] - c0; + int row0 = c0 * 64; + int warp_in_wg = warp % 4; + int k_row = warp_in_wg * 32 + lane; + int m128_row_base = warp_in_wg * 32 << 16; + int lane_half = lane >> 4; + int t_row = warp_in_wg * 16 + (lane & 15); + int m64_row_base = warp_in_wg * 32 << 16; + float init_row[16]; + int dht_slot = dht_index[seg]; + if (dht_slot >= 0) { + int dht_tile = (head / dht_hg * dht_n + dht_slot) * dht_hg + head % dht_hg; + long long dht_base = (long long)dht_tile * (long long)tile_stride * 16384 + (long long)(k_row * 128); + #pragma unroll + for (int blk = 0; blk < 8; blk++) { + #pragma unroll + for (int q = 0; q < 4; q++) { + { + float4 _v4 = *reinterpret_cast(dht + dht_base + (long long)(blk * 16 + q * 4)); + init_row[q * 4 + 0] = _v4.x; + init_row[q * 4 + 1] = _v4.y; + init_row[q * 4 + 2] = _v4.z; + init_row[q * 4 + 3] = _v4.w; + } + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + (unsigned int)(blk * 16), init_row); + } + } else { + init_row[0] = 0.0f; + init_row[1] = 0.0f; + init_row[2] = 0.0f; + init_row[3] = 0.0f; + init_row[4] = 0.0f; + init_row[5] = 0.0f; + init_row[6] = 0.0f; + init_row[7] = 0.0f; + init_row[8] = 0.0f; + init_row[9] = 0.0f; + init_row[10] = 0.0f; + init_row[11] = 0.0f; + init_row[12] = 0.0f; + init_row[13] = 0.0f; + init_row[14] = 0.0f; + init_row[15] = 0.0f; + #pragma unroll + for (int blk_1 = 0; blk_1 < 8; blk_1++) { + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + (unsigned int)(blk_1 * 16), init_row); + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + float gk_pre = 0.0f; + float decay = 1.0f; + unsigned int d1_pk[32]; + d1_pk[0] = 0; + d1_pk[1] = 0; + d1_pk[2] = 0; + d1_pk[3] = 0; + d1_pk[4] = 0; + d1_pk[5] = 0; + d1_pk[6] = 0; + d1_pk[7] = 0; + d1_pk[8] = 0; + d1_pk[9] = 0; + d1_pk[10] = 0; + d1_pk[11] = 0; + d1_pk[12] = 0; + d1_pk[13] = 0; + d1_pk[14] = 0; + d1_pk[15] = 0; + d1_pk[16] = 0; + d1_pk[17] = 0; + d1_pk[18] = 0; + d1_pk[19] = 0; + d1_pk[20] = 0; + d1_pk[21] = 0; + d1_pk[22] = 0; + d1_pk[23] = 0; + d1_pk[24] = 0; + d1_pk[25] = 0; + d1_pk[26] = 0; + d1_pk[27] = 0; + d1_pk[28] = 0; + d1_pk[29] = 0; + d1_pk[30] = 0; + d1_pk[31] = 0; + if (nc > 0) { + long long gn_first = ((long long)(row0 + nc * 64 - 1) * (long long)num_heads + (long long)head) * 128 + (long long)k_row; + gk_pre = gk[gn_first]; + { + uint4 _uv4_1 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(lane_half * 16)); + d1_pk[0 + 0] = _uv4_1.x; + d1_pk[0 + 1] = _uv4_1.y; + d1_pk[0 + 2] = _uv4_1.z; + d1_pk[0 + 3] = _uv4_1.w; + } + { + uint4 _uv4_2 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(lane_half * 16 + 8)); + d1_pk[4 + 0] = _uv4_2.x; + d1_pk[4 + 1] = _uv4_2.y; + d1_pk[4 + 2] = _uv4_2.z; + d1_pk[4 + 3] = _uv4_2.w; + } + { + uint4 _uv4_3 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(32 + lane_half * 16)); + d1_pk[8 + 0] = _uv4_3.x; + d1_pk[8 + 1] = _uv4_3.y; + d1_pk[8 + 2] = _uv4_3.z; + d1_pk[8 + 3] = _uv4_3.w; + } + { + uint4 _uv4_4 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(32 + lane_half * 16 + 8)); + d1_pk[12 + 0] = _uv4_4.x; + d1_pk[12 + 1] = _uv4_4.y; + d1_pk[12 + 2] = _uv4_4.z; + d1_pk[12 + 3] = _uv4_4.w; + } + { + uint4 _uv4_5 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(64 + lane_half * 16)); + d1_pk[16 + 0] = _uv4_5.x; + d1_pk[16 + 1] = _uv4_5.y; + d1_pk[16 + 2] = _uv4_5.z; + d1_pk[16 + 3] = _uv4_5.w; + } + { + uint4 _uv4_6 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(64 + lane_half * 16 + 8)); + d1_pk[20 + 0] = _uv4_6.x; + d1_pk[20 + 1] = _uv4_6.y; + d1_pk[20 + 2] = _uv4_6.z; + d1_pk[20 + 3] = _uv4_6.w; + } + { + uint4 _uv4_7 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(96 + lane_half * 16)); + d1_pk[24 + 0] = _uv4_7.x; + d1_pk[24 + 1] = _uv4_7.y; + d1_pk[24 + 2] = _uv4_7.z; + d1_pk[24 + 3] = _uv4_7.w; + } + { + uint4 _uv4_8 = *reinterpret_cast(dv1 + ((long long)(row0 + (nc - 1) * 64 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(96 + lane_half * 16 + 8)); + d1_pk[28 + 0] = _uv4_8.x; + d1_pk[28 + 1] = _uv4_8.y; + d1_pk[28 + 2] = _uv4_8.z; + d1_pk[28 + 3] = _uv4_8.w; + } + float dh_blk_0[16]; + tmem_ld_x16(&dh_blk_0[0], taddr + (unsigned int)m128_row_base + 0); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_0[_lp*2 + 0], dh_blk_0[_lp*2+1 + 0])); + state_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_1 = 0; q_1 < 2; q_1++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)(q_1 * 8 / 64 * 16384 + k_row * 128 + q_1 * 8 % 64 * 2 ^ (q_1 * 8 / 64 * 16384 + k_row * 128 + q_1 * 8 % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16[q_1 * 4]), "r"(state_bf16[q_1 * 4 + 1]), "r"(state_bf16[q_1 * 4 + 2]), "r"(state_bf16[q_1 * 4 + 3]) : "memory"); + } + float dh_blk_1[16]; + tmem_ld_x16(&dh_blk_1[0], taddr + (unsigned int)m128_row_base + 16); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16_0[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_1[_lp*2 + 0], dh_blk_1[_lp*2+1 + 0])); + state_bf16_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_2 = 0; q_2 < 2; q_2++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((16 + q_2 * 8) / 64 * 16384 + k_row * 128 + (16 + q_2 * 8) % 64 * 2 ^ ((16 + q_2 * 8) / 64 * 16384 + k_row * 128 + (16 + q_2 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16_0[q_2 * 4]), "r"(state_bf16_0[q_2 * 4 + 1]), "r"(state_bf16_0[q_2 * 4 + 2]), "r"(state_bf16_0[q_2 * 4 + 3]) : "memory"); + } + float dh_blk_2[16]; + tmem_ld_x16(&dh_blk_2[0], taddr + (unsigned int)m128_row_base + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16_1[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_2[_lp*2 + 0], dh_blk_2[_lp*2+1 + 0])); + state_bf16_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_3 = 0; q_3 < 2; q_3++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((32 + q_3 * 8) / 64 * 16384 + k_row * 128 + (32 + q_3 * 8) % 64 * 2 ^ ((32 + q_3 * 8) / 64 * 16384 + k_row * 128 + (32 + q_3 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16_1[q_3 * 4]), "r"(state_bf16_1[q_3 * 4 + 1]), "r"(state_bf16_1[q_3 * 4 + 2]), "r"(state_bf16_1[q_3 * 4 + 3]) : "memory"); + } + float dh_blk_3[16]; + tmem_ld_x16(&dh_blk_3[0], taddr + (unsigned int)m128_row_base + 48); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16_2[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_3[_lp*2 + 0], dh_blk_3[_lp*2+1 + 0])); + state_bf16_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_4 = 0; q_4 < 2; q_4++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((48 + q_4 * 8) / 64 * 16384 + k_row * 128 + (48 + q_4 * 8) % 64 * 2 ^ ((48 + q_4 * 8) / 64 * 16384 + k_row * 128 + (48 + q_4 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16_2[q_4 * 4]), "r"(state_bf16_2[q_4 * 4 + 1]), "r"(state_bf16_2[q_4 * 4 + 2]), "r"(state_bf16_2[q_4 * 4 + 3]) : "memory"); + } + float dh_blk_4[16]; + tmem_ld_x16(&dh_blk_4[0], taddr + (unsigned int)m128_row_base + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16_3[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_4[_lp*2 + 0], dh_blk_4[_lp*2+1 + 0])); + state_bf16_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_5 = 0; q_5 < 2; q_5++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((64 + q_5 * 8) / 64 * 16384 + k_row * 128 + (64 + q_5 * 8) % 64 * 2 ^ ((64 + q_5 * 8) / 64 * 16384 + k_row * 128 + (64 + q_5 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16_3[q_5 * 4]), "r"(state_bf16_3[q_5 * 4 + 1]), "r"(state_bf16_3[q_5 * 4 + 2]), "r"(state_bf16_3[q_5 * 4 + 3]) : "memory"); + } + float dh_blk_5[16]; + tmem_ld_x16(&dh_blk_5[0], taddr + (unsigned int)m128_row_base + 80); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16_4[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_5[_lp*2 + 0], dh_blk_5[_lp*2+1 + 0])); + state_bf16_4[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_6 = 0; q_6 < 2; q_6++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((80 + q_6 * 8) / 64 * 16384 + k_row * 128 + (80 + q_6 * 8) % 64 * 2 ^ ((80 + q_6 * 8) / 64 * 16384 + k_row * 128 + (80 + q_6 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16_4[q_6 * 4]), "r"(state_bf16_4[q_6 * 4 + 1]), "r"(state_bf16_4[q_6 * 4 + 2]), "r"(state_bf16_4[q_6 * 4 + 3]) : "memory"); + } + float dh_blk_6[16]; + tmem_ld_x16(&dh_blk_6[0], taddr + (unsigned int)m128_row_base + 96); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16_5[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_6[_lp*2 + 0], dh_blk_6[_lp*2+1 + 0])); + state_bf16_5[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_7 = 0; q_7 < 2; q_7++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((96 + q_7 * 8) / 64 * 16384 + k_row * 128 + (96 + q_7 * 8) % 64 * 2 ^ ((96 + q_7 * 8) / 64 * 16384 + k_row * 128 + (96 + q_7 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16_5[q_7 * 4]), "r"(state_bf16_5[q_7 * 4 + 1]), "r"(state_bf16_5[q_7 * 4 + 2]), "r"(state_bf16_5[q_7 * 4 + 3]) : "memory"); + } + float dh_blk_7[16]; + tmem_ld_x16(&dh_blk_7[0], taddr + (unsigned int)m128_row_base + 112); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int state_bf16_6[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dh_blk_7[_lp*2 + 0], dh_blk_7[_lp*2+1 + 0])); + state_bf16_6[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_8 = 0; q_8 < 2; q_8++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((112 + q_8 * 8) / 64 * 16384 + k_row * 128 + (112 + q_8 * 8) % 64 * 2 ^ ((112 + q_8 * 8) / 64 * 16384 + k_row * 128 + (112 + q_8 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(state_bf16_6[q_8 * 4]), "r"(state_bf16_6[q_8 * 4 + 1]), "r"(state_bf16_6[q_8 * 4 + 2]), "r"(state_bf16_6[q_8 * 4 + 3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(dhs_ready_addr); + float _exp2_9 = approx_exp2(gk_pre); + decay = _exp2_9; + float dh_blk_8[16]; + tmem_ld_x16(&dh_blk_8[0], taddr + (unsigned int)m128_row_base + 0); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_0[16]; + #pragma unroll + for (int j_0 = 0; j_0 < 16; j_0++) { + sc_0[j_0] = dh_blk_8[j_0] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 0, sc_0); + float dh_blk_9[16]; + tmem_ld_x16(&dh_blk_9[0], taddr + (unsigned int)m128_row_base + 16); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_1[16]; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + sc_1[j_1] = dh_blk_9[j_1] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 16, sc_1); + float dh_blk_10[16]; + tmem_ld_x16(&dh_blk_10[0], taddr + (unsigned int)m128_row_base + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_2[16]; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + sc_2[j_2] = dh_blk_10[j_2] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 32, sc_2); + float dh_blk_11[16]; + tmem_ld_x16(&dh_blk_11[0], taddr + (unsigned int)m128_row_base + 48); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_3[16]; + #pragma unroll + for (int j_3 = 0; j_3 < 16; j_3++) { + sc_3[j_3] = dh_blk_11[j_3] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 48, sc_3); + float dh_blk_12[16]; + tmem_ld_x16(&dh_blk_12[0], taddr + (unsigned int)m128_row_base + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_4[16]; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + sc_4[j_4] = dh_blk_12[j_4] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 64, sc_4); + float dh_blk_13[16]; + tmem_ld_x16(&dh_blk_13[0], taddr + (unsigned int)m128_row_base + 80); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_5[16]; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + sc_5[j_5] = dh_blk_13[j_5] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 80, sc_5); + float dh_blk_14[16]; + tmem_ld_x16(&dh_blk_14[0], taddr + (unsigned int)m128_row_base + 96); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_6[16]; + #pragma unroll + for (int j_6 = 0; j_6 < 16; j_6++) { + sc_6[j_6] = dh_blk_14[j_6] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 96, sc_6); + float dh_blk_15[16]; + tmem_ld_x16(&dh_blk_15[0], taddr + (unsigned int)m128_row_base + 112); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float sc_7[16]; + #pragma unroll + for (int j_7 = 0; j_7 < 16; j_7++) { + sc_7[j_7] = dh_blk_15[j_7] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 112, sc_7); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + unsigned int _phase_dv_ready_0 = 0; + unsigned int _phase_dh_ready_0 = 0; + #pragma unroll 1 + for (int rc = 0; rc < nc; rc++) { + int c = nc - 1 - rc; + mbarrier_wait(dv_ready_addr, _phase_dv_ready_0); + _phase_dv_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + long long t_base = ((long long)(row0 + c * 64 + t_row) * (long long)num_heads + (long long)head) * 128; + float dv_a_0[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&dv_a_0[0])), "=r"(*reinterpret_cast(&dv_a_0[1])), "=r"(*reinterpret_cast(&dv_a_0[2])), "=r"(*reinterpret_cast(&dv_a_0[3])), "=r"(*reinterpret_cast(&dv_a_0[4])), "=r"(*reinterpret_cast(&dv_a_0[5])), "=r"(*reinterpret_cast(&dv_a_0[6])), "=r"(*reinterpret_cast(&dv_a_0[7])), "=r"(*reinterpret_cast(&dv_a_0[8])), "=r"(*reinterpret_cast(&dv_a_0[9])), "=r"(*reinterpret_cast(&dv_a_0[10])), "=r"(*reinterpret_cast(&dv_a_0[11])), "=r"(*reinterpret_cast(&dv_a_0[12])), "=r"(*reinterpret_cast(&dv_a_0[13])), "=r"(*reinterpret_cast(&dv_a_0[14])), "=r"(*reinterpret_cast(&dv_a_0[15])) + : "r"(taddr + 128 + (unsigned int)m64_row_base + 0)); + float dv_b_0[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&dv_b_0[0])), "=r"(*reinterpret_cast(&dv_b_0[1])), "=r"(*reinterpret_cast(&dv_b_0[2])), "=r"(*reinterpret_cast(&dv_b_0[3])), "=r"(*reinterpret_cast(&dv_b_0[4])), "=r"(*reinterpret_cast(&dv_b_0[5])), "=r"(*reinterpret_cast(&dv_b_0[6])), "=r"(*reinterpret_cast(&dv_b_0[7])), "=r"(*reinterpret_cast(&dv_b_0[8])), "=r"(*reinterpret_cast(&dv_b_0[9])), "=r"(*reinterpret_cast(&dv_b_0[10])), "=r"(*reinterpret_cast(&dv_b_0[11])), "=r"(*reinterpret_cast(&dv_b_0[12])), "=r"(*reinterpret_cast(&dv_b_0[13])), "=r"(*reinterpret_cast(&dv_b_0[14])), "=r"(*reinterpret_cast(&dv_b_0[15])) + : "r"(taddr + 128 + (unsigned int)m64_row_base + 32)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + int col = lane_half * 16; + float d1[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&d1[_pair * 2])[0] = __uint_as_float(static_cast(d1_pk[_pair]) << 16); + (&d1[_pair * 2])[1] = __uint_as_float(static_cast(d1_pk[_pair]) & 0xffff0000u); + } + float bdvv[16]; + #pragma unroll + for (int j = 0; j < 16; j++) { + bdvv[j] = dv_a_0[j] + d1[j]; + } + unsigned int bdvv_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bdvv[_lp*2 + 0], bdvv[_lp*2+1 + 0])); + bdvv_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_9 = 0; q_9 < 2; q_9++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bdv_addr + (unsigned int)((col + q_9 * 8) / 64 * 8192 + t_row * 128 + (col + q_9 * 8) % 64 * 2 ^ ((col + q_9 * 8) / 64 * 8192 + t_row * 128 + (col + q_9 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bdvv_bf16[q_9 * 4]), "r"(bdvv_bf16[q_9 * 4 + 1]), "r"(bdvv_bf16[q_9 * 4 + 2]), "r"(bdvv_bf16[q_9 * 4 + 3]) : "memory"); + } + int col_0 = 32 + lane_half * 16; + float d1_0[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&d1_0[_pair * 2])[0] = __uint_as_float(static_cast(d1_pk[8 + _pair]) << 16); + (&d1_0[_pair * 2])[1] = __uint_as_float(static_cast(d1_pk[8 + _pair]) & 0xffff0000u); + } + float bdvv_0[16]; + #pragma unroll + for (int j_8 = 0; j_8 < 16; j_8++) { + bdvv_0[j_8] = dv_b_0[j_8] + d1_0[j_8]; + } + unsigned int bdvv_bf16_0[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bdvv_0[_lp*2 + 0], bdvv_0[_lp*2+1 + 0])); + bdvv_bf16_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_10 = 0; q_10 < 2; q_10++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bdv_addr + (unsigned int)((col_0 + q_10 * 8) / 64 * 8192 + t_row * 128 + (col_0 + q_10 * 8) % 64 * 2 ^ ((col_0 + q_10 * 8) / 64 * 8192 + t_row * 128 + (col_0 + q_10 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bdvv_bf16_0[q_10 * 4]), "r"(bdvv_bf16_0[q_10 * 4 + 1]), "r"(bdvv_bf16_0[q_10 * 4 + 2]), "r"(bdvv_bf16_0[q_10 * 4 + 3]) : "memory"); + } + float dv_a_1[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&dv_a_1[0])), "=r"(*reinterpret_cast(&dv_a_1[1])), "=r"(*reinterpret_cast(&dv_a_1[2])), "=r"(*reinterpret_cast(&dv_a_1[3])), "=r"(*reinterpret_cast(&dv_a_1[4])), "=r"(*reinterpret_cast(&dv_a_1[5])), "=r"(*reinterpret_cast(&dv_a_1[6])), "=r"(*reinterpret_cast(&dv_a_1[7])), "=r"(*reinterpret_cast(&dv_a_1[8])), "=r"(*reinterpret_cast(&dv_a_1[9])), "=r"(*reinterpret_cast(&dv_a_1[10])), "=r"(*reinterpret_cast(&dv_a_1[11])), "=r"(*reinterpret_cast(&dv_a_1[12])), "=r"(*reinterpret_cast(&dv_a_1[13])), "=r"(*reinterpret_cast(&dv_a_1[14])), "=r"(*reinterpret_cast(&dv_a_1[15])) + : "r"(taddr + 128 + (unsigned int)m64_row_base + 64)); + float dv_b_1[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&dv_b_1[0])), "=r"(*reinterpret_cast(&dv_b_1[1])), "=r"(*reinterpret_cast(&dv_b_1[2])), "=r"(*reinterpret_cast(&dv_b_1[3])), "=r"(*reinterpret_cast(&dv_b_1[4])), "=r"(*reinterpret_cast(&dv_b_1[5])), "=r"(*reinterpret_cast(&dv_b_1[6])), "=r"(*reinterpret_cast(&dv_b_1[7])), "=r"(*reinterpret_cast(&dv_b_1[8])), "=r"(*reinterpret_cast(&dv_b_1[9])), "=r"(*reinterpret_cast(&dv_b_1[10])), "=r"(*reinterpret_cast(&dv_b_1[11])), "=r"(*reinterpret_cast(&dv_b_1[12])), "=r"(*reinterpret_cast(&dv_b_1[13])), "=r"(*reinterpret_cast(&dv_b_1[14])), "=r"(*reinterpret_cast(&dv_b_1[15])) + : "r"(taddr + 128 + (unsigned int)m64_row_base + 96)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(dv_read_addr); + int col_1 = 64 + lane_half * 16; + float d1_1[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&d1_1[_pair * 2])[0] = __uint_as_float(static_cast(d1_pk[16 + _pair]) << 16); + (&d1_1[_pair * 2])[1] = __uint_as_float(static_cast(d1_pk[16 + _pair]) & 0xffff0000u); + } + float bdvv_1[16]; + #pragma unroll + for (int j_9 = 0; j_9 < 16; j_9++) { + bdvv_1[j_9] = dv_a_1[j_9] + d1_1[j_9]; + } + unsigned int bdvv_bf16_1[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bdvv_1[_lp*2 + 0], bdvv_1[_lp*2+1 + 0])); + bdvv_bf16_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_11 = 0; q_11 < 2; q_11++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bdv_addr + (unsigned int)((col_1 + q_11 * 8) / 64 * 8192 + t_row * 128 + (col_1 + q_11 * 8) % 64 * 2 ^ ((col_1 + q_11 * 8) / 64 * 8192 + t_row * 128 + (col_1 + q_11 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bdvv_bf16_1[q_11 * 4]), "r"(bdvv_bf16_1[q_11 * 4 + 1]), "r"(bdvv_bf16_1[q_11 * 4 + 2]), "r"(bdvv_bf16_1[q_11 * 4 + 3]) : "memory"); + } + int col_2 = 96 + lane_half * 16; + float d1_2[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&d1_2[_pair * 2])[0] = __uint_as_float(static_cast(d1_pk[24 + _pair]) << 16); + (&d1_2[_pair * 2])[1] = __uint_as_float(static_cast(d1_pk[24 + _pair]) & 0xffff0000u); + } + float bdvv_2[16]; + #pragma unroll + for (int j_10 = 0; j_10 < 16; j_10++) { + bdvv_2[j_10] = dv_b_1[j_10] + d1_2[j_10]; + } + unsigned int bdvv_bf16_2[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bdvv_2[_lp*2 + 0], bdvv_2[_lp*2+1 + 0])); + bdvv_bf16_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_12 = 0; q_12 < 2; q_12++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bdv_addr + (unsigned int)((col_2 + q_12 * 8) / 64 * 8192 + t_row * 128 + (col_2 + q_12 * 8) % 64 * 2 ^ ((col_2 + q_12 * 8) / 64 * 8192 + t_row * 128 + (col_2 + q_12 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bdvv_bf16_2[q_12 * 4]), "r"(bdvv_bf16_2[q_12 * 4 + 1]), "r"(bdvv_bf16_2[q_12 * 4 + 2]), "r"(bdvv_bf16_2[q_12 * 4 + 3]) : "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(bdv_ready_addr); + if (nc > rc + 1) { + long long gn_next = ((long long)(row0 + c * 64 - 1) * (long long)num_heads + (long long)head) * 128 + (long long)k_row; + gk_pre = gk[gn_next]; + { + uint4 _uv4_10 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(lane_half * 16)); + d1_pk[0 + 0] = _uv4_10.x; + d1_pk[0 + 1] = _uv4_10.y; + d1_pk[0 + 2] = _uv4_10.z; + d1_pk[0 + 3] = _uv4_10.w; + } + { + uint4 _uv4_11 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(lane_half * 16 + 8)); + d1_pk[4 + 0] = _uv4_11.x; + d1_pk[4 + 1] = _uv4_11.y; + d1_pk[4 + 2] = _uv4_11.z; + d1_pk[4 + 3] = _uv4_11.w; + } + { + uint4 _uv4_12 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(32 + lane_half * 16)); + d1_pk[8 + 0] = _uv4_12.x; + d1_pk[8 + 1] = _uv4_12.y; + d1_pk[8 + 2] = _uv4_12.z; + d1_pk[8 + 3] = _uv4_12.w; + } + { + uint4 _uv4_13 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(32 + lane_half * 16 + 8)); + d1_pk[12 + 0] = _uv4_13.x; + d1_pk[12 + 1] = _uv4_13.y; + d1_pk[12 + 2] = _uv4_13.z; + d1_pk[12 + 3] = _uv4_13.w; + } + { + uint4 _uv4_14 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(64 + lane_half * 16)); + d1_pk[16 + 0] = _uv4_14.x; + d1_pk[16 + 1] = _uv4_14.y; + d1_pk[16 + 2] = _uv4_14.z; + d1_pk[16 + 3] = _uv4_14.w; + } + { + uint4 _uv4_15 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(64 + lane_half * 16 + 8)); + d1_pk[20 + 0] = _uv4_15.x; + d1_pk[20 + 1] = _uv4_15.y; + d1_pk[20 + 2] = _uv4_15.z; + d1_pk[20 + 3] = _uv4_15.w; + } + { + uint4 _uv4_16 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(96 + lane_half * 16)); + d1_pk[24 + 0] = _uv4_16.x; + d1_pk[24 + 1] = _uv4_16.y; + d1_pk[24 + 2] = _uv4_16.z; + d1_pk[24 + 3] = _uv4_16.w; + } + { + uint4 _uv4_17 = *reinterpret_cast(dv1 + t_base - (long long)(64 * num_heads * 128) + (long long)(96 + lane_half * 16 + 8)); + d1_pk[28 + 0] = _uv4_17.x; + d1_pk[28 + 1] = _uv4_17.y; + d1_pk[28 + 2] = _uv4_17.z; + d1_pk[28 + 3] = _uv4_17.w; + } + } + mbarrier_wait(dh_ready_addr, _phase_dh_ready_0); + _phase_dh_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + if (nc > rc + 1) { + float _exp2_18 = approx_exp2(gk_pre); + decay = _exp2_18; + float dh_blk[16]; + tmem_ld_x16(&dh_blk[0], taddr + (unsigned int)m128_row_base + 0); + float p_blk[16]; + tmem_ld_x16(&p_blk[0], taddr + 128 + (unsigned int)m128_row_base + 0); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd[16]; + #pragma unroll + for (int j_11 = 0; j_11 < 16; j_11++) { + upd[j_11] = __fmaf_rn(p_blk[j_11], scale, dh_blk[j_11]); + } + unsigned int upd_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd[_lp*2 + 0], upd[_lp*2+1 + 0])); + upd_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_13 = 0; q_13 < 2; q_13++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)(q_13 * 8 / 64 * 16384 + k_row * 128 + q_13 * 8 % 64 * 2 ^ (q_13 * 8 / 64 * 16384 + k_row * 128 + q_13 * 8 % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16[q_13 * 4]), "r"(upd_bf16[q_13 * 4 + 1]), "r"(upd_bf16[q_13 * 4 + 2]), "r"(upd_bf16[q_13 * 4 + 3]) : "memory"); + } + float sc[16]; + #pragma unroll + for (int j_12 = 0; j_12 < 16; j_12++) { + sc[j_12] = upd[j_12] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 0, sc); + float dh_blk_16[16]; + tmem_ld_x16(&dh_blk_16[0], taddr + (unsigned int)m128_row_base + 16); + float p_blk_0[16]; + tmem_ld_x16(&p_blk_0[0], taddr + 128 + (unsigned int)m128_row_base + 16); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_0[16]; + #pragma unroll + for (int j_13 = 0; j_13 < 16; j_13++) { + upd_0[j_13] = __fmaf_rn(p_blk_0[j_13], scale, dh_blk_16[j_13]); + } + unsigned int upd_bf16_0[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd_0[_lp*2 + 0], upd_0[_lp*2+1 + 0])); + upd_bf16_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_14 = 0; q_14 < 2; q_14++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((16 + q_14 * 8) / 64 * 16384 + k_row * 128 + (16 + q_14 * 8) % 64 * 2 ^ ((16 + q_14 * 8) / 64 * 16384 + k_row * 128 + (16 + q_14 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16_0[q_14 * 4]), "r"(upd_bf16_0[q_14 * 4 + 1]), "r"(upd_bf16_0[q_14 * 4 + 2]), "r"(upd_bf16_0[q_14 * 4 + 3]) : "memory"); + } + float sc_8[16]; + #pragma unroll + for (int j_14 = 0; j_14 < 16; j_14++) { + sc_8[j_14] = upd_0[j_14] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 16, sc_8); + float dh_blk_17[16]; + tmem_ld_x16(&dh_blk_17[0], taddr + (unsigned int)m128_row_base + 32); + float p_blk_1[16]; + tmem_ld_x16(&p_blk_1[0], taddr + 128 + (unsigned int)m128_row_base + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_1[16]; + #pragma unroll + for (int j_15 = 0; j_15 < 16; j_15++) { + upd_1[j_15] = __fmaf_rn(p_blk_1[j_15], scale, dh_blk_17[j_15]); + } + unsigned int upd_bf16_1[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd_1[_lp*2 + 0], upd_1[_lp*2+1 + 0])); + upd_bf16_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_15 = 0; q_15 < 2; q_15++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((32 + q_15 * 8) / 64 * 16384 + k_row * 128 + (32 + q_15 * 8) % 64 * 2 ^ ((32 + q_15 * 8) / 64 * 16384 + k_row * 128 + (32 + q_15 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16_1[q_15 * 4]), "r"(upd_bf16_1[q_15 * 4 + 1]), "r"(upd_bf16_1[q_15 * 4 + 2]), "r"(upd_bf16_1[q_15 * 4 + 3]) : "memory"); + } + float sc_9[16]; + #pragma unroll + for (int j_16 = 0; j_16 < 16; j_16++) { + sc_9[j_16] = upd_1[j_16] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 32, sc_9); + float dh_blk_18[16]; + tmem_ld_x16(&dh_blk_18[0], taddr + (unsigned int)m128_row_base + 48); + float p_blk_2[16]; + tmem_ld_x16(&p_blk_2[0], taddr + 128 + (unsigned int)m128_row_base + 48); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_2[16]; + #pragma unroll + for (int j_17 = 0; j_17 < 16; j_17++) { + upd_2[j_17] = __fmaf_rn(p_blk_2[j_17], scale, dh_blk_18[j_17]); + } + unsigned int upd_bf16_2[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd_2[_lp*2 + 0], upd_2[_lp*2+1 + 0])); + upd_bf16_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_16 = 0; q_16 < 2; q_16++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((48 + q_16 * 8) / 64 * 16384 + k_row * 128 + (48 + q_16 * 8) % 64 * 2 ^ ((48 + q_16 * 8) / 64 * 16384 + k_row * 128 + (48 + q_16 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16_2[q_16 * 4]), "r"(upd_bf16_2[q_16 * 4 + 1]), "r"(upd_bf16_2[q_16 * 4 + 2]), "r"(upd_bf16_2[q_16 * 4 + 3]) : "memory"); + } + float sc_10[16]; + #pragma unroll + for (int j_18 = 0; j_18 < 16; j_18++) { + sc_10[j_18] = upd_2[j_18] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 48, sc_10); + float dh_blk_19[16]; + tmem_ld_x16(&dh_blk_19[0], taddr + (unsigned int)m128_row_base + 64); + float p_blk_3[16]; + tmem_ld_x16(&p_blk_3[0], taddr + 128 + (unsigned int)m128_row_base + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_3[16]; + #pragma unroll + for (int j_19 = 0; j_19 < 16; j_19++) { + upd_3[j_19] = __fmaf_rn(p_blk_3[j_19], scale, dh_blk_19[j_19]); + } + unsigned int upd_bf16_3[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd_3[_lp*2 + 0], upd_3[_lp*2+1 + 0])); + upd_bf16_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_17 = 0; q_17 < 2; q_17++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((64 + q_17 * 8) / 64 * 16384 + k_row * 128 + (64 + q_17 * 8) % 64 * 2 ^ ((64 + q_17 * 8) / 64 * 16384 + k_row * 128 + (64 + q_17 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16_3[q_17 * 4]), "r"(upd_bf16_3[q_17 * 4 + 1]), "r"(upd_bf16_3[q_17 * 4 + 2]), "r"(upd_bf16_3[q_17 * 4 + 3]) : "memory"); + } + float sc_11[16]; + #pragma unroll + for (int j_20 = 0; j_20 < 16; j_20++) { + sc_11[j_20] = upd_3[j_20] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 64, sc_11); + float dh_blk_20[16]; + tmem_ld_x16(&dh_blk_20[0], taddr + (unsigned int)m128_row_base + 80); + float p_blk_4[16]; + tmem_ld_x16(&p_blk_4[0], taddr + 128 + (unsigned int)m128_row_base + 80); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_4[16]; + #pragma unroll + for (int j_21 = 0; j_21 < 16; j_21++) { + upd_4[j_21] = __fmaf_rn(p_blk_4[j_21], scale, dh_blk_20[j_21]); + } + unsigned int upd_bf16_4[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd_4[_lp*2 + 0], upd_4[_lp*2+1 + 0])); + upd_bf16_4[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_18 = 0; q_18 < 2; q_18++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((80 + q_18 * 8) / 64 * 16384 + k_row * 128 + (80 + q_18 * 8) % 64 * 2 ^ ((80 + q_18 * 8) / 64 * 16384 + k_row * 128 + (80 + q_18 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16_4[q_18 * 4]), "r"(upd_bf16_4[q_18 * 4 + 1]), "r"(upd_bf16_4[q_18 * 4 + 2]), "r"(upd_bf16_4[q_18 * 4 + 3]) : "memory"); + } + float sc_12[16]; + #pragma unroll + for (int j_22 = 0; j_22 < 16; j_22++) { + sc_12[j_22] = upd_4[j_22] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 80, sc_12); + float dh_blk_21[16]; + tmem_ld_x16(&dh_blk_21[0], taddr + (unsigned int)m128_row_base + 96); + float p_blk_5[16]; + tmem_ld_x16(&p_blk_5[0], taddr + 128 + (unsigned int)m128_row_base + 96); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_5[16]; + #pragma unroll + for (int j_23 = 0; j_23 < 16; j_23++) { + upd_5[j_23] = __fmaf_rn(p_blk_5[j_23], scale, dh_blk_21[j_23]); + } + unsigned int upd_bf16_5[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd_5[_lp*2 + 0], upd_5[_lp*2+1 + 0])); + upd_bf16_5[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_19 = 0; q_19 < 2; q_19++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((96 + q_19 * 8) / 64 * 16384 + k_row * 128 + (96 + q_19 * 8) % 64 * 2 ^ ((96 + q_19 * 8) / 64 * 16384 + k_row * 128 + (96 + q_19 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16_5[q_19 * 4]), "r"(upd_bf16_5[q_19 * 4 + 1]), "r"(upd_bf16_5[q_19 * 4 + 2]), "r"(upd_bf16_5[q_19 * 4 + 3]) : "memory"); + } + float sc_13[16]; + #pragma unroll + for (int j_24 = 0; j_24 < 16; j_24++) { + sc_13[j_24] = upd_5[j_24] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 96, sc_13); + float dh_blk_22[16]; + tmem_ld_x16(&dh_blk_22[0], taddr + (unsigned int)m128_row_base + 112); + float p_blk_6[16]; + tmem_ld_x16(&p_blk_6[0], taddr + 128 + (unsigned int)m128_row_base + 112); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_6[16]; + #pragma unroll + for (int j_25 = 0; j_25 < 16; j_25++) { + upd_6[j_25] = __fmaf_rn(p_blk_6[j_25], scale, dh_blk_22[j_25]); + } + unsigned int upd_bf16_6[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(upd_6[_lp*2 + 0], upd_6[_lp*2+1 + 0])); + upd_bf16_6[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_20 = 0; q_20 < 2; q_20++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((dhs_addr + (unsigned int)((112 + q_20 * 8) / 64 * 16384 + k_row * 128 + (112 + q_20 * 8) % 64 * 2 ^ ((112 + q_20 * 8) / 64 * 16384 + k_row * 128 + (112 + q_20 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(upd_bf16_6[q_20 * 4]), "r"(upd_bf16_6[q_20 * 4 + 1]), "r"(upd_bf16_6[q_20 * 4 + 2]), "r"(upd_bf16_6[q_20 * 4 + 3]) : "memory"); + } + float sc_14[16]; + #pragma unroll + for (int j_26 = 0; j_26 < 16; j_26++) { + sc_14[j_26] = upd_6[j_26] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 112, sc_14); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(dhs_ready_addr); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + if (nc > 0) { + float dh_blk_23[16]; + tmem_ld_x16(&dh_blk_23[0], taddr + (unsigned int)m128_row_base + 0); + float p_blk_7[16]; + tmem_ld_x16(&p_blk_7[0], taddr + 128 + (unsigned int)m128_row_base + 0); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_7[16]; + #pragma unroll + for (int j_27 = 0; j_27 < 16; j_27++) { + upd_7[j_27] = __fmaf_rn(p_blk_7[j_27], scale, dh_blk_23[j_27]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 0, upd_7); + float dh_blk_23_1[16]; + tmem_ld_x16(&dh_blk_23_1[0], taddr + (unsigned int)m128_row_base + 16); + float p_blk_7_1[16]; + tmem_ld_x16(&p_blk_7_1[0], taddr + 128 + (unsigned int)m128_row_base + 16); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_7_1[16]; + #pragma unroll + for (int j_28 = 0; j_28 < 16; j_28++) { + upd_7_1[j_28] = __fmaf_rn(p_blk_7_1[j_28], scale, dh_blk_23_1[j_28]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 16, upd_7_1); + float dh_blk_24[16]; + tmem_ld_x16(&dh_blk_24[0], taddr + (unsigned int)m128_row_base + 32); + float p_blk_8[16]; + tmem_ld_x16(&p_blk_8[0], taddr + 128 + (unsigned int)m128_row_base + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_8[16]; + #pragma unroll + for (int j_29 = 0; j_29 < 16; j_29++) { + upd_8[j_29] = __fmaf_rn(p_blk_8[j_29], scale, dh_blk_24[j_29]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 32, upd_8); + float dh_blk_25[16]; + tmem_ld_x16(&dh_blk_25[0], taddr + (unsigned int)m128_row_base + 48); + float p_blk_9[16]; + tmem_ld_x16(&p_blk_9[0], taddr + 128 + (unsigned int)m128_row_base + 48); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_9[16]; + #pragma unroll + for (int j_30 = 0; j_30 < 16; j_30++) { + upd_9[j_30] = __fmaf_rn(p_blk_9[j_30], scale, dh_blk_25[j_30]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 48, upd_9); + float dh_blk_26[16]; + tmem_ld_x16(&dh_blk_26[0], taddr + (unsigned int)m128_row_base + 64); + float p_blk_10[16]; + tmem_ld_x16(&p_blk_10[0], taddr + 128 + (unsigned int)m128_row_base + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_10[16]; + #pragma unroll + for (int j_31 = 0; j_31 < 16; j_31++) { + upd_10[j_31] = __fmaf_rn(p_blk_10[j_31], scale, dh_blk_26[j_31]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 64, upd_10); + float dh_blk_27[16]; + tmem_ld_x16(&dh_blk_27[0], taddr + (unsigned int)m128_row_base + 80); + float p_blk_11[16]; + tmem_ld_x16(&p_blk_11[0], taddr + 128 + (unsigned int)m128_row_base + 80); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_11[16]; + #pragma unroll + for (int j_32 = 0; j_32 < 16; j_32++) { + upd_11[j_32] = __fmaf_rn(p_blk_11[j_32], scale, dh_blk_27[j_32]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 80, upd_11); + float dh_blk_28[16]; + tmem_ld_x16(&dh_blk_28[0], taddr + (unsigned int)m128_row_base + 96); + float p_blk_12[16]; + tmem_ld_x16(&p_blk_12[0], taddr + 128 + (unsigned int)m128_row_base + 96); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_12[16]; + #pragma unroll + for (int j_33 = 0; j_33 < 16; j_33++) { + upd_12[j_33] = __fmaf_rn(p_blk_12[j_33], scale, dh_blk_28[j_33]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 96, upd_12); + float dh_blk_29[16]; + tmem_ld_x16(&dh_blk_29[0], taddr + (unsigned int)m128_row_base + 112); + float p_blk_13[16]; + tmem_ld_x16(&p_blk_13[0], taddr + 128 + (unsigned int)m128_row_base + 112); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float upd_13[16]; + #pragma unroll + for (int j_34 = 0; j_34 < 16; j_34++) { + upd_13[j_34] = __fmaf_rn(p_blk_13[j_34], scale, dh_blk_29[j_34]); + } + tmem_st_x16_f32(taddr + (unsigned int)m128_row_base + 112, upd_13); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + int out_slot = out_index[seg]; + if (out_slot >= 0) { + int out_tile = (head / out_hg * out_n + out_slot) * out_hg + head % out_hg; + long long out_base = (long long)out_tile * (long long)tile_stride * 16384 + (long long)(k_row * 128); + #pragma unroll + for (int blk_2 = 0; blk_2 < 8; blk_2++) { + float dh_out[16]; + tmem_ld_x16(&dh_out[0], taddr + (unsigned int)m128_row_base + (unsigned int)(blk_2 * 16)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int q_21 = 0; q_21 < 4; q_21++) { + { + float4 _v4 = make_float4(dh_out[q_21 * 4 + 0], dh_out[q_21 * 4 + 1], dh_out[q_21 * 4 + 2], dh_out[q_21 * 4 + 3]); + *reinterpret_cast(out + out_base + (long long)(blk_2 * 16 + q_21 * 4)) = _v4; + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("barrier.sync 1, 128;" ::: "memory"); + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + } + } + } + // ---- Role: load ---- + if (warp == 4) { + { // load_main + int head_1 = blockIdx.x % num_heads; + int walk_1 = blockIdx.x / num_heads; + int seg_1 = walk_seg[walk_1]; + int c0_1 = seq_chunk_start[seg_1]; + int nc_1 = seq_chunk_start[seg_1 + 1] - c0_1; + int row0_1 = c0_1 * 64; + unsigned int _phase_kg_empty = 1; + unsigned int _phase_qd_empty = 1; + unsigned int _phase_w_empty = 1; + if (elect_sync()) { + #pragma unroll 1 + for (int rc_1 = 0; rc_1 < nc_1; rc_1++) { + int trow = row0_1 + (nc_1 - 1 - rc_1) * 64; + mbarrier_wait(kg_empty_addr, _phase_kg_empty); + mbarrier_arrive_expect_tx(kg_full_addr, 16384); + #pragma unroll + for (int panel = 0; panel < 2; panel++) { + tma_3d_gmem2smem(kg_panel_addr + (unsigned int)(panel * 8192), (&kg_tma), panel * 64, head_1, trow, kg_full_addr); + } + _phase_kg_empty ^= 1; + mbarrier_wait(qd_empty_addr, _phase_qd_empty); + mbarrier_arrive_expect_tx(qd_full_addr, 32768); + #pragma unroll + for (int panel_1 = 0; panel_1 < 2; panel_1++) { + tma_3d_gmem2smem(qg_panel_addr + (unsigned int)(panel_1 * 8192), (&qg_tma), panel_1 * 64, head_1, trow, qd_full_addr); + tma_3d_gmem2smem(do_panel_addr + (unsigned int)(panel_1 * 8192), (&do_tma), panel_1 * 64, head_1, trow, qd_full_addr); + } + _phase_qd_empty ^= 1; + mbarrier_wait(w_empty_addr, _phase_w_empty); + mbarrier_arrive_expect_tx(w_full_addr, 16384); + #pragma unroll + for (int panel_2 = 0; panel_2 < 2; panel_2++) { + tma_3d_gmem2smem(w_panel_addr + (unsigned int)(panel_2 * 8192), (&w_tma), panel_2 * 64, head_1, trow, w_full_addr); + } + _phase_w_empty ^= 1; + } + } + } + } + // ---- Role: mma ---- + if (warp == 5) { + { // mma_main + int head_2 = blockIdx.x % num_heads; + int walk_2 = blockIdx.x / num_heads; + int seg_2 = walk_seg[walk_2]; + int c0_2 = seq_chunk_start[seg_2]; + int nc_2 = seq_chunk_start[seg_2 + 1] - c0_2; + int row0_2 = c0_2 * 64; + unsigned int _phase_kg_full = 0; + unsigned int _phase_dhs_ready_0 = 0; + unsigned int _phase_dv_read_0 = 0; + unsigned int _phase_qd_full = 0; + unsigned int _phase_bdv_ready_0 = 0; + unsigned int _phase_w_full = 0; + if (elect_sync()) { + #pragma unroll 1 + for (int rc_2 = 0; rc_2 < nc_2; rc_2++) { + mbarrier_wait(kg_full_addr, _phase_kg_full); + mbarrier_wait(dhs_ready_addr, _phase_dhs_ready_0); + _phase_dhs_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = ((kg_op_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_1 = (((dhs_mn_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_0 = ((kg_op_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_1 = (((dhs_mn_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_acc_dv), "r"(0) : "memory"); + tcgen05_commit(dv_ready_addr); + tcgen05_commit(kg_empty_addr); + _phase_kg_full ^= 1; + mbarrier_wait(dv_read_addr, _phase_dv_read_0); + _phase_dv_read_0 ^= 1; + mbarrier_wait(qd_full_addr, _phase_qd_full); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = (((qg_mn_addr) >> 4) & 0x7FFF) | 0x2000000; + int _smem_desc_lo_3 = (((do_mn_addr) >> 4) & 0x7FFF) | 0x2000000; +#else + int _smem_desc_lo_2 = (((qg_mn_addr) >> 4) & 0x3FFF) | 0x2000000; + int _smem_desc_lo_3 = (((do_mn_addr) >> 4) & 0x3FFF) | 0x2000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136414352;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_acc_p), "r"(0) : "memory"); + tcgen05_commit(qd_empty_addr); + _phase_qd_full ^= 1; + mbarrier_wait(bdv_ready_addr, _phase_bdv_ready_0); + _phase_bdv_ready_0 ^= 1; + mbarrier_wait(w_full_addr, _phase_w_full); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_4 = (((w_mn_addr) >> 4) & 0x7FFF) | 0x2000000; + int _smem_desc_lo_5 = (((bdv_mn_addr) >> 4) & 0x7FFF) | 0x2000000; +#else + int _smem_desc_lo_4 = (((w_mn_addr) >> 4) & 0x3FFF) | 0x2000000; + int _smem_desc_lo_5 = (((bdv_mn_addr) >> 4) & 0x3FFF) | 0x2000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, %4;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_4), "r"(_smem_desc_lo_5), "r"(tmem_acc_dh), "r"(1), "r"((((uint32_t)(136414352)) | ((uint32_t)(1) << 13) | ((uint32_t)(0) << 14))) : "memory"); + tcgen05_commit(dh_ready_addr); + tcgen05_commit(w_empty_addr); + _phase_w_full ^= 1; + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu deleted file mode 100644 index f8ec15cd2bd..00000000000 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_kernel.cu +++ /dev/null @@ -1,636 +0,0 @@ -/* - * Copyright (c) 2026 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 256 -#define TMEM_ACC_H_OFFSET 0 -#define TMEM_ACC_VP_OFFSET 128 -#define NUM_RING_STAGES 2 -#define NUM_ONE_STAGES 1 -#define SMEM_W_PANEL_OFF 1024 -#define SMEM_W_PANEL_STAGE_BYTES 8192 -#define SMEM_W_PANEL_STRIDE 8192 -#define SMEM_KG_PANEL_OFF 33792 -#define SMEM_KG_PANEL_STAGE_BYTES 8192 -#define SMEM_KG_PANEL_STRIDE 8192 -#define SMEM_W_OP_OFF 1024 -#define SMEM_W_OP_STAGE_BYTES 16384 -#define SMEM_W_OP_STRIDE 16384 -#define SMEM_KG_MN_OFF 33792 -#define SMEM_KG_MN_STAGE_BYTES 16384 -#define SMEM_KG_MN_STRIDE 16384 -#define SMEM_HS_OFF 66560 -#define SMEM_HS_STAGE_BYTES 32768 -#define SMEM_HS_STRIDE 32768 -#define SMEM_HS_MN_OFF 66560 -#define SMEM_HS_MN_STAGE_BYTES 32768 -#define SMEM_HS_MN_STRIDE 32768 -#define SMEM_BV_OFF 99328 -#define SMEM_BV_STAGE_BYTES 16384 -#define SMEM_BV_STRIDE 16384 -#define SMEM_BV_MN_OFF 99328 -#define SMEM_BV_MN_STAGE_BYTES 16384 -#define SMEM_BV_MN_STRIDE 16384 -#define SMEM_TOTAL 115712 -#define THREADS 192 -#define LAUNCH_MIN_BLOCKS 1 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - - - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - - - -__device__ __forceinline__ void tma_3d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.3d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4}], [%5];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -extern "C" { - -__global__ __launch_bounds__(192, LAUNCH_MIN_BLOCKS) void -kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ h0, int* __restrict__ h0_index, int h0_hg, int h0_n, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, __nv_bfloat16* __restrict__ h_tape, __nv_bfloat16* __restrict__ v_new) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; -#if __CUDA_ARCH__ == 1000 - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); -#else - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); -#endif - - const int mbar_base = smem; - #define tiles_full_addr (mbar_base + 0) - #define tiles_empty_addr (mbar_base + 16) - #define hs_ready_addr (mbar_base + 32) - #define vp_ready_addr (mbar_base + 40) - #define bv_ready_addr (mbar_base + 48) - #define h_ready_addr (mbar_base + 56) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* w_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int w_panel_addr = smem + 1024; - __nv_bfloat16* kg_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); - const int kg_panel_addr = smem + 33792; - __nv_bfloat16* w_op = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int w_op_addr = smem + 1024; - __nv_bfloat16* kg_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); - const int kg_mn_addr = smem + 33792; - __nv_bfloat16* hs = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int hs_addr = smem + 66560; - __nv_bfloat16* hs_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int hs_mn_addr = smem + 66560; - __nv_bfloat16* bv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); - const int bv_addr = smem + 99328; - __nv_bfloat16* bv_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); - const int bv_mn_addr = smem + 99328; - - // Mbarrier init (6 pipeline groups, 0 ordered-sequence groups, 8 barriers) - // Mbarriers at smem_raw[0..64) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'ring' --- - // tiles_full: 2 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - mbarrier_init(smem + 8, 1); - // tiles_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - // --- pipeline 'one' --- - // hs_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 32, 128); - // vp_ready: 1 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - // bv_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 48, 128); - // h_ready: 1 barriers, init_count=1 - mbarrier_init(smem + 56, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (256 columns, 256 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 64); - if (warp == 0) { - int _tmem_hold = smem + 64; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_acc_h = taddr; - const int tmem_acc_vp = taddr + 128; - - // ---- Role: compute ---- - if (warp <= 3) { - { // compute_main - int head = blockIdx.x % num_heads; - int walk = blockIdx.x / num_heads; - int seg = walk_seg[walk]; - int c0 = seq_chunk_start[seg]; - int nc = seq_chunk_start[seg + 1] - c0; - int row0 = c0 * 64; - int warp_in_wg = warp % 4; - int k_row = warp_in_wg * 32 + lane; - int h_row_base = warp_in_wg * 32 << 16; - int lane_half = lane >> 4; - int t_row = warp_in_wg * 16 + (lane & 15); - int vp_row_base = warp_in_wg * 32 << 16; - float init_row[16]; - int h0_slot = h0_index[seg]; - if (h0_slot >= 0) { - int h0_tile = (head / h0_hg * h0_n + h0_slot) * h0_hg + head % h0_hg; - long long h0_base = (long long)h0_tile * (long long)tile_stride * 16384 + (long long)(k_row * 128); - #pragma unroll - for (int blk = 0; blk < 8; blk++) { - #pragma unroll - for (int q = 0; q < 4; q++) { - { - float4 _v4 = *reinterpret_cast(h0 + h0_base + (long long)(blk * 16 + q * 4)); - init_row[q * 4 + 0] = _v4.x; - init_row[q * 4 + 1] = _v4.y; - init_row[q * 4 + 2] = _v4.z; - init_row[q * 4 + 3] = _v4.w; - } - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + (unsigned int)(blk * 16), init_row); - } - } else { - init_row[0] = 0.0f; - init_row[1] = 0.0f; - init_row[2] = 0.0f; - init_row[3] = 0.0f; - init_row[4] = 0.0f; - init_row[5] = 0.0f; - init_row[6] = 0.0f; - init_row[7] = 0.0f; - init_row[8] = 0.0f; - init_row[9] = 0.0f; - init_row[10] = 0.0f; - init_row[11] = 0.0f; - init_row[12] = 0.0f; - init_row[13] = 0.0f; - init_row[14] = 0.0f; - init_row[15] = 0.0f; - #pragma unroll - for (int blk_1 = 0; blk_1 < 8; blk_1++) { - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + (unsigned int)(blk_1 * 16), init_row); - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - unsigned int _phase_vp_ready_0 = 0; - unsigned int _phase_h_ready_0 = 0; - #pragma unroll 1 - for (int c = 0; c < nc; c++) { - #pragma unroll - for (int blk_2 = 0; blk_2 < 8; blk_2++) { - float _tmem_load_0[16]; - tmem_ld_x16(&_tmem_load_0[0], taddr + (unsigned int)h_row_base + (unsigned int)(blk_2 * 16)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - unsigned int _tmem_load_0_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; - } - #pragma unroll - for (int q_1 = 0; q_1 < 2; q_1++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs_addr + (unsigned int)((blk_2 * 16 + q_1 * 8) / 64 * 16384 + k_row * 128 + (blk_2 * 16 + q_1 * 8) % 64 * 2 ^ ((blk_2 * 16 + q_1 * 8) / 64 * 16384 + k_row * 128 + (blk_2 * 16 + q_1 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[q_1 * 4]), "r"(_tmem_load_0_bf16[q_1 * 4 + 1]), "r"(_tmem_load_0_bf16[q_1 * 4 + 2]), "r"(_tmem_load_0_bf16[q_1 * 4 + 3]) : "memory"); - } - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(hs_ready_addr); - mbarrier_wait(vp_ready_addr, _phase_vp_ready_0); - _phase_vp_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int trow = row0 + c * 64 + t_row; - long long u_base = ((long long)trow * (long long)num_heads + (long long)head) * 128; - #pragma unroll - for (int blk_3 = 0; blk_3 < 4; blk_3++) { - float _tmem_load_1[16]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" - : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])) - : "r"(taddr + 128 + (unsigned int)vp_row_base + (unsigned int)(blk_3 * 32))); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - int col = blk_3 * 32 + lane_half * 16; - float uv[16]; - { - const uint4* _vptr_1 = reinterpret_cast(u + u_base + (long long)col); - uint4 _vld_1[2]; - #pragma unroll - for (int _blk = 0; _blk < 2; _blk++) { - _vld_1[_blk] = _vptr_1[_blk]; - uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); - #pragma unroll - for (int _pair = 0; _pair < 4; _pair++) { - (&uv[0 + _blk * 8 + _pair * 2])[0] = __uint_as_float(static_cast(_vpairs_1[_pair]) << 16); - (&uv[0 + _blk * 8 + _pair * 2])[1] = __uint_as_float(static_cast(_vpairs_1[_pair]) & 0xffff0000u); - } - } - } - float bvv[16]; - #pragma unroll - for (int j = 0; j < 16; j++) { - bvv[j] = uv[j] - _tmem_load_1[j]; - } - unsigned int bvv_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv[_lp*2 + 0], bvv[_lp*2+1 + 0])); - bvv_bf16[_lp] = *(uint32_t*)&_bf2; - } - #pragma unroll - for (int q_2 = 0; q_2 < 2; q_2++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv_addr + (unsigned int)((col + q_2 * 8) / 64 * 8192 + t_row * 128 + (col + q_2 * 8) % 64 * 2 ^ ((col + q_2 * 8) / 64 * 8192 + t_row * 128 + (col + q_2 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_bf16[q_2 * 4]), "r"(bvv_bf16[q_2 * 4 + 1]), "r"(bvv_bf16[q_2 * 4 + 2]), "r"(bvv_bf16[q_2 * 4 + 3]) : "memory"); - } - } - long long gn_index = ((long long)(row0 + c * 64 + 64 - 1) * (long long)num_heads + (long long)head) * 128 + (long long)k_row; - float decay = approx_exp2(gk[gn_index]); - #pragma unroll - for (int blk_4 = 0; blk_4 < 8; blk_4++) { - float _tmem_load_2[16]; - tmem_ld_x16(&_tmem_load_2[0], taddr + (unsigned int)h_row_base + (unsigned int)(blk_4 * 16)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - float hs2[16]; - #pragma unroll - for (int j_1 = 0; j_1 < 16; j_1++) { - hs2[j_1] = _tmem_load_2[j_1] * decay; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + (unsigned int)(blk_4 * 16), hs2); - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(bv_ready_addr); - mbarrier_wait(h_ready_addr, _phase_h_ready_0); - _phase_h_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int out_slot = out_index[seg]; - if (out_slot >= 0) { - int out_tile = (head / out_hg * out_n + out_slot) * out_hg + head % out_hg; - long long out_base = (long long)out_tile * (long long)tile_stride * 16384 + (long long)(k_row * 128); - #pragma unroll - for (int blk_5 = 0; blk_5 < 8; blk_5++) { - float _tmem_load_3[16]; - tmem_ld_x16(&_tmem_load_3[0], taddr + (unsigned int)h_row_base + (unsigned int)(blk_5 * 16)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int q_3 = 0; q_3 < 4; q_3++) { - { - float4 _v4 = make_float4(_tmem_load_3[q_3 * 4 + 0], _tmem_load_3[q_3 * 4 + 1], _tmem_load_3[q_3 * 4 + 2], _tmem_load_3[q_3 * 4 + 3]); - *reinterpret_cast(out + out_base + (long long)(blk_5 * 16 + q_3 * 4)) = _v4; - } - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - asm volatile("barrier.sync 1, 128;" ::: "memory"); - if (warp == 0) { - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); - } - } - } - // ---- Role: load ---- - if (warp == 4) { - { // load_main - int head_1 = blockIdx.x % num_heads; - int walk_1 = blockIdx.x / num_heads; - int seg_1 = walk_seg[walk_1]; - int c0_1 = seq_chunk_start[seg_1]; - int nc_1 = seq_chunk_start[seg_1 + 1] - c0_1; - int row0_1 = c0_1 * 64; - unsigned int stage = 0; - unsigned int _phase_tiles_empty = 1; - if (elect_sync()) { - #pragma unroll 1 - for (int c_1 = 0; c_1 < nc_1; c_1++) { - mbarrier_wait(tiles_empty_addr + (stage) * 8, _phase_tiles_empty); - mbarrier_arrive_expect_tx(tiles_full_addr + (stage) * 8, 32768); - int trow_1 = row0_1 + c_1 * 64; - #pragma unroll - for (int panel = 0; panel < 2; panel++) { - tma_3d_gmem2smem(w_panel_addr + (stage * 2 + (unsigned int)panel) * 8192, (&w_tma), panel * 64, head_1, trow_1, tiles_full_addr + (stage) * 8); - tma_3d_gmem2smem(kg_panel_addr + (stage * 2 + (unsigned int)panel) * 8192, (&kg_tma), panel * 64, head_1, trow_1, tiles_full_addr + (stage) * 8); - } - stage += 1; - if (stage == 2) { stage = 0; _phase_tiles_empty ^= 1; } - } - } - } - } - // ---- Role: mma ---- - if (warp == 5) { - { // mma_main - int head_2 = blockIdx.x % num_heads; - int walk_2 = blockIdx.x / num_heads; - int seg_2 = walk_seg[walk_2]; - int c0_2 = seq_chunk_start[seg_2]; - int nc_2 = seq_chunk_start[seg_2 + 1] - c0_2; - int row0_2 = c0_2 * 64; - unsigned int stage_1 = 0; - unsigned int _phase_tiles_full = 0; - unsigned int _phase_hs_ready_0 = 0; - unsigned int _phase_bv_ready_0 = 0; - if (elect_sync()) { - #pragma unroll 1 - for (int c_2 = 0; c_2 < nc_2; c_2++) { - mbarrier_wait(tiles_full_addr + (stage_1) * 8, _phase_tiles_full); - mbarrier_wait(hs_ready_addr, _phase_hs_ready_0); - _phase_hs_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x7FFF) + (stage_1) * 1024; - int _smem_desc_lo_1 = (((hs_mn_addr) >> 4) & 0x7FFF) | 0x4000000; -#else - int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x3FFF) + (stage_1) * 1024; - int _smem_desc_lo_1 = (((hs_mn_addr) >> 4) & 0x3FFF) | 0x4000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 506;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_acc_vp), "r"(0) : "memory"); - tcgen05_commit(vp_ready_addr); - mbarrier_wait(bv_ready_addr, _phase_bv_ready_0); - _phase_bv_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_2 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (stage_1) * 1024; - int _smem_desc_lo_3 = (((bv_mn_addr) >> 4) & 0x7FFF) | 0x2000000; -#else - int _smem_desc_lo_2 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (stage_1) * 1024; - int _smem_desc_lo_3 = (((bv_mn_addr) >> 4) & 0x3FFF) | 0x2000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136414352;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_acc_h), "r"(1) : "memory"); - tcgen05_commit(h_ready_addr); - tcgen05_commit(tiles_empty_addr + (stage_1) * 8); - stage_1 += 1; - if (stage_1 == 2) { stage_1 = 0; _phase_tiles_full ^= 1; } - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu similarity index 84% rename from flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu rename to flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu index 4c61e376eda..0bcc946d1d2 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7658fcc297c6fd906d53_kcp_binding.cu +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ h0, int* __restrict__ h0_index, int h0_hg, int h0_n, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, __nv_bfloat16* __restrict__ h_tape, __nv_bfloat16* __restrict__ v_new); +extern "C" __global__ void kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ h0, int* __restrict__ h0_index, int h0_hg, int h0_n, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, __nv_bfloat16* __restrict__ h_tape, __nv_bfloat16* __restrict__ v_new); -namespace cake_binding_kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp { +namespace cake_binding_kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp"; +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -96,8 +96,13 @@ CUtensorMap EncodeTma_w_tma(const TensorView& t) { cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "w_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "w_tma")}; cuuint32_t box_dim[3] = {64u, 1u, 64u}; cuuint32_t elem_strides[3] = {1u, 1u, 1u}; - TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) - << KERNEL_NAME << ": the box of descriptor 'w_tma' exceeds its source extents"; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'w_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) << KERNEL_NAME << ": the strides of descriptor 'w_tma' must be multiples of 16 bytes"; void* address = DataPtr(t); @@ -122,8 +127,13 @@ CUtensorMap EncodeTma_kg_tma(const TensorView& t) { cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "kg_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "kg_tma")}; cuuint32_t box_dim[3] = {64u, 1u, 64u}; cuuint32_t elem_strides[3] = {1u, 1u, 1u}; - TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) - << KERNEL_NAME << ": the box of descriptor 'kg_tma' exceeds its source extents"; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'kg_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) << KERNEL_NAME << ": the strides of descriptor 'kg_tma' must be multiples of 16 bytes"; void* address = DataPtr(t); @@ -225,7 +235,7 @@ void Run(TensorView a_w_tma, TensorView a_kg_tma, TensorView a_u, TensorView a_g const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 115712), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 115712), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -251,14 +261,14 @@ void Run(TensorView a_w_tma, TensorView a_kg_tma, TensorView a_u, TensorView a_g void* kernel_args[] = {&v_w_tma, &v_kg_tma, &v_u, &v_gk, &v_num_heads, &v_seq_chunk_start, &v_num_seqs, &v_walk_seg, &v_h0, &v_h0_index, &v_h0_hg, &v_h0_n, &v_out, &v_out_index, &v_out_hg, &v_out_n, &v_tile_stride, &v_h_tape, &v_v_new}; cudaLaunchConfig_t config{}; config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); - config.blockDim = dim3(192, 1, 1); + config.blockDim = dim3(320, 1, 1); config.dynamicSmemBytes = 115712; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp +} // namespace cake_binding_kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_7658fcc297c6fd906d53_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu new file mode 100644 index 00000000000..1d6a4e81b43 --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_7d43005ed0adb9ffadf2_kcp_kernel.cu @@ -0,0 +1,926 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LAUNCH_MIN_BLOCKS 2 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +extern "C" { + +__global__ __launch_bounds__(320, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_det_train_7d43005ed0adb9ffadf2_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ h0, int* __restrict__ h0_index, int h0_hg, int h0_n, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, __nv_bfloat16* __restrict__ h_tape, __nv_bfloat16* __restrict__ v_new) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define tiles_full_addr (mbar_base + 0) + #define tiles_empty_addr (mbar_base + 16) + #define hs_ready_addr (mbar_base + 32) + #define vp_ready_addr (mbar_base + 40) + #define bv_ready_addr (mbar_base + 48) + #define dec_done_addr (mbar_base + 56) + #define h_ready_addr (mbar_base + 64) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Prologue + __nv_bfloat16* w_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int w_panel_addr = smem + 1024; + __nv_bfloat16* kg_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int kg_panel_addr = smem + 33792; + __nv_bfloat16* w_op = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int w_op_addr = smem + 1024; + __nv_bfloat16* kg_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int kg_mn_addr = smem + 33792; + __nv_bfloat16* hs = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int hs_addr = smem + 66560; + __nv_bfloat16* hs_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int hs_mn_addr = smem + 66560; + __nv_bfloat16* bv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int bv_addr = smem + 99328; + __nv_bfloat16* bv_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int bv_mn_addr = smem + 99328; + + // Mbarrier init (9 barriers at smem_raw[0..72)) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'ring' --- + // tiles_full: 2 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + mbarrier_init(smem + 8, 1); + // tiles_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + // --- pipeline 'one' --- + // hs_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 32, 128); + // vp_ready: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // bv_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 48, 128); + // dec_done: 1 barriers, init_count=128 + mbarrier_init(smem + 56, 128); + // h_ready: 1 barriers, init_count=2 + mbarrier_init(smem + 64, 2); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 72); + if (warp == 0) { + int _tmem_hold = smem + 72; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const unsigned int taddr = tmem_addr_storage[0]; + const unsigned int tmem_acc_h = taddr; + const unsigned int tmem_acc_vp = taddr + 128; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int head = blockIdx.x % num_heads; + int walk = blockIdx.x / num_heads; + int seg = walk_seg[walk]; + int c0 = seq_chunk_start[seg]; + int nc = seq_chunk_start[seg + 1] - c0; + int row0 = c0 * 64; + int warp_in_wg = warp % 4; + int k_row = warp_in_wg * 32 + lane; + int h_row_base = warp_in_wg * 32 << 16; + int lane_half = lane >> 4; + int t_row = warp_in_wg * 16 + (lane & 15); + int vp_row_base = warp_in_wg * 32 << 16; + float init_row[16]; + int h0_slot = h0_index[seg]; + if (h0_slot >= 0) { + int h0_tile = (head / h0_hg * h0_n + h0_slot) * h0_hg + head % h0_hg; + long long h0_base = (long long)h0_tile * (long long)tile_stride * 16384 + (long long)(k_row * 128); + #pragma unroll + for (int blk = 0; blk < 8; blk++) { + #pragma unroll + for (int q = 0; q < 4; q++) { + { + float4 _v4 = *reinterpret_cast(h0 + h0_base + (long long)(blk * 16 + q * 4)); + init_row[q * 4 + 0] = _v4.x; + init_row[q * 4 + 1] = _v4.y; + init_row[q * 4 + 2] = _v4.z; + init_row[q * 4 + 3] = _v4.w; + } + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base + (unsigned int)(blk * 16), init_row); + } + } else { + init_row[0] = 0.0f; + init_row[1] = 0.0f; + init_row[2] = 0.0f; + init_row[3] = 0.0f; + init_row[4] = 0.0f; + init_row[5] = 0.0f; + init_row[6] = 0.0f; + init_row[7] = 0.0f; + init_row[8] = 0.0f; + init_row[9] = 0.0f; + init_row[10] = 0.0f; + init_row[11] = 0.0f; + init_row[12] = 0.0f; + init_row[13] = 0.0f; + init_row[14] = 0.0f; + init_row[15] = 0.0f; + #pragma unroll + for (int blk_1 = 0; blk_1 < 8; blk_1++) { + tmem_st_x16_f32(taddr + (unsigned int)h_row_base + (unsigned int)(blk_1 * 16), init_row); + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + unsigned int u_pk[32]; + u_pk[0] = 0; + u_pk[1] = 0; + u_pk[2] = 0; + u_pk[3] = 0; + u_pk[4] = 0; + u_pk[5] = 0; + u_pk[6] = 0; + u_pk[7] = 0; + u_pk[8] = 0; + u_pk[9] = 0; + u_pk[10] = 0; + u_pk[11] = 0; + u_pk[12] = 0; + u_pk[13] = 0; + u_pk[14] = 0; + u_pk[15] = 0; + u_pk[16] = 0; + u_pk[17] = 0; + u_pk[18] = 0; + u_pk[19] = 0; + u_pk[20] = 0; + u_pk[21] = 0; + u_pk[22] = 0; + u_pk[23] = 0; + u_pk[24] = 0; + u_pk[25] = 0; + u_pk[26] = 0; + u_pk[27] = 0; + u_pk[28] = 0; + u_pk[29] = 0; + u_pk[30] = 0; + u_pk[31] = 0; + if (nc > 0) { + { + uint4 _uv4_1 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(lane_half * 16)); + u_pk[0 + 0] = _uv4_1.x; + u_pk[0 + 1] = _uv4_1.y; + u_pk[0 + 2] = _uv4_1.z; + u_pk[0 + 3] = _uv4_1.w; + } + { + uint4 _uv4_2 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(lane_half * 16 + 8)); + u_pk[4 + 0] = _uv4_2.x; + u_pk[4 + 1] = _uv4_2.y; + u_pk[4 + 2] = _uv4_2.z; + u_pk[4 + 3] = _uv4_2.w; + } + { + uint4 _uv4_3 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(32 + lane_half * 16)); + u_pk[8 + 0] = _uv4_3.x; + u_pk[8 + 1] = _uv4_3.y; + u_pk[8 + 2] = _uv4_3.z; + u_pk[8 + 3] = _uv4_3.w; + } + { + uint4 _uv4_4 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(32 + lane_half * 16 + 8)); + u_pk[12 + 0] = _uv4_4.x; + u_pk[12 + 1] = _uv4_4.y; + u_pk[12 + 2] = _uv4_4.z; + u_pk[12 + 3] = _uv4_4.w; + } + { + uint4 _uv4_5 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(64 + lane_half * 16)); + u_pk[16 + 0] = _uv4_5.x; + u_pk[16 + 1] = _uv4_5.y; + u_pk[16 + 2] = _uv4_5.z; + u_pk[16 + 3] = _uv4_5.w; + } + { + uint4 _uv4_6 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(64 + lane_half * 16 + 8)); + u_pk[20 + 0] = _uv4_6.x; + u_pk[20 + 1] = _uv4_6.y; + u_pk[20 + 2] = _uv4_6.z; + u_pk[20 + 3] = _uv4_6.w; + } + { + uint4 _uv4_7 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(96 + lane_half * 16)); + u_pk[24 + 0] = _uv4_7.x; + u_pk[24 + 1] = _uv4_7.y; + u_pk[24 + 2] = _uv4_7.z; + u_pk[24 + 3] = _uv4_7.w; + } + { + uint4 _uv4_8 = *reinterpret_cast(u + ((long long)(row0 + t_row) * (long long)num_heads + (long long)head) * 128 + (long long)(96 + lane_half * 16 + 8)); + u_pk[28 + 0] = _uv4_8.x; + u_pk[28 + 1] = _uv4_8.y; + u_pk[28 + 2] = _uv4_8.z; + u_pk[28 + 3] = _uv4_8.w; + } + } + unsigned int _phase_vp_ready_0 = 0; + unsigned int _phase_h_ready_0 = 0; + #pragma unroll 1 + for (int c = 0; c < nc; c++) { + #pragma unroll + for (int blk_2 = 0; blk_2 < 8; blk_2++) { + float h_blk[16]; + tmem_ld_x16(&h_blk[0], taddr + (unsigned int)h_row_base + (unsigned int)(blk_2 * 16)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + unsigned int h_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(h_blk[_lp*2 + 0], h_blk[_lp*2+1 + 0])); + h_bf16[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_1 = 0; q_1 < 2; q_1++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs_addr + (unsigned int)((blk_2 * 16 + q_1 * 8) / 64 * 16384 + k_row * 128 + (blk_2 * 16 + q_1 * 8) % 64 * 2 ^ ((blk_2 * 16 + q_1 * 8) / 64 * 16384 + k_row * 128 + (blk_2 * 16 + q_1 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(h_bf16[q_1 * 4]), "r"(h_bf16[q_1 * 4 + 1]), "r"(h_bf16[q_1 * 4 + 2]), "r"(h_bf16[q_1 * 4 + 3]) : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(hs_ready_addr); + mbarrier_wait(vp_ready_addr, _phase_vp_ready_0); + _phase_vp_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + long long u_base = ((long long)(row0 + c * 64 + t_row) * (long long)num_heads + (long long)head) * 128; + float vp_blk_0[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&vp_blk_0[0])), "=r"(*reinterpret_cast(&vp_blk_0[1])), "=r"(*reinterpret_cast(&vp_blk_0[2])), "=r"(*reinterpret_cast(&vp_blk_0[3])), "=r"(*reinterpret_cast(&vp_blk_0[4])), "=r"(*reinterpret_cast(&vp_blk_0[5])), "=r"(*reinterpret_cast(&vp_blk_0[6])), "=r"(*reinterpret_cast(&vp_blk_0[7])), "=r"(*reinterpret_cast(&vp_blk_0[8])), "=r"(*reinterpret_cast(&vp_blk_0[9])), "=r"(*reinterpret_cast(&vp_blk_0[10])), "=r"(*reinterpret_cast(&vp_blk_0[11])), "=r"(*reinterpret_cast(&vp_blk_0[12])), "=r"(*reinterpret_cast(&vp_blk_0[13])), "=r"(*reinterpret_cast(&vp_blk_0[14])), "=r"(*reinterpret_cast(&vp_blk_0[15])) + : "r"(taddr + 128 + (unsigned int)vp_row_base + 0)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int col_0 = lane_half * 16; + float u_blk_0[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&u_blk_0[_pair * 2])[0] = __uint_as_float(static_cast(u_pk[_pair]) << 16); + (&u_blk_0[_pair * 2])[1] = __uint_as_float(static_cast(u_pk[_pair]) & 0xffff0000u); + } + float bvv_0[16]; + #pragma unroll + for (int j_8 = 0; j_8 < 16; j_8++) { + bvv_0[j_8] = u_blk_0[j_8] - vp_blk_0[j_8]; + } + unsigned int bvv_bf16_0[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv_0[_lp*2 + 0], bvv_0[_lp*2+1 + 0])); + bvv_bf16_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_0 = 0; q_0 < 2; q_0++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv_addr + (unsigned int)((col_0 + q_0 * 8) / 64 * 8192 + t_row * 128 + (col_0 + q_0 * 8) % 64 * 2 ^ ((col_0 + q_0 * 8) / 64 * 8192 + t_row * 128 + (col_0 + q_0 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_bf16_0[q_0 * 4]), "r"(bvv_bf16_0[q_0 * 4 + 1]), "r"(bvv_bf16_0[q_0 * 4 + 2]), "r"(bvv_bf16_0[q_0 * 4 + 3]) : "memory"); + } + float vp_blk_1[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&vp_blk_1[0])), "=r"(*reinterpret_cast(&vp_blk_1[1])), "=r"(*reinterpret_cast(&vp_blk_1[2])), "=r"(*reinterpret_cast(&vp_blk_1[3])), "=r"(*reinterpret_cast(&vp_blk_1[4])), "=r"(*reinterpret_cast(&vp_blk_1[5])), "=r"(*reinterpret_cast(&vp_blk_1[6])), "=r"(*reinterpret_cast(&vp_blk_1[7])), "=r"(*reinterpret_cast(&vp_blk_1[8])), "=r"(*reinterpret_cast(&vp_blk_1[9])), "=r"(*reinterpret_cast(&vp_blk_1[10])), "=r"(*reinterpret_cast(&vp_blk_1[11])), "=r"(*reinterpret_cast(&vp_blk_1[12])), "=r"(*reinterpret_cast(&vp_blk_1[13])), "=r"(*reinterpret_cast(&vp_blk_1[14])), "=r"(*reinterpret_cast(&vp_blk_1[15])) + : "r"(taddr + 128 + (unsigned int)vp_row_base + 32)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int col_1 = 32 + lane_half * 16; + float u_blk_1[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&u_blk_1[_pair * 2])[0] = __uint_as_float(static_cast(u_pk[8 + _pair]) << 16); + (&u_blk_1[_pair * 2])[1] = __uint_as_float(static_cast(u_pk[8 + _pair]) & 0xffff0000u); + } + float bvv_1[16]; + #pragma unroll + for (int j_9 = 0; j_9 < 16; j_9++) { + bvv_1[j_9] = u_blk_1[j_9] - vp_blk_1[j_9]; + } + unsigned int bvv_bf16_1[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv_1[_lp*2 + 0], bvv_1[_lp*2+1 + 0])); + bvv_bf16_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_1_1 = 0; q_1_1 < 2; q_1_1++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv_addr + (unsigned int)((col_1 + q_1_1 * 8) / 64 * 8192 + t_row * 128 + (col_1 + q_1_1 * 8) % 64 * 2 ^ ((col_1 + q_1_1 * 8) / 64 * 8192 + t_row * 128 + (col_1 + q_1_1 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_bf16_1[q_1_1 * 4]), "r"(bvv_bf16_1[q_1_1 * 4 + 1]), "r"(bvv_bf16_1[q_1_1 * 4 + 2]), "r"(bvv_bf16_1[q_1_1 * 4 + 3]) : "memory"); + } + float vp_blk_2[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&vp_blk_2[0])), "=r"(*reinterpret_cast(&vp_blk_2[1])), "=r"(*reinterpret_cast(&vp_blk_2[2])), "=r"(*reinterpret_cast(&vp_blk_2[3])), "=r"(*reinterpret_cast(&vp_blk_2[4])), "=r"(*reinterpret_cast(&vp_blk_2[5])), "=r"(*reinterpret_cast(&vp_blk_2[6])), "=r"(*reinterpret_cast(&vp_blk_2[7])), "=r"(*reinterpret_cast(&vp_blk_2[8])), "=r"(*reinterpret_cast(&vp_blk_2[9])), "=r"(*reinterpret_cast(&vp_blk_2[10])), "=r"(*reinterpret_cast(&vp_blk_2[11])), "=r"(*reinterpret_cast(&vp_blk_2[12])), "=r"(*reinterpret_cast(&vp_blk_2[13])), "=r"(*reinterpret_cast(&vp_blk_2[14])), "=r"(*reinterpret_cast(&vp_blk_2[15])) + : "r"(taddr + 128 + (unsigned int)vp_row_base + 64)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int col_2 = 64 + lane_half * 16; + float u_blk_2[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&u_blk_2[_pair * 2])[0] = __uint_as_float(static_cast(u_pk[16 + _pair]) << 16); + (&u_blk_2[_pair * 2])[1] = __uint_as_float(static_cast(u_pk[16 + _pair]) & 0xffff0000u); + } + float bvv_2[16]; + #pragma unroll + for (int j_10 = 0; j_10 < 16; j_10++) { + bvv_2[j_10] = u_blk_2[j_10] - vp_blk_2[j_10]; + } + unsigned int bvv_bf16_2[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv_2[_lp*2 + 0], bvv_2[_lp*2+1 + 0])); + bvv_bf16_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_2 = 0; q_2 < 2; q_2++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv_addr + (unsigned int)((col_2 + q_2 * 8) / 64 * 8192 + t_row * 128 + (col_2 + q_2 * 8) % 64 * 2 ^ ((col_2 + q_2 * 8) / 64 * 8192 + t_row * 128 + (col_2 + q_2 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_bf16_2[q_2 * 4]), "r"(bvv_bf16_2[q_2 * 4 + 1]), "r"(bvv_bf16_2[q_2 * 4 + 2]), "r"(bvv_bf16_2[q_2 * 4 + 3]) : "memory"); + } + float vp_blk_3[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&vp_blk_3[0])), "=r"(*reinterpret_cast(&vp_blk_3[1])), "=r"(*reinterpret_cast(&vp_blk_3[2])), "=r"(*reinterpret_cast(&vp_blk_3[3])), "=r"(*reinterpret_cast(&vp_blk_3[4])), "=r"(*reinterpret_cast(&vp_blk_3[5])), "=r"(*reinterpret_cast(&vp_blk_3[6])), "=r"(*reinterpret_cast(&vp_blk_3[7])), "=r"(*reinterpret_cast(&vp_blk_3[8])), "=r"(*reinterpret_cast(&vp_blk_3[9])), "=r"(*reinterpret_cast(&vp_blk_3[10])), "=r"(*reinterpret_cast(&vp_blk_3[11])), "=r"(*reinterpret_cast(&vp_blk_3[12])), "=r"(*reinterpret_cast(&vp_blk_3[13])), "=r"(*reinterpret_cast(&vp_blk_3[14])), "=r"(*reinterpret_cast(&vp_blk_3[15])) + : "r"(taddr + 128 + (unsigned int)vp_row_base + 96)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int col_3 = 96 + lane_half * 16; + float u_blk_3[16]; + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + (&u_blk_3[_pair * 2])[0] = __uint_as_float(static_cast(u_pk[24 + _pair]) << 16); + (&u_blk_3[_pair * 2])[1] = __uint_as_float(static_cast(u_pk[24 + _pair]) & 0xffff0000u); + } + float bvv_3[16]; + #pragma unroll + for (int j_11 = 0; j_11 < 16; j_11++) { + bvv_3[j_11] = u_blk_3[j_11] - vp_blk_3[j_11]; + } + unsigned int bvv_bf16_3[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv_3[_lp*2 + 0], bvv_3[_lp*2+1 + 0])); + bvv_bf16_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int q_3 = 0; q_3 < 2; q_3++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv_addr + (unsigned int)((col_3 + q_3 * 8) / 64 * 8192 + t_row * 128 + (col_3 + q_3 * 8) % 64 * 2 ^ ((col_3 + q_3 * 8) / 64 * 8192 + t_row * 128 + (col_3 + q_3 * 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_bf16_3[q_3 * 4]), "r"(bvv_bf16_3[q_3 * 4 + 1]), "r"(bvv_bf16_3[q_3 * 4 + 2]), "r"(bvv_bf16_3[q_3 * 4 + 3]) : "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(bv_ready_addr); + if (nc > c + 1) { + long long u_next = u_base + (long long)(64 * num_heads * 128); + { + uint4 _uv4_9 = *reinterpret_cast(u + u_next + (long long)(lane_half * 16)); + u_pk[0 + 0] = _uv4_9.x; + u_pk[0 + 1] = _uv4_9.y; + u_pk[0 + 2] = _uv4_9.z; + u_pk[0 + 3] = _uv4_9.w; + } + { + uint4 _uv4_10 = *reinterpret_cast(u + u_next + (long long)(lane_half * 16 + 8)); + u_pk[4 + 0] = _uv4_10.x; + u_pk[4 + 1] = _uv4_10.y; + u_pk[4 + 2] = _uv4_10.z; + u_pk[4 + 3] = _uv4_10.w; + } + { + uint4 _uv4_11 = *reinterpret_cast(u + u_next + (long long)(32 + lane_half * 16)); + u_pk[8 + 0] = _uv4_11.x; + u_pk[8 + 1] = _uv4_11.y; + u_pk[8 + 2] = _uv4_11.z; + u_pk[8 + 3] = _uv4_11.w; + } + { + uint4 _uv4_12 = *reinterpret_cast(u + u_next + (long long)(32 + lane_half * 16 + 8)); + u_pk[12 + 0] = _uv4_12.x; + u_pk[12 + 1] = _uv4_12.y; + u_pk[12 + 2] = _uv4_12.z; + u_pk[12 + 3] = _uv4_12.w; + } + { + uint4 _uv4_13 = *reinterpret_cast(u + u_next + (long long)(64 + lane_half * 16)); + u_pk[16 + 0] = _uv4_13.x; + u_pk[16 + 1] = _uv4_13.y; + u_pk[16 + 2] = _uv4_13.z; + u_pk[16 + 3] = _uv4_13.w; + } + { + uint4 _uv4_14 = *reinterpret_cast(u + u_next + (long long)(64 + lane_half * 16 + 8)); + u_pk[20 + 0] = _uv4_14.x; + u_pk[20 + 1] = _uv4_14.y; + u_pk[20 + 2] = _uv4_14.z; + u_pk[20 + 3] = _uv4_14.w; + } + { + uint4 _uv4_15 = *reinterpret_cast(u + u_next + (long long)(96 + lane_half * 16)); + u_pk[24 + 0] = _uv4_15.x; + u_pk[24 + 1] = _uv4_15.y; + u_pk[24 + 2] = _uv4_15.z; + u_pk[24 + 3] = _uv4_15.w; + } + { + uint4 _uv4_16 = *reinterpret_cast(u + u_next + (long long)(96 + lane_half * 16 + 8)); + u_pk[28 + 0] = _uv4_16.x; + u_pk[28 + 1] = _uv4_16.y; + u_pk[28 + 2] = _uv4_16.z; + u_pk[28 + 3] = _uv4_16.w; + } + } + mbarrier_wait(h_ready_addr, _phase_h_ready_0); + _phase_h_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int out_slot = out_index[seg]; + if (out_slot >= 0) { + int out_tile = (head / out_hg * out_n + out_slot) * out_hg + head % out_hg; + long long out_base = (long long)out_tile * (long long)tile_stride * 16384 + (long long)(k_row * 128); + #pragma unroll + for (int blk_3 = 0; blk_3 < 8; blk_3++) { + float h_out[16]; + tmem_ld_x16(&h_out[0], taddr + (unsigned int)h_row_base + (unsigned int)(blk_3 * 16)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int q_4 = 0; q_4 < 4; q_4++) { + { + float4 _v4 = make_float4(h_out[q_4 * 4 + 0], h_out[q_4 * 4 + 1], h_out[q_4 * 4 + 2], h_out[q_4 * 4 + 3]); + *reinterpret_cast(out + out_base + (long long)(blk_3 * 16 + q_4 * 4)) = _v4; + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("barrier.sync 1, 256;" ::: "memory"); + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + } + } + } + // ---- Role: decayw ---- + if (warp >= 4 && warp <= 7) { + { // decayw_main + int head_1 = blockIdx.x % num_heads; + int walk_1 = blockIdx.x / num_heads; + int seg_1 = walk_seg[walk_1]; + int c0_1 = seq_chunk_start[seg_1]; + int nc_1 = seq_chunk_start[seg_1 + 1] - c0_1; + int row0_1 = c0_1 * 64; + int warp_in_wg_1 = warp % 4; + int k_row_1 = warp_in_wg_1 * 32 + lane; + int h_row_base_1 = warp_in_wg_1 * 32 << 16; + float gk_pre = 0.0f; + if (nc_1 > 0) { + long long gn_first = ((long long)(row0_1 + 64 - 1) * (long long)num_heads + (long long)head_1) * 128 + (long long)k_row_1; + gk_pre = gk[gn_first]; + } + unsigned int _phase_hs_ready_0 = 0; + #pragma unroll 1 + for (int c_1 = 0; c_1 < nc_1; c_1++) { + mbarrier_wait(hs_ready_addr, _phase_hs_ready_0); + _phase_hs_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float decay = approx_exp2(gk_pre); + float h_blk_0[16]; + tmem_ld_x16(&h_blk_0[0], taddr + (unsigned int)h_row_base_1 + 0); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_0[16]; + #pragma unroll + for (int j_0 = 0; j_0 < 16; j_0++) { + hd_0[j_0] = h_blk_0[j_0] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 0, hd_0); + float h_blk_1[16]; + tmem_ld_x16(&h_blk_1[0], taddr + (unsigned int)h_row_base_1 + 16); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_1[16]; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + hd_1[j_1] = h_blk_1[j_1] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 16, hd_1); + float h_blk_2[16]; + tmem_ld_x16(&h_blk_2[0], taddr + (unsigned int)h_row_base_1 + 32); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_2[16]; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + hd_2[j_2] = h_blk_2[j_2] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 32, hd_2); + float h_blk_3[16]; + tmem_ld_x16(&h_blk_3[0], taddr + (unsigned int)h_row_base_1 + 48); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_3[16]; + #pragma unroll + for (int j_3 = 0; j_3 < 16; j_3++) { + hd_3[j_3] = h_blk_3[j_3] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 48, hd_3); + float h_blk_4[16]; + tmem_ld_x16(&h_blk_4[0], taddr + (unsigned int)h_row_base_1 + 64); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_4[16]; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + hd_4[j_4] = h_blk_4[j_4] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 64, hd_4); + float h_blk_5[16]; + tmem_ld_x16(&h_blk_5[0], taddr + (unsigned int)h_row_base_1 + 80); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_5[16]; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + hd_5[j_5] = h_blk_5[j_5] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 80, hd_5); + float h_blk_6[16]; + tmem_ld_x16(&h_blk_6[0], taddr + (unsigned int)h_row_base_1 + 96); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_6[16]; + #pragma unroll + for (int j_6 = 0; j_6 < 16; j_6++) { + hd_6[j_6] = h_blk_6[j_6] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 96, hd_6); + float h_blk_7[16]; + tmem_ld_x16(&h_blk_7[0], taddr + (unsigned int)h_row_base_1 + 112); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float hd_7[16]; + #pragma unroll + for (int j_7 = 0; j_7 < 16; j_7++) { + hd_7[j_7] = h_blk_7[j_7] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)h_row_base_1 + 112, hd_7); + if (nc_1 > c_1 + 1) { + long long gn_next = ((long long)(row0_1 + c_1 * 64 + 128 - 1) * (long long)num_heads + (long long)head_1) * 128 + (long long)k_row_1; + gk_pre = gk[gn_next]; + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(dec_done_addr); + } + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("barrier.sync 1, 256;" ::: "memory"); + } + } + // ---- Role: load ---- + if (warp == 8) { + { // load_main + int head_2 = blockIdx.x % num_heads; + int walk_2 = blockIdx.x / num_heads; + int seg_2 = walk_seg[walk_2]; + int c0_2 = seq_chunk_start[seg_2]; + int nc_2 = seq_chunk_start[seg_2 + 1] - c0_2; + int row0_2 = c0_2 * 64; + unsigned int stage = 0; + unsigned int _phase_tiles_empty = 1; + if (elect_sync()) { + #pragma unroll 1 + for (int c_2 = 0; c_2 < nc_2; c_2++) { + mbarrier_wait(tiles_empty_addr + (stage) * 8, _phase_tiles_empty); + mbarrier_arrive_expect_tx(tiles_full_addr + (stage) * 8, 32768); + int trow = row0_2 + c_2 * 64; + #pragma unroll + for (int panel = 0; panel < 2; panel++) { + tma_3d_gmem2smem(w_panel_addr + (stage * 2 + (unsigned int)panel) * 8192, (&w_tma), panel * 64, head_2, trow, tiles_full_addr + (stage) * 8); + tma_3d_gmem2smem(kg_panel_addr + (stage * 2 + (unsigned int)panel) * 8192, (&kg_tma), panel * 64, head_2, trow, tiles_full_addr + (stage) * 8); + } + stage += 1; + if (stage == 2) { stage = 0; _phase_tiles_empty ^= 1; } + } + } + } + } + // ---- Role: mma ---- + if (warp == 9) { + { // mma_main + int head_3 = blockIdx.x % num_heads; + int walk_3 = blockIdx.x / num_heads; + int seg_3 = walk_seg[walk_3]; + int c0_3 = seq_chunk_start[seg_3]; + int nc_3 = seq_chunk_start[seg_3 + 1] - c0_3; + int row0_3 = c0_3 * 64; + unsigned int stage_1 = 0; + unsigned int _phase_tiles_full = 0; + unsigned int _phase_hs_ready_0_1 = 0; + unsigned int _phase_bv_ready_0 = 0; + unsigned int _phase_dec_done_0 = 0; + if (elect_sync()) { + #pragma unroll 1 + for (int c_3 = 0; c_3 < nc_3; c_3++) { + mbarrier_wait(tiles_full_addr + (stage_1) * 8, _phase_tiles_full); + mbarrier_wait(hs_ready_addr, _phase_hs_ready_0_1); + _phase_hs_ready_0_1 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x7FFF) + (stage_1) * 1024; + int _smem_desc_lo_1 = (((hs_mn_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x3FFF) + (stage_1) * 1024; + int _smem_desc_lo_1 = (((hs_mn_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_acc_vp), "r"(0) : "memory"); + tcgen05_commit(vp_ready_addr); + mbarrier_wait(bv_ready_addr, _phase_bv_ready_0); + _phase_bv_ready_0 ^= 1; + mbarrier_wait(dec_done_addr, _phase_dec_done_0); + _phase_dec_done_0 ^= 1; + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(h_ready_addr); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (stage_1) * 1024; + int _smem_desc_lo_3 = (((bv_mn_addr) >> 4) & 0x7FFF) | 0x2000000; +#else + int _smem_desc_lo_2 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (stage_1) * 1024; + int _smem_desc_lo_3 = (((bv_mn_addr) >> 4) & 0x3FFF) | 0x2000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136414352;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_acc_h), "r"(1) : "memory"); + tcgen05_commit(h_ready_addr); + tcgen05_commit(tiles_empty_addr + (stage_1) * 8); + stage_1 += 1; + if (stage_1 == 2) { stage_1 = 0; _phase_tiles_full ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" From 9f3182225f2dcad6d763afde708febbfa1c58e0a Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 18:10:18 -0800 Subject: [PATCH 08/10] cake_kda_det_train_kcp: the shipped transfer walk takes the two-warpgroup BF16x3 build The shipped half-width BF16x3 transfer walk record points at the build with two compute warpgroups (320 threads, same bits). The test-only transfer walk records stay as they are. --- .../experimental/cake_kda_det_train_kcp/cake_kcp_jit.py | 8 ++++---- tests/experimental/test_cake_kda_det_train_kcp_build.py | 9 +++++---- tests/experimental/test_cake_kda_train_kcp_build.py | 6 ++++-- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index 223e636eb1f..69330255b22 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -785,15 +785,15 @@ "ffi_entry": "run", # Two CTAs per (segment, head): each walks one 64-column half of the state. "grid_rule": {"ctas_per_head": 2, "kind": "kcp_segments_heads"}, - "kernel_symbol": "kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp", + "kernel_symbol": "kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp", "launch": { - "block": (192, 1, 1), + "block": (320, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 107520, + "dynamic_smem_bytes": 0, "use_pdl": False, }, - "library": "cake_kda_det_train_9a3f306bacf8c8a1547d_kcp", + "library": "cake_kda_det_train_8ffd931ca0595948b608_kcp", "logical": { "gk": "gk", "init_mode": "init_mode", diff --git a/tests/experimental/test_cake_kda_det_train_kcp_build.py b/tests/experimental/test_cake_kda_det_train_kcp_build.py index ef8858188fa..5457674f7d3 100644 --- a/tests/experimental/test_cake_kda_det_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_det_train_kcp_build.py @@ -136,8 +136,9 @@ "kcp_dhu_tape_full": DHU_PLAN, } # Launch block of the records that are not 192 threads wide (the merge records: MERGE_RECORDS); the full-width fwdh -# offset walk runs the split-role schedule (four state, four decay, one TMA and one MMA warp). -BLOCKS = {"kcp_fwdh_offset": 320} +# offset walk runs the split-role schedule (four state, four decay, one TMA and one MMA warp), the shipped transfer +# walk two compute warpgroups (eight compute, one load and one MMA warp). +BLOCKS = {"kcp_fwdh_offset": 320, "kcp_transfer_walk": 320} # The identity rule: ``cake_kda_det_train_<20 hex of sha256("\\n")>``, # then ``_kcp`` for a program only the context-parallel route runs (every one but the single-rank ``dav`` stage). # The digest hashes the source module path the program shipped under, so the digests are those of the previous names @@ -173,7 +174,7 @@ {"program": "tape", "vh": 64, "tape_store": "tma"}, ), "kcp_transfer_walk": ( - f"tart.kda_chunk_kcp_transfer_walk.{_W}.kda_chunk_kcp_transfer_walk_bf16x3_v1", + f"tart.kda_chunk_kcp_transfer_walk.{_W}.kda_chunk_det_transfer_walk_bf16x3_v2_kcp", {"vh": 64}, ), **{name: (MERGE_V2_MODULE, row[0]) for name, row in MERGE_RECORDS.items()}, @@ -517,7 +518,7 @@ def test_walk_heads_selection() -> None: def test_the_shipped_transfer_walk_is_the_half_width_bf16x3_program() -> None: record = kcp_jit.MODULES["kcp_transfer_walk"] assert record["grid_rule"] == {"ctas_per_head": 2, "kind": "kcp_segments_heads"} - assert record["library"] == "cake_kda_det_train_9a3f306bacf8c8a1547d_kcp" + assert record["library"] == "cake_kda_det_train_8ffd931ca0595948b608_kcp" previous = kcp_jit.TEST_MODULES["kcp_transfer_walk_tf32_hilo"] assert previous["grid_rule"] == {"kind": "kcp_segments_heads"} assert previous["library"] == "cake_kda_det_train_1ac972f0cb7cc7b13cf7_kcp" diff --git a/tests/experimental/test_cake_kda_train_kcp_build.py b/tests/experimental/test_cake_kda_train_kcp_build.py index 7bb08a96760..3a86bc35beb 100644 --- a/tests/experimental/test_cake_kda_train_kcp_build.py +++ b/tests/experimental/test_cake_kda_train_kcp_build.py @@ -491,13 +491,15 @@ def test_test_manifest_is_the_test_only_programs_under_the_opt_in( DET_PREFIX, PREFIX = "cake_kda_det_train_", "cake_kda_train_" # The records whose builds differ between the two packages: the deterministic package ships newer builds of the -# affine merge (selected by its MERGE_RULE) and of the full-width offset walks (the adjoint one beside the plain build, -# selected by its WALK_HEADS_RULE); this package keeps its own builds of these records. The deterministic package's +# affine merge (selected by its MERGE_RULE), of the full-width offset walks (the adjoint one beside the plain build, +# selected by its WALK_HEADS_RULE) and of the shipped transfer walk; this package keeps its own builds of these +# records. The deterministic package's # extra records belong to those programs. Every other record and unit is the deterministic package's under this # package's names. OWN = { "kcp_fwdh_offset": PREFIX + "7658fcc297c6fd906d53_kcp", "kcp_dhu_offset": PREFIX + "1c66573e87150245374b_kcp", + "kcp_transfer_walk": PREFIX + "9a3f306bacf8c8a1547d_kcp", "kcp_affine_merge_c8": PREFIX + "c6958a0a4ed5cb55a9a1_kcp", "kcp_affine_merge_c16": PREFIX + "89f5ff59f03eb170e3bf_kcp", } From 4599e4d7b7eafc37597fab5a4252b940dc1a5714 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 18:16:17 -0800 Subject: [PATCH 09/10] cake_kda_det_train_kcp: render the two-warpgroup BF16x3 transfer walk build --- .../cake_kda_det_train_kcp/build.json | 42 +- .../cake_kda_det_train_kcp/cake_kcp_jit.py | 2 +- ...train_8ffd931ca0595948b608_kcp_binding.cu} | 36 +- ...t_train_8ffd931ca0595948b608_kcp_kernel.cu | 1094 +++++++++++++ ...t_train_9a3f306bacf8c8a1547d_kcp_kernel.cu | 1452 ----------------- 5 files changed, 1139 insertions(+), 1487 deletions(-) rename flashinfer/experimental/cake_kda_det_train_kcp/csrc/{cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu => cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu} (83%) create mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu delete mode 100644 flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/build.json b/flashinfer/experimental/cake_kda_det_train_kcp/build.json index ac691d1a85a..d2578f41a30 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/build.json +++ b/flashinfer/experimental/cake_kda_det_train_kcp/build.json @@ -545,65 +545,65 @@ }, { "arch": "sm_100a", - "closure_sha256": "35fb18a6795791394cabdf8faea9e544ab7c5722683afa1fddd95176ed5b45cc", + "closure_sha256": "6c2aa715acc6d6f5f5077dda6c1c57e369d2f9f18e4f30d9f6a30d67ec86d4a0", "entry": "run", "gencode": "-gencode=arch=compute_100a,code=sm_100a", "min_cuda": "12.8", - "name": "cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_sm_100a", + "name": "cake_kda_det_train_8ffd931ca0595948b608_kcp_sm_100a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_sm_100a.so", + "path": "aot/cake_kda_det_train_8ffd931ca0595948b608_kcp_sm_100a.so", "sha256": { - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu": "005ae78cfc5e19d1c3e37f42c761ff8d3c70bd9d8d815d6fa39b7b3b992b32a6", - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu": "1dee65abf64a4ad4587e58340ad370909e6775c557d6ffa7527e20b21099dd7f" + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu": "86eeacf9a351d290f18d36d48df3fc3e9e0e4a9fb85443dcea23e9b8d0937d39", + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu": "4ad74ee862d8af722b5cf78fbf9865a2d68c3b59c58daceaab43b15d7928addd" }, "sources": [ - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu", - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu" + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu", + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu" ] }, { "arch": "sm_103a", - "closure_sha256": "39832b33709f6143abcc02e064e5191376acaa60050fb365608855fc26b7c3bd", + "closure_sha256": "550a2892705ea2d0e7e731bddaa5a5a87f41d14f18c477c0a0c795aff1e1afa2", "entry": "run", "gencode": "-gencode=arch=compute_103a,code=sm_103a", "min_cuda": "12.9", - "name": "cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_sm_103a", + "name": "cake_kda_det_train_8ffd931ca0595948b608_kcp_sm_103a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_sm_103a.so", + "path": "aot/cake_kda_det_train_8ffd931ca0595948b608_kcp_sm_103a.so", "sha256": { - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu": "005ae78cfc5e19d1c3e37f42c761ff8d3c70bd9d8d815d6fa39b7b3b992b32a6", - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu": "1dee65abf64a4ad4587e58340ad370909e6775c557d6ffa7527e20b21099dd7f" + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu": "86eeacf9a351d290f18d36d48df3fc3e9e0e4a9fb85443dcea23e9b8d0937d39", + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu": "4ad74ee862d8af722b5cf78fbf9865a2d68c3b59c58daceaab43b15d7928addd" }, "sources": [ - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu", - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu" + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu", + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu" ] }, { "arch": "sm_107a", - "closure_sha256": "e9964d549ef099c4351c4ea7afac73593bfa87cf295857f63c6cf0b516e9db14", + "closure_sha256": "09723d6ddd2deeefee819e8da66174f13f0957f2d76cdac42ce21fedd9579143", "entry": "run", "gencode": "-gencode=arch=compute_107a,code=sm_107a", "min_cuda": "13.4", - "name": "cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_sm_107a", + "name": "cake_kda_det_train_8ffd931ca0595948b608_kcp_sm_107a", "nvcc_flags": [ "--use_fast_math" ], "nvcc_flags_if": [], - "path": "aot/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_sm_107a.so", + "path": "aot/cake_kda_det_train_8ffd931ca0595948b608_kcp_sm_107a.so", "sha256": { - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu": "005ae78cfc5e19d1c3e37f42c761ff8d3c70bd9d8d815d6fa39b7b3b992b32a6", - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu": "1dee65abf64a4ad4587e58340ad370909e6775c557d6ffa7527e20b21099dd7f" + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu": "86eeacf9a351d290f18d36d48df3fc3e9e0e4a9fb85443dcea23e9b8d0937d39", + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu": "4ad74ee862d8af722b5cf78fbf9865a2d68c3b59c58daceaab43b15d7928addd" }, "sources": [ - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu", - "csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu" + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu", + "csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu" ] }, { diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py index 69330255b22..40bf34b33b1 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py @@ -790,7 +790,7 @@ "block": (320, 1, 1), "cluster": (1, 1, 1), "cooperative": False, - "dynamic_smem_bytes": 0, + "dynamic_smem_bytes": 107520, "use_pdl": False, }, "library": "cake_kda_det_train_8ffd931ca0595948b608_kcp", diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu similarity index 83% rename from flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu rename to flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu index d442ffb0a10..1d23dc27662 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_binding.cu +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, int init_mode, int use_u, int write_transposed); +extern "C" __global__ void kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, int init_mode, int use_u, int write_transposed); -namespace cake_binding_kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp { +namespace cake_binding_kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp"; +constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -96,8 +96,13 @@ CUtensorMap EncodeTma_w_tma(const TensorView& t) { cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "w_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "w_tma")}; cuuint32_t box_dim[3] = {64u, 1u, 64u}; cuuint32_t elem_strides[3] = {1u, 1u, 1u}; - TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) - << KERNEL_NAME << ": the box of descriptor 'w_tma' exceeds its source extents"; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'w_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) << KERNEL_NAME << ": the strides of descriptor 'w_tma' must be multiples of 16 bytes"; void* address = DataPtr(t); @@ -122,8 +127,13 @@ CUtensorMap EncodeTma_kg_tma(const TensorView& t) { cuuint64_t global_strides[2] = {Stride(CheckedMul(d1, 2LL), "kg_tma"), Stride(CheckedMul(CheckedMul(d1, d2), 2LL), "kg_tma")}; cuuint32_t box_dim[3] = {64u, 1u, 64u}; cuuint32_t elem_strides[3] = {1u, 1u, 1u}; - TVM_FFI_CHECK(box_dim[0] <= global_dim[0] && box_dim[1] <= global_dim[1] && box_dim[2] <= global_dim[2], ValueError) - << KERNEL_NAME << ": the box of descriptor 'kg_tma' exceeds its source extents"; + const bool allow_oob_box[3] = {false, false, false}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'kg_tma' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) << KERNEL_NAME << ": the strides of descriptor 'kg_tma' must be multiples of 16 bytes"; void* address = DataPtr(t); @@ -206,7 +216,7 @@ void Run(TensorView a_w_tma, TensorView a_kg_tma, TensorView a_u, TensorView a_g const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 107520), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 107520), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -229,14 +239,14 @@ void Run(TensorView a_w_tma, TensorView a_kg_tma, TensorView a_u, TensorView a_g void* kernel_args[] = {&v_w_tma, &v_kg_tma, &v_u, &v_gk, &v_num_heads, &v_seq_chunk_start, &v_num_seqs, &v_walk_seg, &v_out, &v_out_index, &v_out_hg, &v_out_n, &v_tile_stride, &v_init_mode, &v_use_u, &v_write_transposed}; cudaLaunchConfig_t config{}; config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); - config.blockDim = dim3(192, 1, 1); + config.blockDim = dim3(320, 1, 1); config.dynamicSmemBytes = 107520; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp +} // namespace cake_binding_kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu new file mode 100644 index 00000000000..8014ee96dd8 --- /dev/null +++ b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_8ffd931ca0595948b608_kcp_kernel.cu @@ -0,0 +1,1094 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LAUNCH_MIN_BLOCKS 2 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +extern "C" { + +__global__ __launch_bounds__(320, LAUNCH_MIN_BLOCKS) void +kernel_cake_kda_det_train_8ffd931ca0595948b608_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, int init_mode, int use_u, int write_transposed) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define w_full_addr (mbar_base + 0) + #define w_empty_addr (mbar_base + 8) + #define kg_full_addr (mbar_base + 16) + #define kg_empty_addr (mbar_base + 24) + #define hs_ready_addr (mbar_base + 32) + #define vp_ready_addr (mbar_base + 40) + #define bv_ready_addr (mbar_base + 48) + #define h_ready_addr (mbar_base + 56) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Prologue + __nv_bfloat16* w_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int w_panel_addr = smem + 1024; + __nv_bfloat16* kg_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int kg_panel_addr = smem + 17408; + __nv_bfloat16* w_op = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int w_op_addr = smem + 1024; + __nv_bfloat16* kg_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int kg_mn_addr = smem + 17408; + __nv_bfloat16* hs0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int hs0_addr = smem + 33792; + __nv_bfloat16* hs1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int hs1_addr = smem + 50176; + __nv_bfloat16* hs2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int hs2_addr = smem + 66560; + __nv_bfloat16* hs_mn0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int hs_mn0_addr = smem + 33792; + __nv_bfloat16* hs_mn1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int hs_mn1_addr = smem + 50176; + __nv_bfloat16* hs_mn2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int hs_mn2_addr = smem + 66560; + __nv_bfloat16* bv0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 82944); + const int bv0_addr = smem + 82944; + __nv_bfloat16* bv1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 91136); + const int bv1_addr = smem + 91136; + __nv_bfloat16* bv2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int bv2_addr = smem + 99328; + __nv_bfloat16* bv_mn0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 82944); + const int bv_mn0_addr = smem + 82944; + __nv_bfloat16* bv_mn1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 91136); + const int bv_mn1_addr = smem + 91136; + __nv_bfloat16* bv_mn2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int bv_mn2_addr = smem + 99328; + + // Mbarrier init (8 barriers at smem_raw[0..64)) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'ring_pipe' --- + // w_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // w_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // kg_full: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + // kg_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 24, 1); + // --- pipeline 'one' --- + // hs_ready: 1 barriers, init_count=256 + mbarrier_init(smem + 32, 256); + // vp_ready: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // bv_ready: 1 barriers, init_count=256 + mbarrier_init(smem + 48, 256); + // h_ready: 1 barriers, init_count=1 + mbarrier_init(smem + 56, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (256 columns, 192 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 64); + if (warp == 0) { + int _tmem_hold = smem + 64; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const unsigned int taddr = tmem_addr_storage[0]; + const unsigned int tmem_acc_h = taddr; + const unsigned int tmem_acc_vp = taddr + 64; + const unsigned int tmem_acc_dx = taddr + 128; + + // ---- Role: compute ---- + if (warp <= 7) { + { // compute_main + int vslice = blockIdx.x % 2; + int head = blockIdx.x / 2 % num_heads; + int seg = walk_seg[blockIdx.x / (2 * num_heads)]; + int c0 = seq_chunk_start[seg]; + int nc = seq_chunk_start[seg + 1] - c0; + int row0 = c0 * 64; + int warp_in_wg = warp % 4; + int k_row = warp_in_wg * 32 + lane; + int h_row_base = warp_in_wg * 32 << 16; + int lane_half = lane >> 4; + int t_row = warp_in_wg * 16 + (lane & 15); + int vp_row_base = warp_in_wg * 32 << 16; + int v0 = vslice * 64; + unsigned int warp_group_idx = ((unsigned int)tid - 0) / 128; + int cb = (int)warp_group_idx * 32; + unsigned int mask_hi = 4294901760u; + float init_row[16]; + #pragma unroll + for (int j_1_0 = 0; j_1_0 < 16; j_1_0++) { + float one_here_1_0 = 0.0f; + if (init_mode == 0) { + if (v0 + cb + j_1_0 == k_row) { + one_here_1_0 = 1.0f; + } + } + init_row[j_1_0] = one_here_1_0; + } + tmem_st_x16_f32(taddr + (unsigned int)(h_row_base + cb), init_row); + #pragma unroll + for (int j_1_1 = 0; j_1_1 < 16; j_1_1++) { + float one_here_1_1 = 0.0f; + if (init_mode == 0) { + if (v0 + cb + 16 + j_1_1 == k_row) { + one_here_1_1 = 1.0f; + } + } + init_row[j_1_1] = one_here_1_1; + } + tmem_st_x16_f32(taddr + (unsigned int)(h_row_base + cb) + 16, init_row); + float zeros[16]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + zeros[8] = 0.0f; + zeros[9] = 0.0f; + zeros[10] = 0.0f; + zeros[11] = 0.0f; + zeros[12] = 0.0f; + zeros[13] = 0.0f; + zeros[14] = 0.0f; + zeros[15] = 0.0f; + tmem_st_x16_f32(taddr + 128 + (unsigned int)(h_row_base + cb), zeros); + tmem_st_x16_f32(taddr + 128 + (unsigned int)(h_row_base + cb) + 16, zeros); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + float h_a[16]; + float h_b[16]; + float d_a[16]; + float d_b[16]; + float vp_a[16]; + float vp_b[16]; + float gk_cur = 0.0f; + if (nc > 0) { + gk_cur = gk[((long long)(row0 + 64 - 1) * (long long)num_heads + (long long)head) * 128 + (long long)k_row]; + } + unsigned int _phase_vp_ready_0 = 0; + unsigned int _phase_h_ready_0 = 0; + #pragma unroll 1 + for (int c = 0; c < nc; c++) { + float decay = approx_exp2(gk_cur); + if (nc > c + 1) { + gk_cur = gk[((long long)(row0 + (c + 1) * 64 + 64 - 1) * (long long)num_heads + (long long)head) * 128 + (long long)k_row]; + } + tmem_ld_x16(&h_a[0], taddr + (unsigned int)(h_row_base + cb)); + tmem_ld_x16(&d_a[0], taddr + 128 + (unsigned int)(h_row_base + cb)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + tmem_ld_x16(&h_b[0], taddr + (unsigned int)(h_row_base + cb) + 16); + tmem_ld_x16(&d_b[0], taddr + 128 + (unsigned int)(h_row_base + cb) + 16); + #pragma unroll + for (int j = 0; j < 16; j++) { + h_a[j] = h_a[j] + d_a[j]; + } + unsigned int tmem_load_0_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(h_a[_lp*2 + 0], h_a[_lp*2+1 + 0])); + tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + float r1[16]; + r1[0] = h_a[0] - __uint_as_float(tmem_load_0_bf16[0] << 16); + r1[1] = h_a[1] - __uint_as_float(tmem_load_0_bf16[0] & mask_hi); + r1[2] = h_a[2] - __uint_as_float(tmem_load_0_bf16[1] << 16); + r1[3] = h_a[3] - __uint_as_float(tmem_load_0_bf16[1] & mask_hi); + r1[4] = h_a[4] - __uint_as_float(tmem_load_0_bf16[2] << 16); + r1[5] = h_a[5] - __uint_as_float(tmem_load_0_bf16[2] & mask_hi); + r1[6] = h_a[6] - __uint_as_float(tmem_load_0_bf16[3] << 16); + r1[7] = h_a[7] - __uint_as_float(tmem_load_0_bf16[3] & mask_hi); + r1[8] = h_a[8] - __uint_as_float(tmem_load_0_bf16[4] << 16); + r1[9] = h_a[9] - __uint_as_float(tmem_load_0_bf16[4] & mask_hi); + r1[10] = h_a[10] - __uint_as_float(tmem_load_0_bf16[5] << 16); + r1[11] = h_a[11] - __uint_as_float(tmem_load_0_bf16[5] & mask_hi); + r1[12] = h_a[12] - __uint_as_float(tmem_load_0_bf16[6] << 16); + r1[13] = h_a[13] - __uint_as_float(tmem_load_0_bf16[6] & mask_hi); + r1[14] = h_a[14] - __uint_as_float(tmem_load_0_bf16[7] << 16); + r1[15] = h_a[15] - __uint_as_float(tmem_load_0_bf16[7] & mask_hi); + unsigned int r1_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1[_lp*2 + 0], r1[_lp*2+1 + 0])); + r1_bf16[_lp] = *(uint32_t*)&_bf2; + } + float r2[16]; + r2[0] = r1[0] - __uint_as_float(r1_bf16[0] << 16); + r2[1] = r1[1] - __uint_as_float(r1_bf16[0] & mask_hi); + r2[2] = r1[2] - __uint_as_float(r1_bf16[1] << 16); + r2[3] = r1[3] - __uint_as_float(r1_bf16[1] & mask_hi); + r2[4] = r1[4] - __uint_as_float(r1_bf16[2] << 16); + r2[5] = r1[5] - __uint_as_float(r1_bf16[2] & mask_hi); + r2[6] = r1[6] - __uint_as_float(r1_bf16[3] << 16); + r2[7] = r1[7] - __uint_as_float(r1_bf16[3] & mask_hi); + r2[8] = r1[8] - __uint_as_float(r1_bf16[4] << 16); + r2[9] = r1[9] - __uint_as_float(r1_bf16[4] & mask_hi); + r2[10] = r1[10] - __uint_as_float(r1_bf16[5] << 16); + r2[11] = r1[11] - __uint_as_float(r1_bf16[5] & mask_hi); + r2[12] = r1[12] - __uint_as_float(r1_bf16[6] << 16); + r2[13] = r1[13] - __uint_as_float(r1_bf16[6] & mask_hi); + r2[14] = r1[14] - __uint_as_float(r1_bf16[7] << 16); + r2[15] = r1[15] - __uint_as_float(r1_bf16[7] & mask_hi); + unsigned int r2_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2[_lp*2 + 0], r2[_lp*2+1 + 0])); + r2_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(cb / 64 * 16384 + k_row * 128 + cb % 64 * 2 ^ (cb / 64 * 16384 + k_row * 128 + cb % 64 * 2 >> 7 & 7) << 4))), "r"(tmem_load_0_bf16[0]), "r"(tmem_load_0_bf16[1]), "r"(tmem_load_0_bf16[2]), "r"(tmem_load_0_bf16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)((cb + 8) / 64 * 16384 + k_row * 128 + (cb + 8) % 64 * 2 ^ ((cb + 8) / 64 * 16384 + k_row * 128 + (cb + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(tmem_load_0_bf16[4]), "r"(tmem_load_0_bf16[5]), "r"(tmem_load_0_bf16[6]), "r"(tmem_load_0_bf16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(cb / 64 * 16384 + k_row * 128 + cb % 64 * 2 ^ (cb / 64 * 16384 + k_row * 128 + cb % 64 * 2 >> 7 & 7) << 4))), "r"(r1_bf16[0]), "r"(r1_bf16[1]), "r"(r1_bf16[2]), "r"(r1_bf16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)((cb + 8) / 64 * 16384 + k_row * 128 + (cb + 8) % 64 * 2 ^ ((cb + 8) / 64 * 16384 + k_row * 128 + (cb + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r1_bf16[4]), "r"(r1_bf16[5]), "r"(r1_bf16[6]), "r"(r1_bf16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(cb / 64 * 16384 + k_row * 128 + cb % 64 * 2 ^ (cb / 64 * 16384 + k_row * 128 + cb % 64 * 2 >> 7 & 7) << 4))), "r"(r2_bf16[0]), "r"(r2_bf16[1]), "r"(r2_bf16[2]), "r"(r2_bf16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)((cb + 8) / 64 * 16384 + k_row * 128 + (cb + 8) % 64 * 2 ^ ((cb + 8) / 64 * 16384 + k_row * 128 + (cb + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r2_bf16[4]), "r"(r2_bf16[5]), "r"(r2_bf16[6]), "r"(r2_bf16[7]) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + h_b[j_1] = h_b[j_1] + d_b[j_1]; + } + unsigned int tmem_load_0_bf16_0[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(h_b[_lp*2 + 0], h_b[_lp*2+1 + 0])); + tmem_load_0_bf16_0[_lp] = *(uint32_t*)&_bf2; + } + float r1_0[16]; + r1_0[0] = h_b[0] - __uint_as_float(tmem_load_0_bf16_0[0] << 16); + r1_0[1] = h_b[1] - __uint_as_float(tmem_load_0_bf16_0[0] & mask_hi); + r1_0[2] = h_b[2] - __uint_as_float(tmem_load_0_bf16_0[1] << 16); + r1_0[3] = h_b[3] - __uint_as_float(tmem_load_0_bf16_0[1] & mask_hi); + r1_0[4] = h_b[4] - __uint_as_float(tmem_load_0_bf16_0[2] << 16); + r1_0[5] = h_b[5] - __uint_as_float(tmem_load_0_bf16_0[2] & mask_hi); + r1_0[6] = h_b[6] - __uint_as_float(tmem_load_0_bf16_0[3] << 16); + r1_0[7] = h_b[7] - __uint_as_float(tmem_load_0_bf16_0[3] & mask_hi); + r1_0[8] = h_b[8] - __uint_as_float(tmem_load_0_bf16_0[4] << 16); + r1_0[9] = h_b[9] - __uint_as_float(tmem_load_0_bf16_0[4] & mask_hi); + r1_0[10] = h_b[10] - __uint_as_float(tmem_load_0_bf16_0[5] << 16); + r1_0[11] = h_b[11] - __uint_as_float(tmem_load_0_bf16_0[5] & mask_hi); + r1_0[12] = h_b[12] - __uint_as_float(tmem_load_0_bf16_0[6] << 16); + r1_0[13] = h_b[13] - __uint_as_float(tmem_load_0_bf16_0[6] & mask_hi); + r1_0[14] = h_b[14] - __uint_as_float(tmem_load_0_bf16_0[7] << 16); + r1_0[15] = h_b[15] - __uint_as_float(tmem_load_0_bf16_0[7] & mask_hi); + unsigned int r1_bf16_0[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1_0[_lp*2 + 0], r1_0[_lp*2+1 + 0])); + r1_bf16_0[_lp] = *(uint32_t*)&_bf2; + } + float r2_0[16]; + r2_0[0] = r1_0[0] - __uint_as_float(r1_bf16_0[0] << 16); + r2_0[1] = r1_0[1] - __uint_as_float(r1_bf16_0[0] & mask_hi); + r2_0[2] = r1_0[2] - __uint_as_float(r1_bf16_0[1] << 16); + r2_0[3] = r1_0[3] - __uint_as_float(r1_bf16_0[1] & mask_hi); + r2_0[4] = r1_0[4] - __uint_as_float(r1_bf16_0[2] << 16); + r2_0[5] = r1_0[5] - __uint_as_float(r1_bf16_0[2] & mask_hi); + r2_0[6] = r1_0[6] - __uint_as_float(r1_bf16_0[3] << 16); + r2_0[7] = r1_0[7] - __uint_as_float(r1_bf16_0[3] & mask_hi); + r2_0[8] = r1_0[8] - __uint_as_float(r1_bf16_0[4] << 16); + r2_0[9] = r1_0[9] - __uint_as_float(r1_bf16_0[4] & mask_hi); + r2_0[10] = r1_0[10] - __uint_as_float(r1_bf16_0[5] << 16); + r2_0[11] = r1_0[11] - __uint_as_float(r1_bf16_0[5] & mask_hi); + r2_0[12] = r1_0[12] - __uint_as_float(r1_bf16_0[6] << 16); + r2_0[13] = r1_0[13] - __uint_as_float(r1_bf16_0[6] & mask_hi); + r2_0[14] = r1_0[14] - __uint_as_float(r1_bf16_0[7] << 16); + r2_0[15] = r1_0[15] - __uint_as_float(r1_bf16_0[7] & mask_hi); + unsigned int r2_bf16_0[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2_0[_lp*2 + 0], r2_0[_lp*2+1 + 0])); + r2_bf16_0[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)((cb + 16) / 64 * 16384 + k_row * 128 + (cb + 16) % 64 * 2 ^ ((cb + 16) / 64 * 16384 + k_row * 128 + (cb + 16) % 64 * 2 >> 7 & 7) << 4))), "r"(tmem_load_0_bf16_0[0]), "r"(tmem_load_0_bf16_0[1]), "r"(tmem_load_0_bf16_0[2]), "r"(tmem_load_0_bf16_0[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)((cb + 16 + 8) / 64 * 16384 + k_row * 128 + (cb + 16 + 8) % 64 * 2 ^ ((cb + 16 + 8) / 64 * 16384 + k_row * 128 + (cb + 16 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(tmem_load_0_bf16_0[4]), "r"(tmem_load_0_bf16_0[5]), "r"(tmem_load_0_bf16_0[6]), "r"(tmem_load_0_bf16_0[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)((cb + 16) / 64 * 16384 + k_row * 128 + (cb + 16) % 64 * 2 ^ ((cb + 16) / 64 * 16384 + k_row * 128 + (cb + 16) % 64 * 2 >> 7 & 7) << 4))), "r"(r1_bf16_0[0]), "r"(r1_bf16_0[1]), "r"(r1_bf16_0[2]), "r"(r1_bf16_0[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)((cb + 16 + 8) / 64 * 16384 + k_row * 128 + (cb + 16 + 8) % 64 * 2 ^ ((cb + 16 + 8) / 64 * 16384 + k_row * 128 + (cb + 16 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r1_bf16_0[4]), "r"(r1_bf16_0[5]), "r"(r1_bf16_0[6]), "r"(r1_bf16_0[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)((cb + 16) / 64 * 16384 + k_row * 128 + (cb + 16) % 64 * 2 ^ ((cb + 16) / 64 * 16384 + k_row * 128 + (cb + 16) % 64 * 2 >> 7 & 7) << 4))), "r"(r2_bf16_0[0]), "r"(r2_bf16_0[1]), "r"(r2_bf16_0[2]), "r"(r2_bf16_0[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)((cb + 16 + 8) / 64 * 16384 + k_row * 128 + (cb + 16 + 8) % 64 * 2 ^ ((cb + 16 + 8) / 64 * 16384 + k_row * 128 + (cb + 16 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r2_bf16_0[4]), "r"(r2_bf16_0[5]), "r"(r2_bf16_0[6]), "r"(r2_bf16_0[7]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(hs_ready_addr); + tmem_ld_x16(&h_a[0], taddr + (unsigned int)(h_row_base + cb)); + tmem_ld_x16(&d_a[0], taddr + 128 + (unsigned int)(h_row_base + cb)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + tmem_ld_x16(&h_b[0], taddr + (unsigned int)(h_row_base + cb) + 16); + tmem_ld_x16(&d_b[0], taddr + 128 + (unsigned int)(h_row_base + cb) + 16); + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + h_a[j_2] = h_a[j_2] + d_a[j_2]; + } + #pragma unroll + for (int j_3 = 0; j_3 < 16; j_3++) { + h_a[j_3] = h_a[j_3] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)(h_row_base + cb), h_a); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + h_b[j_4] = h_b[j_4] + d_b[j_4]; + } + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + h_b[j_5] = h_b[j_5] * decay; + } + tmem_st_x16_f32(taddr + (unsigned int)(h_row_base + cb) + 16, h_b); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_wait(vp_ready_addr, _phase_vp_ready_0); + _phase_vp_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int trow = row0 + c * 64 + t_row; + long long u_base = ((long long)trow * (long long)num_heads + (long long)head) * 128 + (long long)v0; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" + : "=r"(*reinterpret_cast(&vp_a[0])), "=r"(*reinterpret_cast(&vp_a[1])), "=r"(*reinterpret_cast(&vp_a[2])), "=r"(*reinterpret_cast(&vp_a[3])), "=r"(*reinterpret_cast(&vp_a[4])), "=r"(*reinterpret_cast(&vp_a[5])), "=r"(*reinterpret_cast(&vp_a[6])), "=r"(*reinterpret_cast(&vp_a[7])), "=r"(*reinterpret_cast(&vp_a[8])), "=r"(*reinterpret_cast(&vp_a[9])), "=r"(*reinterpret_cast(&vp_a[10])), "=r"(*reinterpret_cast(&vp_a[11])), "=r"(*reinterpret_cast(&vp_a[12])), "=r"(*reinterpret_cast(&vp_a[13])), "=r"(*reinterpret_cast(&vp_a[14])), "=r"(*reinterpret_cast(&vp_a[15])) + : "r"(taddr + 64 + (unsigned int)(vp_row_base + cb))); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int col_1 = cb + lane_half * 16; + float uv_1[16]; + uv_1[0] = 0.0f; + uv_1[1] = 0.0f; + uv_1[2] = 0.0f; + uv_1[3] = 0.0f; + uv_1[4] = 0.0f; + uv_1[5] = 0.0f; + uv_1[6] = 0.0f; + uv_1[7] = 0.0f; + uv_1[8] = 0.0f; + uv_1[9] = 0.0f; + uv_1[10] = 0.0f; + uv_1[11] = 0.0f; + uv_1[12] = 0.0f; + uv_1[13] = 0.0f; + uv_1[14] = 0.0f; + uv_1[15] = 0.0f; + if (use_u == 1) { + { + const uint4* _vptr_0 = reinterpret_cast(u + u_base + (long long)col_1); + uint4 _vld_0[2]; + #pragma unroll + for (int _blk = 0; _blk < 2; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + (&uv_1[0 + _blk * 8 + _pair * 2])[0] = __uint_as_float(static_cast(_vpairs_0[_pair]) << 16); + (&uv_1[0 + _blk * 8 + _pair * 2])[1] = __uint_as_float(static_cast(_vpairs_0[_pair]) & 0xffff0000u); + } + } + } + } + float bvv_1[16]; + #pragma unroll + for (int j_1_2 = 0; j_1_2 < 16; j_1_2++) { + bvv_1[j_1_2] = uv_1[j_1_2] - vp_a[j_1_2]; + } + unsigned int tmem_load_0_bf16_1[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv_1[_lp*2 + 0], bvv_1[_lp*2+1 + 0])); + tmem_load_0_bf16_1[_lp] = *(uint32_t*)&_bf2; + } + float r1_1[16]; + r1_1[0] = bvv_1[0] - __uint_as_float(tmem_load_0_bf16_1[0] << 16); + r1_1[1] = bvv_1[1] - __uint_as_float(tmem_load_0_bf16_1[0] & mask_hi); + r1_1[2] = bvv_1[2] - __uint_as_float(tmem_load_0_bf16_1[1] << 16); + r1_1[3] = bvv_1[3] - __uint_as_float(tmem_load_0_bf16_1[1] & mask_hi); + r1_1[4] = bvv_1[4] - __uint_as_float(tmem_load_0_bf16_1[2] << 16); + r1_1[5] = bvv_1[5] - __uint_as_float(tmem_load_0_bf16_1[2] & mask_hi); + r1_1[6] = bvv_1[6] - __uint_as_float(tmem_load_0_bf16_1[3] << 16); + r1_1[7] = bvv_1[7] - __uint_as_float(tmem_load_0_bf16_1[3] & mask_hi); + r1_1[8] = bvv_1[8] - __uint_as_float(tmem_load_0_bf16_1[4] << 16); + r1_1[9] = bvv_1[9] - __uint_as_float(tmem_load_0_bf16_1[4] & mask_hi); + r1_1[10] = bvv_1[10] - __uint_as_float(tmem_load_0_bf16_1[5] << 16); + r1_1[11] = bvv_1[11] - __uint_as_float(tmem_load_0_bf16_1[5] & mask_hi); + r1_1[12] = bvv_1[12] - __uint_as_float(tmem_load_0_bf16_1[6] << 16); + r1_1[13] = bvv_1[13] - __uint_as_float(tmem_load_0_bf16_1[6] & mask_hi); + r1_1[14] = bvv_1[14] - __uint_as_float(tmem_load_0_bf16_1[7] << 16); + r1_1[15] = bvv_1[15] - __uint_as_float(tmem_load_0_bf16_1[7] & mask_hi); + unsigned int r1_bf16_2[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1_1[_lp*2 + 0], r1_1[_lp*2+1 + 0])); + r1_bf16_2[_lp] = *(uint32_t*)&_bf2; + } + float r2_1[16]; + r2_1[0] = r1_1[0] - __uint_as_float(r1_bf16_2[0] << 16); + r2_1[1] = r1_1[1] - __uint_as_float(r1_bf16_2[0] & mask_hi); + r2_1[2] = r1_1[2] - __uint_as_float(r1_bf16_2[1] << 16); + r2_1[3] = r1_1[3] - __uint_as_float(r1_bf16_2[1] & mask_hi); + r2_1[4] = r1_1[4] - __uint_as_float(r1_bf16_2[2] << 16); + r2_1[5] = r1_1[5] - __uint_as_float(r1_bf16_2[2] & mask_hi); + r2_1[6] = r1_1[6] - __uint_as_float(r1_bf16_2[3] << 16); + r2_1[7] = r1_1[7] - __uint_as_float(r1_bf16_2[3] & mask_hi); + r2_1[8] = r1_1[8] - __uint_as_float(r1_bf16_2[4] << 16); + r2_1[9] = r1_1[9] - __uint_as_float(r1_bf16_2[4] & mask_hi); + r2_1[10] = r1_1[10] - __uint_as_float(r1_bf16_2[5] << 16); + r2_1[11] = r1_1[11] - __uint_as_float(r1_bf16_2[5] & mask_hi); + r2_1[12] = r1_1[12] - __uint_as_float(r1_bf16_2[6] << 16); + r2_1[13] = r1_1[13] - __uint_as_float(r1_bf16_2[6] & mask_hi); + r2_1[14] = r1_1[14] - __uint_as_float(r1_bf16_2[7] << 16); + r2_1[15] = r1_1[15] - __uint_as_float(r1_bf16_2[7] & mask_hi); + unsigned int r2_bf16_2[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2_1[_lp*2 + 0], r2_1[_lp*2+1 + 0])); + r2_bf16_2[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv0_addr + (unsigned int)(col_1 / 64 * 8192 + t_row * 128 + col_1 % 64 * 2 ^ (col_1 / 64 * 8192 + t_row * 128 + col_1 % 64 * 2 >> 7 & 7) << 4))), "r"(tmem_load_0_bf16_1[0]), "r"(tmem_load_0_bf16_1[1]), "r"(tmem_load_0_bf16_1[2]), "r"(tmem_load_0_bf16_1[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv0_addr + (unsigned int)((col_1 + 8) / 64 * 8192 + t_row * 128 + (col_1 + 8) % 64 * 2 ^ ((col_1 + 8) / 64 * 8192 + t_row * 128 + (col_1 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(tmem_load_0_bf16_1[4]), "r"(tmem_load_0_bf16_1[5]), "r"(tmem_load_0_bf16_1[6]), "r"(tmem_load_0_bf16_1[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv1_addr + (unsigned int)(col_1 / 64 * 8192 + t_row * 128 + col_1 % 64 * 2 ^ (col_1 / 64 * 8192 + t_row * 128 + col_1 % 64 * 2 >> 7 & 7) << 4))), "r"(r1_bf16_2[0]), "r"(r1_bf16_2[1]), "r"(r1_bf16_2[2]), "r"(r1_bf16_2[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv1_addr + (unsigned int)((col_1 + 8) / 64 * 8192 + t_row * 128 + (col_1 + 8) % 64 * 2 ^ ((col_1 + 8) / 64 * 8192 + t_row * 128 + (col_1 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r1_bf16_2[4]), "r"(r1_bf16_2[5]), "r"(r1_bf16_2[6]), "r"(r1_bf16_2[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv2_addr + (unsigned int)(col_1 / 64 * 8192 + t_row * 128 + col_1 % 64 * 2 ^ (col_1 / 64 * 8192 + t_row * 128 + col_1 % 64 * 2 >> 7 & 7) << 4))), "r"(r2_bf16_2[0]), "r"(r2_bf16_2[1]), "r"(r2_bf16_2[2]), "r"(r2_bf16_2[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv2_addr + (unsigned int)((col_1 + 8) / 64 * 8192 + t_row * 128 + (col_1 + 8) % 64 * 2 ^ ((col_1 + 8) / 64 * 8192 + t_row * 128 + (col_1 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r2_bf16_2[4]), "r"(r2_bf16_2[5]), "r"(r2_bf16_2[6]), "r"(r2_bf16_2[7]) : "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(bv_ready_addr); + mbarrier_wait(h_ready_addr, _phase_h_ready_0); + _phase_h_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int slot = out_index[seg]; + if (slot >= 0) { + long long tile = (long long)((head / out_hg * out_n + slot) * out_hg + head % out_hg) * (long long)tile_stride * 16384; + float _tmem_load_0[16]; + tmem_ld_x16(&_tmem_load_0[0], taddr + (unsigned int)(h_row_base + cb)); + float _tmem_load_36_0[16]; + tmem_ld_x16(&_tmem_load_36_0[0], taddr + 128 + (unsigned int)(h_row_base + cb)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int j_1_2_1 = 0; j_1_2_1 < 16; j_1_2_1++) { + _tmem_load_0[j_1_2_1] = _tmem_load_0[j_1_2_1] + _tmem_load_36_0[j_1_2_1]; + } + if (write_transposed == 1) { + long long row_at_1_0 = tile + (long long)(k_row * 128 + v0 + cb); + { + float4 _v4 = make_float4(_tmem_load_0[0 + 0], _tmem_load_0[0 + 1], _tmem_load_0[0 + 2], _tmem_load_0[0 + 3]); + *reinterpret_cast(out + row_at_1_0) = _v4; + } + { + float4 _v4 = make_float4(_tmem_load_0[4 + 0], _tmem_load_0[4 + 1], _tmem_load_0[4 + 2], _tmem_load_0[4 + 3]); + *reinterpret_cast(out + row_at_1_0 + 4) = _v4; + } + { + float4 _v4 = make_float4(_tmem_load_0[8 + 0], _tmem_load_0[8 + 1], _tmem_load_0[8 + 2], _tmem_load_0[8 + 3]); + *reinterpret_cast(out + row_at_1_0 + 8) = _v4; + } + { + float4 _v4 = make_float4(_tmem_load_0[12 + 0], _tmem_load_0[12 + 1], _tmem_load_0[12 + 2], _tmem_load_0[12 + 3]); + *reinterpret_cast(out + row_at_1_0 + 12) = _v4; + } + } else { + #pragma unroll + for (int j_1_3 = 0; j_1_3 < 16; j_1_3++) { + *(reinterpret_cast(out) + (tile + (long long)((v0 + cb + j_1_3) * 128 + k_row))) = _tmem_load_0[j_1_3]; + } + } + float _tmem_load_1[16]; + tmem_ld_x16(&_tmem_load_1[0], taddr + (unsigned int)(h_row_base + cb) + 16); + float _tmem_load_36_1[16]; + tmem_ld_x16(&_tmem_load_36_1[0], taddr + 128 + (unsigned int)(h_row_base + cb) + 16); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int j_1_4 = 0; j_1_4 < 16; j_1_4++) { + _tmem_load_1[j_1_4] = _tmem_load_1[j_1_4] + _tmem_load_36_1[j_1_4]; + } + if (write_transposed == 1) { + long long row_at_1_1 = tile + (long long)(k_row * 128 + v0 + cb + 16); + { + float4 _v4 = make_float4(_tmem_load_1[0 + 0], _tmem_load_1[0 + 1], _tmem_load_1[0 + 2], _tmem_load_1[0 + 3]); + *reinterpret_cast(out + row_at_1_1) = _v4; + } + { + float4 _v4 = make_float4(_tmem_load_1[4 + 0], _tmem_load_1[4 + 1], _tmem_load_1[4 + 2], _tmem_load_1[4 + 3]); + *reinterpret_cast(out + row_at_1_1 + 4) = _v4; + } + { + float4 _v4 = make_float4(_tmem_load_1[8 + 0], _tmem_load_1[8 + 1], _tmem_load_1[8 + 2], _tmem_load_1[8 + 3]); + *reinterpret_cast(out + row_at_1_1 + 8) = _v4; + } + { + float4 _v4 = make_float4(_tmem_load_1[12 + 0], _tmem_load_1[12 + 1], _tmem_load_1[12 + 2], _tmem_load_1[12 + 3]); + *reinterpret_cast(out + row_at_1_1 + 12) = _v4; + } + } else { + #pragma unroll + for (int j_1_5 = 0; j_1_5 < 16; j_1_5++) { + *(reinterpret_cast(out) + (tile + (long long)((v0 + cb + 16 + j_1_5) * 128 + k_row))) = _tmem_load_1[j_1_5]; + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("barrier.sync 1, 256;" ::: "memory"); + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + } + } + } + // ---- Role: load ---- + if (warp == 8) { + { // load_main + int vslice_1 = blockIdx.x % 2; + int head_1 = blockIdx.x / 2 % num_heads; + int seg_1 = walk_seg[blockIdx.x / (2 * num_heads)]; + int c0_1 = seq_chunk_start[seg_1]; + int nc_1 = seq_chunk_start[seg_1 + 1] - c0_1; + int row0_1 = c0_1 * 64; + unsigned int _phase_w_empty = 1; + if (elect_sync()) { + #pragma unroll 1 + for (int c_1 = 0; c_1 < nc_1; c_1++) { + int trow_1 = row0_1 + c_1 * 64; + mbarrier_wait(w_empty_addr, _phase_w_empty); + mbarrier_arrive_expect_tx(w_full_addr, 16384); + #pragma unroll + for (int half = 0; half < 2; half++) { + tma_3d_gmem2smem(w_panel_addr + (0 * 2 + (unsigned int)half) * 8192, (&w_tma), half * 64, head_1, trow_1, w_full_addr); + } + if (nc_1 > c_1 + 1) { + int trow_next = row0_1 + (c_1 + 1) * 64; + #pragma unroll + for (int half_1 = 0; half_1 < 2; half_1++) { + asm volatile("cp.async.bulk.prefetch.tensor.3d.L2.global.tile [%0, {%1, %2, %3}];" :: "l"((uint64_t)((&w_tma))), "r"((int)(half_1 * 64)), "r"((int)(head_1)), "r"((int)(trow_next)) : "memory"); + asm volatile("cp.async.bulk.prefetch.tensor.3d.L2.global.tile [%0, {%1, %2, %3}];" :: "l"((uint64_t)((&kg_tma))), "r"((int)(half_1 * 64)), "r"((int)(head_1)), "r"((int)(trow_next)) : "memory"); + } + } + mbarrier_wait(kg_empty_addr, _phase_w_empty); + mbarrier_arrive_expect_tx(kg_full_addr, 16384); + #pragma unroll + for (int half_2 = 0; half_2 < 2; half_2++) { + tma_3d_gmem2smem(kg_panel_addr + (0 * 2 + (unsigned int)half_2) * 8192, (&kg_tma), half_2 * 64, head_1, trow_1, kg_full_addr); + } + _phase_w_empty ^= 1; + } + } + } + } + // ---- Role: mma ---- + if (warp == 9) { + { // mma_main + int vslice_2 = blockIdx.x % 2; + int head_2 = blockIdx.x / 2 % num_heads; + int seg_2 = walk_seg[blockIdx.x / (2 * num_heads)]; + int c0_2 = seq_chunk_start[seg_2]; + int nc_2 = seq_chunk_start[seg_2 + 1] - c0_2; + int row0_2 = c0_2 * 64; + unsigned int _phase_w_full = 0; + unsigned int _phase_hs_ready_0 = 0; + unsigned int _phase_bv_ready_0 = 0; + if (elect_sync()) { + #pragma unroll 1 + for (int c_2 = 0; c_2 < nc_2; c_2++) { + mbarrier_wait(w_full_addr, _phase_w_full); + mbarrier_wait(hs_ready_addr, _phase_hs_ready_0); + _phase_hs_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x7FFF) + (0) * 1024; + int _smem_desc_lo_1 = (((hs_mn2_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x3FFF) + (0) * 1024; + int _smem_desc_lo_1 = (((hs_mn2_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 68224144;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_acc_vp), "r"(0) : "memory"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = (((w_op_addr) >> 4) & 0x7FFF) + (0) * 1024; + int _smem_desc_lo_3 = (((hs_mn1_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_2 = (((w_op_addr) >> 4) & 0x3FFF) + (0) * 1024; + int _smem_desc_lo_3 = (((hs_mn1_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 68224144;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_acc_vp), "r"(1) : "memory"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_4 = (((w_op_addr) >> 4) & 0x7FFF) + (0) * 1024; + int _smem_desc_lo_5 = (((hs_mn0_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_4 = (((w_op_addr) >> 4) & 0x3FFF) + (0) * 1024; + int _smem_desc_lo_5 = (((hs_mn0_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 68224144;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_4), "r"(_smem_desc_lo_5), "r"(tmem_acc_vp), "r"(1) : "memory"); + tcgen05_commit(vp_ready_addr); + tcgen05_commit(w_empty_addr); + mbarrier_wait(kg_full_addr, _phase_w_full); + mbarrier_wait(bv_ready_addr, _phase_bv_ready_0); + _phase_bv_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_6 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (0) * 1024; + int _smem_desc_lo_7 = (((bv_mn2_addr) >> 4) & 0x7FFF) | 0x2000000; +#else + int _smem_desc_lo_6 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (0) * 1024; + int _smem_desc_lo_7 = (((bv_mn2_addr) >> 4) & 0x3FFF) | 0x2000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 135365776;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_6), "r"(_smem_desc_lo_7), "r"(tmem_acc_dx), "r"(0) : "memory"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_8 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (0) * 1024; + int _smem_desc_lo_9 = (((bv_mn1_addr) >> 4) & 0x7FFF) | 0x2000000; +#else + int _smem_desc_lo_8 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (0) * 1024; + int _smem_desc_lo_9 = (((bv_mn1_addr) >> 4) & 0x3FFF) | 0x2000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 135365776;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_8), "r"(_smem_desc_lo_9), "r"(tmem_acc_dx), "r"(1) : "memory"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_10 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (0) * 1024; + int _smem_desc_lo_11 = (((bv_mn0_addr) >> 4) & 0x7FFF) | 0x2000000; +#else + int _smem_desc_lo_10 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (0) * 1024; + int _smem_desc_lo_11 = (((bv_mn0_addr) >> 4) & 0x3FFF) | 0x2000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 135365776;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_10), "r"(_smem_desc_lo_11), "r"(tmem_acc_dx), "r"(1) : "memory"); + tcgen05_commit(h_ready_addr); + tcgen05_commit(kg_empty_addr); + _phase_w_full ^= 1; + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu b/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu deleted file mode 100644 index 25c103e5a8d..00000000000 --- a/flashinfer/experimental/cake_kda_det_train_kcp/csrc/cake_kda_det_train_9a3f306bacf8c8a1547d_kcp_kernel.cu +++ /dev/null @@ -1,1452 +0,0 @@ -/* - * Copyright (c) 2026 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 192 -#define TMEM_ACC_H_OFFSET 0 -#define TMEM_ACC_VP_OFFSET 64 -#define TMEM_ACC_DX_OFFSET 128 -#define NUM_RING_PIPE_STAGES 1 -#define NUM_ONE_STAGES 1 -#define SMEM_W_PANEL_OFF 1024 -#define SMEM_W_PANEL_STAGE_BYTES 8192 -#define SMEM_W_PANEL_STRIDE 8192 -#define SMEM_KG_PANEL_OFF 17408 -#define SMEM_KG_PANEL_STAGE_BYTES 8192 -#define SMEM_KG_PANEL_STRIDE 8192 -#define SMEM_W_OP_OFF 1024 -#define SMEM_W_OP_STAGE_BYTES 16384 -#define SMEM_W_OP_STRIDE 16384 -#define SMEM_KG_MN_OFF 17408 -#define SMEM_KG_MN_STAGE_BYTES 16384 -#define SMEM_KG_MN_STRIDE 16384 -#define SMEM_HS0_OFF 33792 -#define SMEM_HS0_STAGE_BYTES 16384 -#define SMEM_HS0_STRIDE 16384 -#define SMEM_HS1_OFF 50176 -#define SMEM_HS1_STAGE_BYTES 16384 -#define SMEM_HS1_STRIDE 16384 -#define SMEM_HS2_OFF 66560 -#define SMEM_HS2_STAGE_BYTES 16384 -#define SMEM_HS2_STRIDE 16384 -#define SMEM_HS_MN0_OFF 33792 -#define SMEM_HS_MN0_STAGE_BYTES 16384 -#define SMEM_HS_MN0_STRIDE 16384 -#define SMEM_HS_MN1_OFF 50176 -#define SMEM_HS_MN1_STAGE_BYTES 16384 -#define SMEM_HS_MN1_STRIDE 16384 -#define SMEM_HS_MN2_OFF 66560 -#define SMEM_HS_MN2_STAGE_BYTES 16384 -#define SMEM_HS_MN2_STRIDE 16384 -#define SMEM_BV0_OFF 82944 -#define SMEM_BV0_STAGE_BYTES 8192 -#define SMEM_BV0_STRIDE 8192 -#define SMEM_BV1_OFF 91136 -#define SMEM_BV1_STAGE_BYTES 8192 -#define SMEM_BV1_STRIDE 8192 -#define SMEM_BV2_OFF 99328 -#define SMEM_BV2_STAGE_BYTES 8192 -#define SMEM_BV2_STRIDE 8192 -#define SMEM_BV_MN0_OFF 82944 -#define SMEM_BV_MN0_STAGE_BYTES 8192 -#define SMEM_BV_MN0_STRIDE 8192 -#define SMEM_BV_MN1_OFF 91136 -#define SMEM_BV_MN1_STAGE_BYTES 8192 -#define SMEM_BV_MN1_STRIDE 8192 -#define SMEM_BV_MN2_OFF 99328 -#define SMEM_BV_MN2_STAGE_BYTES 8192 -#define SMEM_BV_MN2_STRIDE 8192 -#define SMEM_TOTAL 107520 -#define THREADS 192 -#define LAUNCH_MIN_BLOCKS 1 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - - - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - - - -__device__ __forceinline__ void tma_3d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.3d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4}], [%5];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -extern "C" { - -__global__ __launch_bounds__(192, LAUNCH_MIN_BLOCKS) void -kernel_cake_kda_det_train_9a3f306bacf8c8a1547d_kcp(const __grid_constant__ CUtensorMap w_tma, const __grid_constant__ CUtensorMap kg_tma, __nv_bfloat16* __restrict__ u, float* __restrict__ gk, int num_heads, int* __restrict__ seq_chunk_start, int num_seqs, int* __restrict__ walk_seg, float* __restrict__ out, int* __restrict__ out_index, int out_hg, int out_n, int tile_stride, int init_mode, int use_u, int write_transposed) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; -#if __CUDA_ARCH__ == 1000 - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); -#else - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); -#endif - - const int mbar_base = smem; - #define tiles_full_addr (mbar_base + 0) - #define tiles_empty_addr (mbar_base + 8) - #define hs_ready_addr (mbar_base + 16) - #define vp_ready_addr (mbar_base + 24) - #define bv_ready_addr (mbar_base + 32) - #define h_ready_addr (mbar_base + 40) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* w_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int w_panel_addr = smem + 1024; - __nv_bfloat16* kg_panel = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); - const int kg_panel_addr = smem + 17408; - __nv_bfloat16* w_op = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int w_op_addr = smem + 1024; - __nv_bfloat16* kg_mn = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); - const int kg_mn_addr = smem + 17408; - __nv_bfloat16* hs0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); - const int hs0_addr = smem + 33792; - __nv_bfloat16* hs1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); - const int hs1_addr = smem + 50176; - __nv_bfloat16* hs2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int hs2_addr = smem + 66560; - __nv_bfloat16* hs_mn0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); - const int hs_mn0_addr = smem + 33792; - __nv_bfloat16* hs_mn1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); - const int hs_mn1_addr = smem + 50176; - __nv_bfloat16* hs_mn2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int hs_mn2_addr = smem + 66560; - __nv_bfloat16* bv0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 82944); - const int bv0_addr = smem + 82944; - __nv_bfloat16* bv1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 91136); - const int bv1_addr = smem + 91136; - __nv_bfloat16* bv2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); - const int bv2_addr = smem + 99328; - __nv_bfloat16* bv_mn0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 82944); - const int bv_mn0_addr = smem + 82944; - __nv_bfloat16* bv_mn1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 91136); - const int bv_mn1_addr = smem + 91136; - __nv_bfloat16* bv_mn2 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); - const int bv_mn2_addr = smem + 99328; - - // Mbarrier init (6 pipeline groups, 0 ordered-sequence groups, 6 barriers) - // Mbarriers at smem_raw[0..48) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'ring_pipe' --- - // tiles_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // tiles_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'one' --- - // hs_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 16, 128); - // vp_ready: 1 barriers, init_count=1 - mbarrier_init(smem + 24, 1); - // bv_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 32, 128); - // h_ready: 1 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (256 columns, 192 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); - if (warp == 0) { - int _tmem_hold = smem + 48; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_acc_h = taddr; - const int tmem_acc_vp = taddr + 64; - const int tmem_acc_dx = taddr + 128; - - // ---- Role: compute ---- - if (warp <= 3) { - { // compute_main - int vslice = blockIdx.x % 2; - int head = blockIdx.x / 2 % num_heads; - int seg = walk_seg[blockIdx.x / (2 * num_heads)]; - int c0 = seq_chunk_start[seg]; - int nc = seq_chunk_start[seg + 1] - c0; - int row0 = c0 * 64; - int warp_in_wg = warp % 4; - int k_row = warp_in_wg * 32 + lane; - int h_row_base = warp_in_wg * 32 << 16; - int lane_half = lane >> 4; - int t_row = warp_in_wg * 16 + (lane & 15); - int vp_row_base = warp_in_wg * 32 << 16; - int v0 = vslice * 64; - unsigned int mask_hi = 4294901760u; - float init_row[16]; - #pragma unroll - for (int j = 0; j < 16; j++) { - float one_here = 0.0f; - if (init_mode == 0) { - if (v0 + j == k_row) { - one_here = 1.0f; - } - } - init_row[j] = one_here; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base, init_row); - #pragma unroll - for (int j_1 = 0; j_1 < 16; j_1++) { - float one_here_1 = 0.0f; - if (init_mode == 0) { - if (v0 + 16 + j_1 == k_row) { - one_here_1 = 1.0f; - } - } - init_row[j_1] = one_here_1; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + 16, init_row); - #pragma unroll - for (int j_2 = 0; j_2 < 16; j_2++) { - float one_here_2 = 0.0f; - if (init_mode == 0) { - if (v0 + 32 + j_2 == k_row) { - one_here_2 = 1.0f; - } - } - init_row[j_2] = one_here_2; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + 32, init_row); - #pragma unroll - for (int j_3 = 0; j_3 < 16; j_3++) { - float one_here_3 = 0.0f; - if (init_mode == 0) { - if (v0 + 48 + j_3 == k_row) { - one_here_3 = 1.0f; - } - } - init_row[j_3] = one_here_3; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + 48, init_row); - float zeros[16]; - zeros[0] = 0.0f; - zeros[1] = 0.0f; - zeros[2] = 0.0f; - zeros[3] = 0.0f; - zeros[4] = 0.0f; - zeros[5] = 0.0f; - zeros[6] = 0.0f; - zeros[7] = 0.0f; - zeros[8] = 0.0f; - zeros[9] = 0.0f; - zeros[10] = 0.0f; - zeros[11] = 0.0f; - zeros[12] = 0.0f; - zeros[13] = 0.0f; - zeros[14] = 0.0f; - zeros[15] = 0.0f; - tmem_st_x16_f32(taddr + 128 + (unsigned int)h_row_base, zeros); - tmem_st_x16_f32(taddr + 128 + (unsigned int)h_row_base + 16, zeros); - tmem_st_x16_f32(taddr + 128 + (unsigned int)h_row_base + 32, zeros); - tmem_st_x16_f32(taddr + 128 + (unsigned int)h_row_base + 48, zeros); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - unsigned int _phase_vp_ready_0 = 0; - unsigned int _phase_h_ready_0 = 0; - #pragma unroll 1 - for (int c = 0; c < nc; c++) { - long long gn_index = ((long long)(row0 + c * 64 + 64 - 1) * (long long)num_heads + (long long)head) * 128 + (long long)k_row; - float decay = approx_exp2(gk[gn_index]); - float _tmem_load_0[16]; - tmem_ld_x16(&_tmem_load_0[0], taddr + (unsigned int)h_row_base); - float _tmem_load_1[16]; - tmem_ld_x16(&_tmem_load_1[0], taddr + 128 + (unsigned int)h_row_base); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_4 = 0; j_4 < 16; j_4++) { - _tmem_load_0[j_4] = _tmem_load_0[j_4] + _tmem_load_1[j_4]; - } - unsigned int _tmem_load_0_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r1[16]; - r1[0] = _tmem_load_0[0] - __uint_as_float(_tmem_load_0_bf16[0] << 16); - r1[1] = _tmem_load_0[1] - __uint_as_float(_tmem_load_0_bf16[0] & mask_hi); - r1[2] = _tmem_load_0[2] - __uint_as_float(_tmem_load_0_bf16[1] << 16); - r1[3] = _tmem_load_0[3] - __uint_as_float(_tmem_load_0_bf16[1] & mask_hi); - r1[4] = _tmem_load_0[4] - __uint_as_float(_tmem_load_0_bf16[2] << 16); - r1[5] = _tmem_load_0[5] - __uint_as_float(_tmem_load_0_bf16[2] & mask_hi); - r1[6] = _tmem_load_0[6] - __uint_as_float(_tmem_load_0_bf16[3] << 16); - r1[7] = _tmem_load_0[7] - __uint_as_float(_tmem_load_0_bf16[3] & mask_hi); - r1[8] = _tmem_load_0[8] - __uint_as_float(_tmem_load_0_bf16[4] << 16); - r1[9] = _tmem_load_0[9] - __uint_as_float(_tmem_load_0_bf16[4] & mask_hi); - r1[10] = _tmem_load_0[10] - __uint_as_float(_tmem_load_0_bf16[5] << 16); - r1[11] = _tmem_load_0[11] - __uint_as_float(_tmem_load_0_bf16[5] & mask_hi); - r1[12] = _tmem_load_0[12] - __uint_as_float(_tmem_load_0_bf16[6] << 16); - r1[13] = _tmem_load_0[13] - __uint_as_float(_tmem_load_0_bf16[6] & mask_hi); - r1[14] = _tmem_load_0[14] - __uint_as_float(_tmem_load_0_bf16[7] << 16); - r1[15] = _tmem_load_0[15] - __uint_as_float(_tmem_load_0_bf16[7] & mask_hi); - unsigned int r1_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1[_lp*2 + 0], r1[_lp*2+1 + 0])); - r1_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r2[16]; - r2[0] = r1[0] - __uint_as_float(r1_bf16[0] << 16); - r2[1] = r1[1] - __uint_as_float(r1_bf16[0] & mask_hi); - r2[2] = r1[2] - __uint_as_float(r1_bf16[1] << 16); - r2[3] = r1[3] - __uint_as_float(r1_bf16[1] & mask_hi); - r2[4] = r1[4] - __uint_as_float(r1_bf16[2] << 16); - r2[5] = r1[5] - __uint_as_float(r1_bf16[2] & mask_hi); - r2[6] = r1[6] - __uint_as_float(r1_bf16[3] << 16); - r2[7] = r1[7] - __uint_as_float(r1_bf16[3] & mask_hi); - r2[8] = r1[8] - __uint_as_float(r1_bf16[4] << 16); - r2[9] = r1[9] - __uint_as_float(r1_bf16[4] & mask_hi); - r2[10] = r1[10] - __uint_as_float(r1_bf16[5] << 16); - r2[11] = r1[11] - __uint_as_float(r1_bf16[5] & mask_hi); - r2[12] = r1[12] - __uint_as_float(r1_bf16[6] << 16); - r2[13] = r1[13] - __uint_as_float(r1_bf16[6] & mask_hi); - r2[14] = r1[14] - __uint_as_float(r1_bf16[7] << 16); - r2[15] = r1[15] - __uint_as_float(r1_bf16[7] & mask_hi); - unsigned int r2_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2[_lp*2 + 0], r2[_lp*2+1 + 0])); - r2_bf16[_lp] = *(uint32_t*)&_bf2; - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 ^ (k_row * 128 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[0]), "r"(_tmem_load_0_bf16[1]), "r"(_tmem_load_0_bf16[2]), "r"(_tmem_load_0_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 + 16 ^ (k_row * 128 + 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[4]), "r"(_tmem_load_0_bf16[5]), "r"(_tmem_load_0_bf16[6]), "r"(_tmem_load_0_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 ^ (k_row * 128 >> 7 & 7) << 4))), "r"(r1_bf16[0]), "r"(r1_bf16[1]), "r"(r1_bf16[2]), "r"(r1_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 + 16 ^ (k_row * 128 + 16 >> 7 & 7) << 4))), "r"(r1_bf16[4]), "r"(r1_bf16[5]), "r"(r1_bf16[6]), "r"(r1_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 ^ (k_row * 128 >> 7 & 7) << 4))), "r"(r2_bf16[0]), "r"(r2_bf16[1]), "r"(r2_bf16[2]), "r"(r2_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 + 16 ^ (k_row * 128 + 16 >> 7 & 7) << 4))), "r"(r2_bf16[4]), "r"(r2_bf16[5]), "r"(r2_bf16[6]), "r"(r2_bf16[7]) : "memory"); - float hs2_0[16]; - #pragma unroll - for (int j_5 = 0; j_5 < 16; j_5++) { - hs2_0[j_5] = _tmem_load_0[j_5] * decay; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base, hs2_0); - float _tmem_load_2[16]; - tmem_ld_x16(&_tmem_load_2[0], taddr + (unsigned int)h_row_base + 16); - float _tmem_load_3[16]; - tmem_ld_x16(&_tmem_load_3[0], taddr + 128 + (unsigned int)h_row_base + 16); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_6 = 0; j_6 < 16; j_6++) { - _tmem_load_2[j_6] = _tmem_load_2[j_6] + _tmem_load_3[j_6]; - } - unsigned int _tmem_load_2_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_2[_lp*2 + 0], _tmem_load_2[_lp*2+1 + 0])); - _tmem_load_2_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r1_1[16]; - r1_1[0] = _tmem_load_2[0] - __uint_as_float(_tmem_load_2_bf16[0] << 16); - r1_1[1] = _tmem_load_2[1] - __uint_as_float(_tmem_load_2_bf16[0] & mask_hi); - r1_1[2] = _tmem_load_2[2] - __uint_as_float(_tmem_load_2_bf16[1] << 16); - r1_1[3] = _tmem_load_2[3] - __uint_as_float(_tmem_load_2_bf16[1] & mask_hi); - r1_1[4] = _tmem_load_2[4] - __uint_as_float(_tmem_load_2_bf16[2] << 16); - r1_1[5] = _tmem_load_2[5] - __uint_as_float(_tmem_load_2_bf16[2] & mask_hi); - r1_1[6] = _tmem_load_2[6] - __uint_as_float(_tmem_load_2_bf16[3] << 16); - r1_1[7] = _tmem_load_2[7] - __uint_as_float(_tmem_load_2_bf16[3] & mask_hi); - r1_1[8] = _tmem_load_2[8] - __uint_as_float(_tmem_load_2_bf16[4] << 16); - r1_1[9] = _tmem_load_2[9] - __uint_as_float(_tmem_load_2_bf16[4] & mask_hi); - r1_1[10] = _tmem_load_2[10] - __uint_as_float(_tmem_load_2_bf16[5] << 16); - r1_1[11] = _tmem_load_2[11] - __uint_as_float(_tmem_load_2_bf16[5] & mask_hi); - r1_1[12] = _tmem_load_2[12] - __uint_as_float(_tmem_load_2_bf16[6] << 16); - r1_1[13] = _tmem_load_2[13] - __uint_as_float(_tmem_load_2_bf16[6] & mask_hi); - r1_1[14] = _tmem_load_2[14] - __uint_as_float(_tmem_load_2_bf16[7] << 16); - r1_1[15] = _tmem_load_2[15] - __uint_as_float(_tmem_load_2_bf16[7] & mask_hi); - unsigned int r1_1_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1_1[_lp*2 + 0], r1_1[_lp*2+1 + 0])); - r1_1_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r2_2[16]; - r2_2[0] = r1_1[0] - __uint_as_float(r1_1_bf16[0] << 16); - r2_2[1] = r1_1[1] - __uint_as_float(r1_1_bf16[0] & mask_hi); - r2_2[2] = r1_1[2] - __uint_as_float(r1_1_bf16[1] << 16); - r2_2[3] = r1_1[3] - __uint_as_float(r1_1_bf16[1] & mask_hi); - r2_2[4] = r1_1[4] - __uint_as_float(r1_1_bf16[2] << 16); - r2_2[5] = r1_1[5] - __uint_as_float(r1_1_bf16[2] & mask_hi); - r2_2[6] = r1_1[6] - __uint_as_float(r1_1_bf16[3] << 16); - r2_2[7] = r1_1[7] - __uint_as_float(r1_1_bf16[3] & mask_hi); - r2_2[8] = r1_1[8] - __uint_as_float(r1_1_bf16[4] << 16); - r2_2[9] = r1_1[9] - __uint_as_float(r1_1_bf16[4] & mask_hi); - r2_2[10] = r1_1[10] - __uint_as_float(r1_1_bf16[5] << 16); - r2_2[11] = r1_1[11] - __uint_as_float(r1_1_bf16[5] & mask_hi); - r2_2[12] = r1_1[12] - __uint_as_float(r1_1_bf16[6] << 16); - r2_2[13] = r1_1[13] - __uint_as_float(r1_1_bf16[6] & mask_hi); - r2_2[14] = r1_1[14] - __uint_as_float(r1_1_bf16[7] << 16); - r2_2[15] = r1_1[15] - __uint_as_float(r1_1_bf16[7] & mask_hi); - unsigned int r2_2_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2_2[_lp*2 + 0], r2_2[_lp*2+1 + 0])); - r2_2_bf16[_lp] = *(uint32_t*)&_bf2; - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 + 32 ^ (k_row * 128 + 32 >> 7 & 7) << 4))), "r"(_tmem_load_2_bf16[0]), "r"(_tmem_load_2_bf16[1]), "r"(_tmem_load_2_bf16[2]), "r"(_tmem_load_2_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 + 48 ^ (k_row * 128 + 48 >> 7 & 7) << 4))), "r"(_tmem_load_2_bf16[4]), "r"(_tmem_load_2_bf16[5]), "r"(_tmem_load_2_bf16[6]), "r"(_tmem_load_2_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 + 32 ^ (k_row * 128 + 32 >> 7 & 7) << 4))), "r"(r1_1_bf16[0]), "r"(r1_1_bf16[1]), "r"(r1_1_bf16[2]), "r"(r1_1_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 + 48 ^ (k_row * 128 + 48 >> 7 & 7) << 4))), "r"(r1_1_bf16[4]), "r"(r1_1_bf16[5]), "r"(r1_1_bf16[6]), "r"(r1_1_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 + 32 ^ (k_row * 128 + 32 >> 7 & 7) << 4))), "r"(r2_2_bf16[0]), "r"(r2_2_bf16[1]), "r"(r2_2_bf16[2]), "r"(r2_2_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 + 48 ^ (k_row * 128 + 48 >> 7 & 7) << 4))), "r"(r2_2_bf16[4]), "r"(r2_2_bf16[5]), "r"(r2_2_bf16[6]), "r"(r2_2_bf16[7]) : "memory"); - float hs2_3[16]; - #pragma unroll - for (int j_7 = 0; j_7 < 16; j_7++) { - hs2_3[j_7] = _tmem_load_2[j_7] * decay; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + 16, hs2_3); - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], taddr + (unsigned int)h_row_base + 32); - float _tmem_load_5[16]; - tmem_ld_x16(&_tmem_load_5[0], taddr + 128 + (unsigned int)h_row_base + 32); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_8 = 0; j_8 < 16; j_8++) { - _tmem_load_4[j_8] = _tmem_load_4[j_8] + _tmem_load_5[j_8]; - } - unsigned int _tmem_load_4_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); - _tmem_load_4_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r1_4[16]; - r1_4[0] = _tmem_load_4[0] - __uint_as_float(_tmem_load_4_bf16[0] << 16); - r1_4[1] = _tmem_load_4[1] - __uint_as_float(_tmem_load_4_bf16[0] & mask_hi); - r1_4[2] = _tmem_load_4[2] - __uint_as_float(_tmem_load_4_bf16[1] << 16); - r1_4[3] = _tmem_load_4[3] - __uint_as_float(_tmem_load_4_bf16[1] & mask_hi); - r1_4[4] = _tmem_load_4[4] - __uint_as_float(_tmem_load_4_bf16[2] << 16); - r1_4[5] = _tmem_load_4[5] - __uint_as_float(_tmem_load_4_bf16[2] & mask_hi); - r1_4[6] = _tmem_load_4[6] - __uint_as_float(_tmem_load_4_bf16[3] << 16); - r1_4[7] = _tmem_load_4[7] - __uint_as_float(_tmem_load_4_bf16[3] & mask_hi); - r1_4[8] = _tmem_load_4[8] - __uint_as_float(_tmem_load_4_bf16[4] << 16); - r1_4[9] = _tmem_load_4[9] - __uint_as_float(_tmem_load_4_bf16[4] & mask_hi); - r1_4[10] = _tmem_load_4[10] - __uint_as_float(_tmem_load_4_bf16[5] << 16); - r1_4[11] = _tmem_load_4[11] - __uint_as_float(_tmem_load_4_bf16[5] & mask_hi); - r1_4[12] = _tmem_load_4[12] - __uint_as_float(_tmem_load_4_bf16[6] << 16); - r1_4[13] = _tmem_load_4[13] - __uint_as_float(_tmem_load_4_bf16[6] & mask_hi); - r1_4[14] = _tmem_load_4[14] - __uint_as_float(_tmem_load_4_bf16[7] << 16); - r1_4[15] = _tmem_load_4[15] - __uint_as_float(_tmem_load_4_bf16[7] & mask_hi); - unsigned int r1_4_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1_4[_lp*2 + 0], r1_4[_lp*2+1 + 0])); - r1_4_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r2_5[16]; - r2_5[0] = r1_4[0] - __uint_as_float(r1_4_bf16[0] << 16); - r2_5[1] = r1_4[1] - __uint_as_float(r1_4_bf16[0] & mask_hi); - r2_5[2] = r1_4[2] - __uint_as_float(r1_4_bf16[1] << 16); - r2_5[3] = r1_4[3] - __uint_as_float(r1_4_bf16[1] & mask_hi); - r2_5[4] = r1_4[4] - __uint_as_float(r1_4_bf16[2] << 16); - r2_5[5] = r1_4[5] - __uint_as_float(r1_4_bf16[2] & mask_hi); - r2_5[6] = r1_4[6] - __uint_as_float(r1_4_bf16[3] << 16); - r2_5[7] = r1_4[7] - __uint_as_float(r1_4_bf16[3] & mask_hi); - r2_5[8] = r1_4[8] - __uint_as_float(r1_4_bf16[4] << 16); - r2_5[9] = r1_4[9] - __uint_as_float(r1_4_bf16[4] & mask_hi); - r2_5[10] = r1_4[10] - __uint_as_float(r1_4_bf16[5] << 16); - r2_5[11] = r1_4[11] - __uint_as_float(r1_4_bf16[5] & mask_hi); - r2_5[12] = r1_4[12] - __uint_as_float(r1_4_bf16[6] << 16); - r2_5[13] = r1_4[13] - __uint_as_float(r1_4_bf16[6] & mask_hi); - r2_5[14] = r1_4[14] - __uint_as_float(r1_4_bf16[7] << 16); - r2_5[15] = r1_4[15] - __uint_as_float(r1_4_bf16[7] & mask_hi); - unsigned int r2_5_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2_5[_lp*2 + 0], r2_5[_lp*2+1 + 0])); - r2_5_bf16[_lp] = *(uint32_t*)&_bf2; - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 + 64 ^ (k_row * 128 + 64 >> 7 & 7) << 4))), "r"(_tmem_load_4_bf16[0]), "r"(_tmem_load_4_bf16[1]), "r"(_tmem_load_4_bf16[2]), "r"(_tmem_load_4_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 + 80 ^ (k_row * 128 + 80 >> 7 & 7) << 4))), "r"(_tmem_load_4_bf16[4]), "r"(_tmem_load_4_bf16[5]), "r"(_tmem_load_4_bf16[6]), "r"(_tmem_load_4_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 + 64 ^ (k_row * 128 + 64 >> 7 & 7) << 4))), "r"(r1_4_bf16[0]), "r"(r1_4_bf16[1]), "r"(r1_4_bf16[2]), "r"(r1_4_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 + 80 ^ (k_row * 128 + 80 >> 7 & 7) << 4))), "r"(r1_4_bf16[4]), "r"(r1_4_bf16[5]), "r"(r1_4_bf16[6]), "r"(r1_4_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 + 64 ^ (k_row * 128 + 64 >> 7 & 7) << 4))), "r"(r2_5_bf16[0]), "r"(r2_5_bf16[1]), "r"(r2_5_bf16[2]), "r"(r2_5_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 + 80 ^ (k_row * 128 + 80 >> 7 & 7) << 4))), "r"(r2_5_bf16[4]), "r"(r2_5_bf16[5]), "r"(r2_5_bf16[6]), "r"(r2_5_bf16[7]) : "memory"); - float hs2_6[16]; - #pragma unroll - for (int j_9 = 0; j_9 < 16; j_9++) { - hs2_6[j_9] = _tmem_load_4[j_9] * decay; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + 32, hs2_6); - float _tmem_load_6[16]; - tmem_ld_x16(&_tmem_load_6[0], taddr + (unsigned int)h_row_base + 48); - float _tmem_load_7[16]; - tmem_ld_x16(&_tmem_load_7[0], taddr + 128 + (unsigned int)h_row_base + 48); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_10 = 0; j_10 < 16; j_10++) { - _tmem_load_6[j_10] = _tmem_load_6[j_10] + _tmem_load_7[j_10]; - } - unsigned int _tmem_load_6_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_6[_lp*2 + 0], _tmem_load_6[_lp*2+1 + 0])); - _tmem_load_6_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r1_7[16]; - r1_7[0] = _tmem_load_6[0] - __uint_as_float(_tmem_load_6_bf16[0] << 16); - r1_7[1] = _tmem_load_6[1] - __uint_as_float(_tmem_load_6_bf16[0] & mask_hi); - r1_7[2] = _tmem_load_6[2] - __uint_as_float(_tmem_load_6_bf16[1] << 16); - r1_7[3] = _tmem_load_6[3] - __uint_as_float(_tmem_load_6_bf16[1] & mask_hi); - r1_7[4] = _tmem_load_6[4] - __uint_as_float(_tmem_load_6_bf16[2] << 16); - r1_7[5] = _tmem_load_6[5] - __uint_as_float(_tmem_load_6_bf16[2] & mask_hi); - r1_7[6] = _tmem_load_6[6] - __uint_as_float(_tmem_load_6_bf16[3] << 16); - r1_7[7] = _tmem_load_6[7] - __uint_as_float(_tmem_load_6_bf16[3] & mask_hi); - r1_7[8] = _tmem_load_6[8] - __uint_as_float(_tmem_load_6_bf16[4] << 16); - r1_7[9] = _tmem_load_6[9] - __uint_as_float(_tmem_load_6_bf16[4] & mask_hi); - r1_7[10] = _tmem_load_6[10] - __uint_as_float(_tmem_load_6_bf16[5] << 16); - r1_7[11] = _tmem_load_6[11] - __uint_as_float(_tmem_load_6_bf16[5] & mask_hi); - r1_7[12] = _tmem_load_6[12] - __uint_as_float(_tmem_load_6_bf16[6] << 16); - r1_7[13] = _tmem_load_6[13] - __uint_as_float(_tmem_load_6_bf16[6] & mask_hi); - r1_7[14] = _tmem_load_6[14] - __uint_as_float(_tmem_load_6_bf16[7] << 16); - r1_7[15] = _tmem_load_6[15] - __uint_as_float(_tmem_load_6_bf16[7] & mask_hi); - unsigned int r1_7_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1_7[_lp*2 + 0], r1_7[_lp*2+1 + 0])); - r1_7_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r2_8[16]; - r2_8[0] = r1_7[0] - __uint_as_float(r1_7_bf16[0] << 16); - r2_8[1] = r1_7[1] - __uint_as_float(r1_7_bf16[0] & mask_hi); - r2_8[2] = r1_7[2] - __uint_as_float(r1_7_bf16[1] << 16); - r2_8[3] = r1_7[3] - __uint_as_float(r1_7_bf16[1] & mask_hi); - r2_8[4] = r1_7[4] - __uint_as_float(r1_7_bf16[2] << 16); - r2_8[5] = r1_7[5] - __uint_as_float(r1_7_bf16[2] & mask_hi); - r2_8[6] = r1_7[6] - __uint_as_float(r1_7_bf16[3] << 16); - r2_8[7] = r1_7[7] - __uint_as_float(r1_7_bf16[3] & mask_hi); - r2_8[8] = r1_7[8] - __uint_as_float(r1_7_bf16[4] << 16); - r2_8[9] = r1_7[9] - __uint_as_float(r1_7_bf16[4] & mask_hi); - r2_8[10] = r1_7[10] - __uint_as_float(r1_7_bf16[5] << 16); - r2_8[11] = r1_7[11] - __uint_as_float(r1_7_bf16[5] & mask_hi); - r2_8[12] = r1_7[12] - __uint_as_float(r1_7_bf16[6] << 16); - r2_8[13] = r1_7[13] - __uint_as_float(r1_7_bf16[6] & mask_hi); - r2_8[14] = r1_7[14] - __uint_as_float(r1_7_bf16[7] << 16); - r2_8[15] = r1_7[15] - __uint_as_float(r1_7_bf16[7] & mask_hi); - unsigned int r2_8_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2_8[_lp*2 + 0], r2_8[_lp*2+1 + 0])); - r2_8_bf16[_lp] = *(uint32_t*)&_bf2; - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 + 96 ^ (k_row * 128 + 96 >> 7 & 7) << 4))), "r"(_tmem_load_6_bf16[0]), "r"(_tmem_load_6_bf16[1]), "r"(_tmem_load_6_bf16[2]), "r"(_tmem_load_6_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs0_addr + (unsigned int)(k_row * 128 + 112 ^ (k_row * 128 + 112 >> 7 & 7) << 4))), "r"(_tmem_load_6_bf16[4]), "r"(_tmem_load_6_bf16[5]), "r"(_tmem_load_6_bf16[6]), "r"(_tmem_load_6_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 + 96 ^ (k_row * 128 + 96 >> 7 & 7) << 4))), "r"(r1_7_bf16[0]), "r"(r1_7_bf16[1]), "r"(r1_7_bf16[2]), "r"(r1_7_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs1_addr + (unsigned int)(k_row * 128 + 112 ^ (k_row * 128 + 112 >> 7 & 7) << 4))), "r"(r1_7_bf16[4]), "r"(r1_7_bf16[5]), "r"(r1_7_bf16[6]), "r"(r1_7_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 + 96 ^ (k_row * 128 + 96 >> 7 & 7) << 4))), "r"(r2_8_bf16[0]), "r"(r2_8_bf16[1]), "r"(r2_8_bf16[2]), "r"(r2_8_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((hs2_addr + (unsigned int)(k_row * 128 + 112 ^ (k_row * 128 + 112 >> 7 & 7) << 4))), "r"(r2_8_bf16[4]), "r"(r2_8_bf16[5]), "r"(r2_8_bf16[6]), "r"(r2_8_bf16[7]) : "memory"); - float hs2_9[16]; - #pragma unroll - for (int j_11 = 0; j_11 < 16; j_11++) { - hs2_9[j_11] = _tmem_load_6[j_11] * decay; - } - tmem_st_x16_f32(taddr + (unsigned int)h_row_base + 48, hs2_9); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(hs_ready_addr); - mbarrier_wait(vp_ready_addr, _phase_vp_ready_0); - _phase_vp_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int trow = row0 + c * 64 + t_row; - long long u_base = ((long long)trow * (long long)num_heads + (long long)head) * 128 + (long long)v0; - float _tmem_load_8[16]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" - : "=r"(*reinterpret_cast(&_tmem_load_8[0])), "=r"(*reinterpret_cast(&_tmem_load_8[1])), "=r"(*reinterpret_cast(&_tmem_load_8[2])), "=r"(*reinterpret_cast(&_tmem_load_8[3])), "=r"(*reinterpret_cast(&_tmem_load_8[4])), "=r"(*reinterpret_cast(&_tmem_load_8[5])), "=r"(*reinterpret_cast(&_tmem_load_8[6])), "=r"(*reinterpret_cast(&_tmem_load_8[7])), "=r"(*reinterpret_cast(&_tmem_load_8[8])), "=r"(*reinterpret_cast(&_tmem_load_8[9])), "=r"(*reinterpret_cast(&_tmem_load_8[10])), "=r"(*reinterpret_cast(&_tmem_load_8[11])), "=r"(*reinterpret_cast(&_tmem_load_8[12])), "=r"(*reinterpret_cast(&_tmem_load_8[13])), "=r"(*reinterpret_cast(&_tmem_load_8[14])), "=r"(*reinterpret_cast(&_tmem_load_8[15])) - : "r"(taddr + 64 + (unsigned int)vp_row_base)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - int col = lane_half * 16; - float uv[16]; - uv[0] = 0.0f; - uv[1] = 0.0f; - uv[2] = 0.0f; - uv[3] = 0.0f; - uv[4] = 0.0f; - uv[5] = 0.0f; - uv[6] = 0.0f; - uv[7] = 0.0f; - uv[8] = 0.0f; - uv[9] = 0.0f; - uv[10] = 0.0f; - uv[11] = 0.0f; - uv[12] = 0.0f; - uv[13] = 0.0f; - uv[14] = 0.0f; - uv[15] = 0.0f; - if (use_u == 1) { - { - const uint4* _vptr_0 = reinterpret_cast(u + u_base + (long long)col); - uint4 _vld_0[2]; - #pragma unroll - for (int _blk = 0; _blk < 2; _blk++) { - _vld_0[_blk] = _vptr_0[_blk]; - uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); - #pragma unroll - for (int _pair = 0; _pair < 4; _pair++) { - (&uv[0 + _blk * 8 + _pair * 2])[0] = __uint_as_float(static_cast(_vpairs_0[_pair]) << 16); - (&uv[0 + _blk * 8 + _pair * 2])[1] = __uint_as_float(static_cast(_vpairs_0[_pair]) & 0xffff0000u); - } - } - } - } - float bvv[16]; - #pragma unroll - for (int j_12 = 0; j_12 < 16; j_12++) { - bvv[j_12] = uv[j_12] - _tmem_load_8[j_12]; - } - unsigned int bvv_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv[_lp*2 + 0], bvv[_lp*2+1 + 0])); - bvv_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r1_10[16]; - r1_10[0] = bvv[0] - __uint_as_float(bvv_bf16[0] << 16); - r1_10[1] = bvv[1] - __uint_as_float(bvv_bf16[0] & mask_hi); - r1_10[2] = bvv[2] - __uint_as_float(bvv_bf16[1] << 16); - r1_10[3] = bvv[3] - __uint_as_float(bvv_bf16[1] & mask_hi); - r1_10[4] = bvv[4] - __uint_as_float(bvv_bf16[2] << 16); - r1_10[5] = bvv[5] - __uint_as_float(bvv_bf16[2] & mask_hi); - r1_10[6] = bvv[6] - __uint_as_float(bvv_bf16[3] << 16); - r1_10[7] = bvv[7] - __uint_as_float(bvv_bf16[3] & mask_hi); - r1_10[8] = bvv[8] - __uint_as_float(bvv_bf16[4] << 16); - r1_10[9] = bvv[9] - __uint_as_float(bvv_bf16[4] & mask_hi); - r1_10[10] = bvv[10] - __uint_as_float(bvv_bf16[5] << 16); - r1_10[11] = bvv[11] - __uint_as_float(bvv_bf16[5] & mask_hi); - r1_10[12] = bvv[12] - __uint_as_float(bvv_bf16[6] << 16); - r1_10[13] = bvv[13] - __uint_as_float(bvv_bf16[6] & mask_hi); - r1_10[14] = bvv[14] - __uint_as_float(bvv_bf16[7] << 16); - r1_10[15] = bvv[15] - __uint_as_float(bvv_bf16[7] & mask_hi); - unsigned int r1_10_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1_10[_lp*2 + 0], r1_10[_lp*2+1 + 0])); - r1_10_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r2_11[16]; - r2_11[0] = r1_10[0] - __uint_as_float(r1_10_bf16[0] << 16); - r2_11[1] = r1_10[1] - __uint_as_float(r1_10_bf16[0] & mask_hi); - r2_11[2] = r1_10[2] - __uint_as_float(r1_10_bf16[1] << 16); - r2_11[3] = r1_10[3] - __uint_as_float(r1_10_bf16[1] & mask_hi); - r2_11[4] = r1_10[4] - __uint_as_float(r1_10_bf16[2] << 16); - r2_11[5] = r1_10[5] - __uint_as_float(r1_10_bf16[2] & mask_hi); - r2_11[6] = r1_10[6] - __uint_as_float(r1_10_bf16[3] << 16); - r2_11[7] = r1_10[7] - __uint_as_float(r1_10_bf16[3] & mask_hi); - r2_11[8] = r1_10[8] - __uint_as_float(r1_10_bf16[4] << 16); - r2_11[9] = r1_10[9] - __uint_as_float(r1_10_bf16[4] & mask_hi); - r2_11[10] = r1_10[10] - __uint_as_float(r1_10_bf16[5] << 16); - r2_11[11] = r1_10[11] - __uint_as_float(r1_10_bf16[5] & mask_hi); - r2_11[12] = r1_10[12] - __uint_as_float(r1_10_bf16[6] << 16); - r2_11[13] = r1_10[13] - __uint_as_float(r1_10_bf16[6] & mask_hi); - r2_11[14] = r1_10[14] - __uint_as_float(r1_10_bf16[7] << 16); - r2_11[15] = r1_10[15] - __uint_as_float(r1_10_bf16[7] & mask_hi); - unsigned int r2_11_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2_11[_lp*2 + 0], r2_11[_lp*2+1 + 0])); - r2_11_bf16[_lp] = *(uint32_t*)&_bf2; - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv0_addr + (unsigned int)(col / 64 * 8192 + t_row * 128 + col % 64 * 2 ^ (col / 64 * 8192 + t_row * 128 + col % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_bf16[0]), "r"(bvv_bf16[1]), "r"(bvv_bf16[2]), "r"(bvv_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv0_addr + (unsigned int)((col + 8) / 64 * 8192 + t_row * 128 + (col + 8) % 64 * 2 ^ ((col + 8) / 64 * 8192 + t_row * 128 + (col + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_bf16[4]), "r"(bvv_bf16[5]), "r"(bvv_bf16[6]), "r"(bvv_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv1_addr + (unsigned int)(col / 64 * 8192 + t_row * 128 + col % 64 * 2 ^ (col / 64 * 8192 + t_row * 128 + col % 64 * 2 >> 7 & 7) << 4))), "r"(r1_10_bf16[0]), "r"(r1_10_bf16[1]), "r"(r1_10_bf16[2]), "r"(r1_10_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv1_addr + (unsigned int)((col + 8) / 64 * 8192 + t_row * 128 + (col + 8) % 64 * 2 ^ ((col + 8) / 64 * 8192 + t_row * 128 + (col + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r1_10_bf16[4]), "r"(r1_10_bf16[5]), "r"(r1_10_bf16[6]), "r"(r1_10_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv2_addr + (unsigned int)(col / 64 * 8192 + t_row * 128 + col % 64 * 2 ^ (col / 64 * 8192 + t_row * 128 + col % 64 * 2 >> 7 & 7) << 4))), "r"(r2_11_bf16[0]), "r"(r2_11_bf16[1]), "r"(r2_11_bf16[2]), "r"(r2_11_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv2_addr + (unsigned int)((col + 8) / 64 * 8192 + t_row * 128 + (col + 8) % 64 * 2 ^ ((col + 8) / 64 * 8192 + t_row * 128 + (col + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r2_11_bf16[4]), "r"(r2_11_bf16[5]), "r"(r2_11_bf16[6]), "r"(r2_11_bf16[7]) : "memory"); - float _tmem_load_9[16]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 16;" - : "=r"(*reinterpret_cast(&_tmem_load_9[0])), "=r"(*reinterpret_cast(&_tmem_load_9[1])), "=r"(*reinterpret_cast(&_tmem_load_9[2])), "=r"(*reinterpret_cast(&_tmem_load_9[3])), "=r"(*reinterpret_cast(&_tmem_load_9[4])), "=r"(*reinterpret_cast(&_tmem_load_9[5])), "=r"(*reinterpret_cast(&_tmem_load_9[6])), "=r"(*reinterpret_cast(&_tmem_load_9[7])), "=r"(*reinterpret_cast(&_tmem_load_9[8])), "=r"(*reinterpret_cast(&_tmem_load_9[9])), "=r"(*reinterpret_cast(&_tmem_load_9[10])), "=r"(*reinterpret_cast(&_tmem_load_9[11])), "=r"(*reinterpret_cast(&_tmem_load_9[12])), "=r"(*reinterpret_cast(&_tmem_load_9[13])), "=r"(*reinterpret_cast(&_tmem_load_9[14])), "=r"(*reinterpret_cast(&_tmem_load_9[15])) - : "r"(taddr + 64 + (unsigned int)vp_row_base + 32)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - int col_12 = 32 + lane_half * 16; - float uv_13[16]; - uv_13[0] = 0.0f; - uv_13[1] = 0.0f; - uv_13[2] = 0.0f; - uv_13[3] = 0.0f; - uv_13[4] = 0.0f; - uv_13[5] = 0.0f; - uv_13[6] = 0.0f; - uv_13[7] = 0.0f; - uv_13[8] = 0.0f; - uv_13[9] = 0.0f; - uv_13[10] = 0.0f; - uv_13[11] = 0.0f; - uv_13[12] = 0.0f; - uv_13[13] = 0.0f; - uv_13[14] = 0.0f; - uv_13[15] = 0.0f; - if (use_u == 1) { - { - const uint4* _vptr_1 = reinterpret_cast(u + u_base + (long long)col_12); - uint4 _vld_1[2]; - #pragma unroll - for (int _blk = 0; _blk < 2; _blk++) { - _vld_1[_blk] = _vptr_1[_blk]; - uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); - #pragma unroll - for (int _pair = 0; _pair < 4; _pair++) { - (&uv_13[0 + _blk * 8 + _pair * 2])[0] = __uint_as_float(static_cast(_vpairs_1[_pair]) << 16); - (&uv_13[0 + _blk * 8 + _pair * 2])[1] = __uint_as_float(static_cast(_vpairs_1[_pair]) & 0xffff0000u); - } - } - } - } - float bvv_14[16]; - #pragma unroll - for (int j_13 = 0; j_13 < 16; j_13++) { - bvv_14[j_13] = uv_13[j_13] - _tmem_load_9[j_13]; - } - unsigned int bvv_14_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(bvv_14[_lp*2 + 0], bvv_14[_lp*2+1 + 0])); - bvv_14_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r1_15[16]; - r1_15[0] = bvv_14[0] - __uint_as_float(bvv_14_bf16[0] << 16); - r1_15[1] = bvv_14[1] - __uint_as_float(bvv_14_bf16[0] & mask_hi); - r1_15[2] = bvv_14[2] - __uint_as_float(bvv_14_bf16[1] << 16); - r1_15[3] = bvv_14[3] - __uint_as_float(bvv_14_bf16[1] & mask_hi); - r1_15[4] = bvv_14[4] - __uint_as_float(bvv_14_bf16[2] << 16); - r1_15[5] = bvv_14[5] - __uint_as_float(bvv_14_bf16[2] & mask_hi); - r1_15[6] = bvv_14[6] - __uint_as_float(bvv_14_bf16[3] << 16); - r1_15[7] = bvv_14[7] - __uint_as_float(bvv_14_bf16[3] & mask_hi); - r1_15[8] = bvv_14[8] - __uint_as_float(bvv_14_bf16[4] << 16); - r1_15[9] = bvv_14[9] - __uint_as_float(bvv_14_bf16[4] & mask_hi); - r1_15[10] = bvv_14[10] - __uint_as_float(bvv_14_bf16[5] << 16); - r1_15[11] = bvv_14[11] - __uint_as_float(bvv_14_bf16[5] & mask_hi); - r1_15[12] = bvv_14[12] - __uint_as_float(bvv_14_bf16[6] << 16); - r1_15[13] = bvv_14[13] - __uint_as_float(bvv_14_bf16[6] & mask_hi); - r1_15[14] = bvv_14[14] - __uint_as_float(bvv_14_bf16[7] << 16); - r1_15[15] = bvv_14[15] - __uint_as_float(bvv_14_bf16[7] & mask_hi); - unsigned int r1_15_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r1_15[_lp*2 + 0], r1_15[_lp*2+1 + 0])); - r1_15_bf16[_lp] = *(uint32_t*)&_bf2; - } - float r2_16[16]; - r2_16[0] = r1_15[0] - __uint_as_float(r1_15_bf16[0] << 16); - r2_16[1] = r1_15[1] - __uint_as_float(r1_15_bf16[0] & mask_hi); - r2_16[2] = r1_15[2] - __uint_as_float(r1_15_bf16[1] << 16); - r2_16[3] = r1_15[3] - __uint_as_float(r1_15_bf16[1] & mask_hi); - r2_16[4] = r1_15[4] - __uint_as_float(r1_15_bf16[2] << 16); - r2_16[5] = r1_15[5] - __uint_as_float(r1_15_bf16[2] & mask_hi); - r2_16[6] = r1_15[6] - __uint_as_float(r1_15_bf16[3] << 16); - r2_16[7] = r1_15[7] - __uint_as_float(r1_15_bf16[3] & mask_hi); - r2_16[8] = r1_15[8] - __uint_as_float(r1_15_bf16[4] << 16); - r2_16[9] = r1_15[9] - __uint_as_float(r1_15_bf16[4] & mask_hi); - r2_16[10] = r1_15[10] - __uint_as_float(r1_15_bf16[5] << 16); - r2_16[11] = r1_15[11] - __uint_as_float(r1_15_bf16[5] & mask_hi); - r2_16[12] = r1_15[12] - __uint_as_float(r1_15_bf16[6] << 16); - r2_16[13] = r1_15[13] - __uint_as_float(r1_15_bf16[6] & mask_hi); - r2_16[14] = r1_15[14] - __uint_as_float(r1_15_bf16[7] << 16); - r2_16[15] = r1_15[15] - __uint_as_float(r1_15_bf16[7] & mask_hi); - unsigned int r2_16_bf16[8]; - #pragma unroll - for (int _lp = 0; _lp < 8; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(r2_16[_lp*2 + 0], r2_16[_lp*2+1 + 0])); - r2_16_bf16[_lp] = *(uint32_t*)&_bf2; - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv0_addr + (unsigned int)(col_12 / 64 * 8192 + t_row * 128 + col_12 % 64 * 2 ^ (col_12 / 64 * 8192 + t_row * 128 + col_12 % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_14_bf16[0]), "r"(bvv_14_bf16[1]), "r"(bvv_14_bf16[2]), "r"(bvv_14_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv0_addr + (unsigned int)((col_12 + 8) / 64 * 8192 + t_row * 128 + (col_12 + 8) % 64 * 2 ^ ((col_12 + 8) / 64 * 8192 + t_row * 128 + (col_12 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(bvv_14_bf16[4]), "r"(bvv_14_bf16[5]), "r"(bvv_14_bf16[6]), "r"(bvv_14_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv1_addr + (unsigned int)(col_12 / 64 * 8192 + t_row * 128 + col_12 % 64 * 2 ^ (col_12 / 64 * 8192 + t_row * 128 + col_12 % 64 * 2 >> 7 & 7) << 4))), "r"(r1_15_bf16[0]), "r"(r1_15_bf16[1]), "r"(r1_15_bf16[2]), "r"(r1_15_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv1_addr + (unsigned int)((col_12 + 8) / 64 * 8192 + t_row * 128 + (col_12 + 8) % 64 * 2 ^ ((col_12 + 8) / 64 * 8192 + t_row * 128 + (col_12 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r1_15_bf16[4]), "r"(r1_15_bf16[5]), "r"(r1_15_bf16[6]), "r"(r1_15_bf16[7]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv2_addr + (unsigned int)(col_12 / 64 * 8192 + t_row * 128 + col_12 % 64 * 2 ^ (col_12 / 64 * 8192 + t_row * 128 + col_12 % 64 * 2 >> 7 & 7) << 4))), "r"(r2_16_bf16[0]), "r"(r2_16_bf16[1]), "r"(r2_16_bf16[2]), "r"(r2_16_bf16[3]) : "memory"); - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((bv2_addr + (unsigned int)((col_12 + 8) / 64 * 8192 + t_row * 128 + (col_12 + 8) % 64 * 2 ^ ((col_12 + 8) / 64 * 8192 + t_row * 128 + (col_12 + 8) % 64 * 2 >> 7 & 7) << 4))), "r"(r2_16_bf16[4]), "r"(r2_16_bf16[5]), "r"(r2_16_bf16[6]), "r"(r2_16_bf16[7]) : "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(bv_ready_addr); - mbarrier_wait(h_ready_addr, _phase_h_ready_0); - _phase_h_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int slot = out_index[seg]; - if (slot >= 0) { - long long tile = (long long)((head / out_hg * out_n + slot) * out_hg + head % out_hg) * (long long)tile_stride * 16384; - float _tmem_load_10[16]; - tmem_ld_x16(&_tmem_load_10[0], taddr + (unsigned int)h_row_base); - float _tmem_load_11[16]; - tmem_ld_x16(&_tmem_load_11[0], taddr + 128 + (unsigned int)h_row_base); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_14 = 0; j_14 < 16; j_14++) { - _tmem_load_10[j_14] = _tmem_load_10[j_14] + _tmem_load_11[j_14]; - } - if (write_transposed == 1) { - long long row_at = tile + (long long)(k_row * 128 + v0); - { - float4 _v4 = make_float4(_tmem_load_10[0 + 0], _tmem_load_10[0 + 1], _tmem_load_10[0 + 2], _tmem_load_10[0 + 3]); - *reinterpret_cast(out + row_at) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_10[4 + 0], _tmem_load_10[4 + 1], _tmem_load_10[4 + 2], _tmem_load_10[4 + 3]); - *reinterpret_cast(out + row_at + 4) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_10[8 + 0], _tmem_load_10[8 + 1], _tmem_load_10[8 + 2], _tmem_load_10[8 + 3]); - *reinterpret_cast(out + row_at + 8) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_10[12 + 0], _tmem_load_10[12 + 1], _tmem_load_10[12 + 2], _tmem_load_10[12 + 3]); - *reinterpret_cast(out + row_at + 12) = _v4; - } - } else { - #pragma unroll - for (int j_15 = 0; j_15 < 16; j_15++) { - *(reinterpret_cast(out) + (tile + (long long)((v0 + j_15) * 128 + k_row))) = _tmem_load_10[j_15]; - } - } - float _tmem_load_12[16]; - tmem_ld_x16(&_tmem_load_12[0], taddr + (unsigned int)h_row_base + 16); - float _tmem_load_13[16]; - tmem_ld_x16(&_tmem_load_13[0], taddr + 128 + (unsigned int)h_row_base + 16); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_16 = 0; j_16 < 16; j_16++) { - _tmem_load_12[j_16] = _tmem_load_12[j_16] + _tmem_load_13[j_16]; - } - if (write_transposed == 1) { - long long row_at_1 = tile + (long long)(k_row * 128 + v0 + 16); - { - float4 _v4 = make_float4(_tmem_load_12[0 + 0], _tmem_load_12[0 + 1], _tmem_load_12[0 + 2], _tmem_load_12[0 + 3]); - *reinterpret_cast(out + row_at_1) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_12[4 + 0], _tmem_load_12[4 + 1], _tmem_load_12[4 + 2], _tmem_load_12[4 + 3]); - *reinterpret_cast(out + row_at_1 + 4) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_12[8 + 0], _tmem_load_12[8 + 1], _tmem_load_12[8 + 2], _tmem_load_12[8 + 3]); - *reinterpret_cast(out + row_at_1 + 8) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_12[12 + 0], _tmem_load_12[12 + 1], _tmem_load_12[12 + 2], _tmem_load_12[12 + 3]); - *reinterpret_cast(out + row_at_1 + 12) = _v4; - } - } else { - #pragma unroll - for (int j_17 = 0; j_17 < 16; j_17++) { - *(reinterpret_cast(out) + (tile + (long long)((v0 + 16 + j_17) * 128 + k_row))) = _tmem_load_12[j_17]; - } - } - float _tmem_load_14[16]; - tmem_ld_x16(&_tmem_load_14[0], taddr + (unsigned int)h_row_base + 32); - float _tmem_load_15[16]; - tmem_ld_x16(&_tmem_load_15[0], taddr + 128 + (unsigned int)h_row_base + 32); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_18 = 0; j_18 < 16; j_18++) { - _tmem_load_14[j_18] = _tmem_load_14[j_18] + _tmem_load_15[j_18]; - } - if (write_transposed == 1) { - long long row_at_2 = tile + (long long)(k_row * 128 + v0 + 32); - { - float4 _v4 = make_float4(_tmem_load_14[0 + 0], _tmem_load_14[0 + 1], _tmem_load_14[0 + 2], _tmem_load_14[0 + 3]); - *reinterpret_cast(out + row_at_2) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_14[4 + 0], _tmem_load_14[4 + 1], _tmem_load_14[4 + 2], _tmem_load_14[4 + 3]); - *reinterpret_cast(out + row_at_2 + 4) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_14[8 + 0], _tmem_load_14[8 + 1], _tmem_load_14[8 + 2], _tmem_load_14[8 + 3]); - *reinterpret_cast(out + row_at_2 + 8) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_14[12 + 0], _tmem_load_14[12 + 1], _tmem_load_14[12 + 2], _tmem_load_14[12 + 3]); - *reinterpret_cast(out + row_at_2 + 12) = _v4; - } - } else { - #pragma unroll - for (int j_19 = 0; j_19 < 16; j_19++) { - *(reinterpret_cast(out) + (tile + (long long)((v0 + 32 + j_19) * 128 + k_row))) = _tmem_load_14[j_19]; - } - } - float _tmem_load_16[16]; - tmem_ld_x16(&_tmem_load_16[0], taddr + (unsigned int)h_row_base + 48); - float _tmem_load_17[16]; - tmem_ld_x16(&_tmem_load_17[0], taddr + 128 + (unsigned int)h_row_base + 48); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - #pragma unroll - for (int j_20 = 0; j_20 < 16; j_20++) { - _tmem_load_16[j_20] = _tmem_load_16[j_20] + _tmem_load_17[j_20]; - } - if (write_transposed == 1) { - long long row_at_3 = tile + (long long)(k_row * 128 + v0 + 48); - { - float4 _v4 = make_float4(_tmem_load_16[0 + 0], _tmem_load_16[0 + 1], _tmem_load_16[0 + 2], _tmem_load_16[0 + 3]); - *reinterpret_cast(out + row_at_3) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_16[4 + 0], _tmem_load_16[4 + 1], _tmem_load_16[4 + 2], _tmem_load_16[4 + 3]); - *reinterpret_cast(out + row_at_3 + 4) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_16[8 + 0], _tmem_load_16[8 + 1], _tmem_load_16[8 + 2], _tmem_load_16[8 + 3]); - *reinterpret_cast(out + row_at_3 + 8) = _v4; - } - { - float4 _v4 = make_float4(_tmem_load_16[12 + 0], _tmem_load_16[12 + 1], _tmem_load_16[12 + 2], _tmem_load_16[12 + 3]); - *reinterpret_cast(out + row_at_3 + 12) = _v4; - } - } else { - #pragma unroll - for (int j_21 = 0; j_21 < 16; j_21++) { - *(reinterpret_cast(out) + (tile + (long long)((v0 + 48 + j_21) * 128 + k_row))) = _tmem_load_16[j_21]; - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - asm volatile("barrier.sync 1, 128;" ::: "memory"); - if (warp == 0) { - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); - } - } - } - // ---- Role: load ---- - if (warp == 4) { - { // load_main - int vslice_1 = blockIdx.x % 2; - int head_1 = blockIdx.x / 2 % num_heads; - int seg_1 = walk_seg[blockIdx.x / (2 * num_heads)]; - int c0_1 = seq_chunk_start[seg_1]; - int nc_1 = seq_chunk_start[seg_1 + 1] - c0_1; - int row0_1 = c0_1 * 64; - unsigned int _phase_tiles_empty = 1; - if (elect_sync()) { - #pragma unroll 1 - for (int c_1 = 0; c_1 < nc_1; c_1++) { - mbarrier_wait(tiles_empty_addr, _phase_tiles_empty); - mbarrier_arrive_expect_tx(tiles_full_addr, 32768); - int trow_1 = row0_1 + c_1 * 64; - #pragma unroll - for (int half = 0; half < 2; half++) { - tma_3d_gmem2smem(w_panel_addr + (0 * 2 + (unsigned int)half) * 8192, (&w_tma), half * 64, head_1, trow_1, tiles_full_addr); - tma_3d_gmem2smem(kg_panel_addr + (0 * 2 + (unsigned int)half) * 8192, (&kg_tma), half * 64, head_1, trow_1, tiles_full_addr); - } - _phase_tiles_empty ^= 1; - } - } - } - } - // ---- Role: mma ---- - if (warp == 5) { - { // mma_main - int vslice_2 = blockIdx.x % 2; - int head_2 = blockIdx.x / 2 % num_heads; - int seg_2 = walk_seg[blockIdx.x / (2 * num_heads)]; - int c0_2 = seq_chunk_start[seg_2]; - int nc_2 = seq_chunk_start[seg_2 + 1] - c0_2; - int row0_2 = c0_2 * 64; - unsigned int _phase_tiles_full = 0; - unsigned int _phase_hs_ready_0 = 0; - unsigned int _phase_bv_ready_0 = 0; - if (elect_sync()) { - #pragma unroll 1 - for (int c_2 = 0; c_2 < nc_2; c_2++) { - mbarrier_wait(tiles_full_addr, _phase_tiles_full); - mbarrier_wait(hs_ready_addr, _phase_hs_ready_0); - _phase_hs_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x7FFF) + (0) * 1024; - int _smem_desc_lo_1 = (((hs_mn2_addr) >> 4) & 0x7FFF) | 0x4000000; -#else - int _smem_desc_lo_0 = (((w_op_addr) >> 4) & 0x3FFF) + (0) * 1024; - int _smem_desc_lo_1 = (((hs_mn2_addr) >> 4) & 0x3FFF) | 0x4000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 68224144;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 506;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_acc_vp), "r"(0) : "memory"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_2 = (((w_op_addr) >> 4) & 0x7FFF) + (0) * 1024; - int _smem_desc_lo_3 = (((hs_mn1_addr) >> 4) & 0x7FFF) | 0x4000000; -#else - int _smem_desc_lo_2 = (((w_op_addr) >> 4) & 0x3FFF) + (0) * 1024; - int _smem_desc_lo_3 = (((hs_mn1_addr) >> 4) & 0x3FFF) | 0x4000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 68224144;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 506;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_acc_vp), "r"(1) : "memory"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_4 = (((w_op_addr) >> 4) & 0x7FFF) + (0) * 1024; - int _smem_desc_lo_5 = (((hs_mn0_addr) >> 4) & 0x7FFF) | 0x4000000; -#else - int _smem_desc_lo_4 = (((w_op_addr) >> 4) & 0x3FFF) + (0) * 1024; - int _smem_desc_lo_5 = (((hs_mn0_addr) >> 4) & 0x3FFF) | 0x4000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 68224144;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 506;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_4), "r"(_smem_desc_lo_5), "r"(tmem_acc_vp), "r"(1) : "memory"); - tcgen05_commit(vp_ready_addr); - mbarrier_wait(bv_ready_addr, _phase_bv_ready_0); - _phase_bv_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_6 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (0) * 1024; - int _smem_desc_lo_7 = (((bv_mn2_addr) >> 4) & 0x7FFF) | 0x2000000; -#else - int _smem_desc_lo_6 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (0) * 1024; - int _smem_desc_lo_7 = (((bv_mn2_addr) >> 4) & 0x3FFF) | 0x2000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 135365776;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_6), "r"(_smem_desc_lo_7), "r"(tmem_acc_dx), "r"(0) : "memory"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_8 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (0) * 1024; - int _smem_desc_lo_9 = (((bv_mn1_addr) >> 4) & 0x7FFF) | 0x2000000; -#else - int _smem_desc_lo_8 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (0) * 1024; - int _smem_desc_lo_9 = (((bv_mn1_addr) >> 4) & 0x3FFF) | 0x2000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 135365776;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_8), "r"(_smem_desc_lo_9), "r"(tmem_acc_dx), "r"(1) : "memory"); -#if __CUDA_ARCH__ == 1070 - int _smem_desc_lo_10 = ((((kg_mn_addr) >> 4) & 0x7FFF) | 0x2000000) + (0) * 1024; - int _smem_desc_lo_11 = (((bv_mn0_addr) >> 4) & 0x7FFF) | 0x2000000; -#else - int _smem_desc_lo_10 = ((((kg_mn_addr) >> 4) & 0x3FFF) | 0x2000000) + (0) * 1024; - int _smem_desc_lo_11 = (((bv_mn0_addr) >> 4) & 0x3FFF) | 0x2000000; -#endif - asm volatile( - "{\n\t" - ".reg .pred p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" - ".reg .b64 da, db;\n\t" - "" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 135365776;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "add.u32 alo, alo, 128;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" - "}\n" - :: "r"(_smem_desc_lo_10), "r"(_smem_desc_lo_11), "r"(tmem_acc_dx), "r"(1) : "memory"); - tcgen05_commit(h_ready_addr); - tcgen05_commit(tiles_empty_addr); - _phase_tiles_full ^= 1; - } - } - } - } - - // Cleanup -} - -} // extern "C" From 621375bfcd6bf6ed852ae1f865aa842e2acd1768 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 4 Nov 2026 19:32:17 -0800 Subject: [PATCH 10/10] cake_kda_det_train_kcp: the walk-launch geometry comment names the 320-thread split-role fwdh offset build --- .../cake_kda_det_train_kcp/cake_kcp.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py index 1f2fb6fe950..409806f157b 100644 --- a/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py +++ b/flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py @@ -1525,16 +1525,16 @@ def merge_cols_per_cta(ncols: int, n_chains: int, rule: str) -> int | None: # Walk-launch geometry one_cta_per_sm / tape_one_cta_per_sm: the K1 / K2 walk kernels' occupancy capped from the host. -# Their launch footprint (192 threads, 115712 B of dynamic shared memory) admits two CTAs per SM only at the full 228 KiB -# carveout; the smallest configuration that fits one CTA admits exactly one. The runtime picks a kernel's configuration at -# launch from its function attribute (unset here) and otherwise from the device-wide preference (cudaDeviceSetCacheConfig): -# cudaFuncCachePreferL1 selects the smallest fitting shared-memory configuration, i.e. one CTA per SM. The generated -# binding compiles the launch (cudaLaunchKernelExC with the fixed dynamic shared memory) and exports no symbol for the -# kernel's registered host function (its wrapper survives only as the mangled `__device_stub__`, which the runtime does -# not accept), so the per-function attribute cannot be set from the host; the knob brackets each capped launch with the -# device-wide preference instead (PreferL1 before, PreferNone after). The preference is host-side state read at launch: -# launches already queued keep their configuration and only the bracketed launch takes PreferL1. Host side only; no -# kernel, launch-configuration or argument change. +# Their launch footprint (192 threads, 320 for the split-role fwdh offset build; 115712 B of dynamic shared memory) +# admits two CTAs per SM only at the full 228 KiB carveout; the smallest configuration that fits one CTA admits exactly +# one. The runtime picks a kernel's configuration at launch from its function attribute (unset here) and otherwise from +# the device-wide preference (cudaDeviceSetCacheConfig): cudaFuncCachePreferL1 selects the smallest fitting +# shared-memory configuration, i.e. one CTA per SM. The generated binding compiles the launch (cudaLaunchKernelExC with +# the fixed dynamic shared memory) and exports no symbol for the kernel's registered host function (its wrapper survives +# only as the mangled `__device_stub__`, which the runtime does not accept), so the per-function attribute cannot be set +# from the host; the knob brackets each capped launch with the device-wide preference instead (PreferL1 before, +# PreferNone after). The preference is host-side state read at launch: launches already queued keep their configuration +# and only the bracketed launch takes PreferL1. Host side only; no kernel, launch-configuration or argument change. WALK_CACHE_CONFIG_CAP = 2 # cudaFuncCachePreferL1 (driver_types.h): the smallest shared-memory configuration that fits WALK_CACHE_CONFIG_DEFAULT = ( 0 # cudaFuncCachePreferNone: the runtime's own choice (today's launches)