gemini-code-assist[bot] commented on code in PR #19895: URL: https://github.com/apache/tvm/pull/19895#discussion_r3484501321
########## docs/tirx/native_basics/cuda/profiling.rst: ########## @@ -0,0 +1,238 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you 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. + +In-kernel profiling with CudaProfiler +===================================== + +Once a kernel is correct and you have seen how it compiles (see +:doc:`compiling`), the next question is usually *where the cycles go*. Host-side +timers and ``nsys`` tell you how long a launch took, but not how that time splits +across the regions *inside* one kernel — the TMA loads, the mainloop MMAs, the +softmax, the epilogue. + +``tvm.tirx.bench.CudaProfiler`` is a lightweight, in-kernel event tracer for +exactly this. You bracket regions of device code with ``start`` / ``end`` +markers; at runtime one leader thread per block stamps the GPU global timer into +a buffer you pass in as an ordinary kernel argument. After the launch you read +the buffer back and decode it into per-region durations or a Perfetto timeline. + +It is *not* zero cost — every event is a ``%globaltimer`` read plus a global +store, and every thread in the region pays a block fence — so it is a +profiling/debugging tool, not something you leave on in production. + +The kernel +---------- + +The kernel below brackets a ``load`` / ``compute`` / ``store`` sequence. The +``compute`` region runs a 4000-iteration FMA loop so it clearly dominates. Events +are a plain ``enum.Enum`` whose integer values start at 0 and index a names list. + +.. code-block:: python + + from enum import Enum + import numpy as np + import tvm + from tvm.script import tirx as T + from tvm.tirx.bench import CudaProfiler, export_to_perfetto_trace + + NUM_BLOCKS, BLOCK, NUM_GROUPS = 4, 128, 1 + WRITE_STRIDE = NUM_BLOCKS * NUM_GROUPS # >= number of (block, group) lanes + PROF_SIZE = 4096 # uint64 slots in the profiler buffer + N = NUM_BLOCKS * BLOCK + + class Ev(Enum): + Load = 0 + Compute = 1 + Store = 2 Review Comment:  Using `enum.Enum` for `Ev` might cause TVM compilation or FFI errors because standard `Enum` members are not instances of `int` and cannot be automatically converted by TVM's FFI. Changing `Ev` to inherit from `enum.IntEnum` ensures that the enum members behave as integers and are compatible with TVM Script and FFI. ```suggestion class Ev(IntEnum): Load = 0 Compute = 1 Store = 2 ``` ########## docs/tirx/native_basics/cuda/profiling.rst: ########## @@ -0,0 +1,238 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you 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. + +In-kernel profiling with CudaProfiler +===================================== + +Once a kernel is correct and you have seen how it compiles (see +:doc:`compiling`), the next question is usually *where the cycles go*. Host-side +timers and ``nsys`` tell you how long a launch took, but not how that time splits +across the regions *inside* one kernel — the TMA loads, the mainloop MMAs, the +softmax, the epilogue. + +``tvm.tirx.bench.CudaProfiler`` is a lightweight, in-kernel event tracer for +exactly this. You bracket regions of device code with ``start`` / ``end`` +markers; at runtime one leader thread per block stamps the GPU global timer into +a buffer you pass in as an ordinary kernel argument. After the launch you read +the buffer back and decode it into per-region durations or a Perfetto timeline. + +It is *not* zero cost — every event is a ``%globaltimer`` read plus a global +store, and every thread in the region pays a block fence — so it is a +profiling/debugging tool, not something you leave on in production. + +The kernel +---------- + +The kernel below brackets a ``load`` / ``compute`` / ``store`` sequence. The +``compute`` region runs a 4000-iteration FMA loop so it clearly dominates. Events +are a plain ``enum.Enum`` whose integer values start at 0 and index a names list. + +.. code-block:: python + + from enum import Enum + import numpy as np + import tvm + from tvm.script import tirx as T + from tvm.tirx.bench import CudaProfiler, export_to_perfetto_trace + + NUM_BLOCKS, BLOCK, NUM_GROUPS = 4, 128, 1 + WRITE_STRIDE = NUM_BLOCKS * NUM_GROUPS # >= number of (block, group) lanes + PROF_SIZE = 4096 # uint64 slots in the profiler buffer + N = NUM_BLOCKS * BLOCK + + class Ev(Enum): + Load = 0 + Compute = 1 + Store = 2 + + EV_NAMES = ["load", "compute", "store"] + + @T.prim_func + def profiled_kernel(out_ptr: T.handle, inp_ptr: T.handle, prof_ptr: T.handle): + out = T.match_buffer(out_ptr, (N,), "float32") + inp = T.match_buffer(inp_ptr, (N,), "float32") + prof = T.match_buffer(prof_ptr, (PROF_SIZE,), "uint64") + T.device_entry() + bid = T.cta_id([NUM_BLOCKS]) + tid = T.thread_id([BLOCK]) + idx = bid * BLOCK + tid + + # Construct the profiler inside the kernel; only the leader thread writes. + p = CudaProfiler(prof, write_stride=WRITE_STRIDE, num_groups=NUM_GROUPS, + default_leader=(tid == 0)) + p.init(0) # group_id = 0; also stamps the buffer header at slot 0 + + p.start(Ev.Load) + x: T.f32 = inp[idx] + p.end(Ev.Load) + + p.start(Ev.Compute) + acc: T.f32 = T.float32(0) + for _ in range(4000): + acc = acc * T.float32(1.0001) + x + p.end(Ev.Compute) + + p.start(Ev.Store) + out[idx] = acc + p.end(Ev.Store) + + p.finalize() # mark this (block, group) lane done + +Run it and read the trace +------------------------- + +Allocate a zeroed ``uint64`` buffer, pass it as the last argument, then read it +back. Each record is one ``uint64``: the high 32 bits are the timestamp, the low +32 bits a packed tag, so decoding is plain bit-twiddling on the host. + +.. code-block:: python + + dev = tvm.cuda(0) + exe = tvm.compile(tvm.IRModule({"main": profiled_kernel}), + target=tvm.target.Target("cuda"), tir_pipeline="tirx") + + inp = tvm.runtime.tensor(np.ones(N, "float32"), device=dev) + out = tvm.runtime.tensor(np.zeros(N, "float32"), device=dev) + prof = tvm.runtime.tensor(np.zeros(PROF_SIZE, "uint64"), device=dev) Review Comment:  In TVM, `tvm.runtime.tensor` is not a standard API and will raise an `AttributeError`. Use `tvm.nd.array` instead to allocate and initialize the NDArrays. ```suggestion inp = tvm.nd.array(np.ones(N, "float32"), device=dev) out = tvm.nd.array(np.zeros(N, "float32"), device=dev) prof = tvm.nd.array(np.zeros(PROF_SIZE, "uint64"), device=dev) ``` ########## docs/tirx/native_basics/cuda/profiling.rst: ########## @@ -0,0 +1,238 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you 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. + +In-kernel profiling with CudaProfiler +===================================== + +Once a kernel is correct and you have seen how it compiles (see +:doc:`compiling`), the next question is usually *where the cycles go*. Host-side +timers and ``nsys`` tell you how long a launch took, but not how that time splits +across the regions *inside* one kernel — the TMA loads, the mainloop MMAs, the +softmax, the epilogue. + +``tvm.tirx.bench.CudaProfiler`` is a lightweight, in-kernel event tracer for +exactly this. You bracket regions of device code with ``start`` / ``end`` +markers; at runtime one leader thread per block stamps the GPU global timer into +a buffer you pass in as an ordinary kernel argument. After the launch you read +the buffer back and decode it into per-region durations or a Perfetto timeline. + +It is *not* zero cost — every event is a ``%globaltimer`` read plus a global +store, and every thread in the region pays a block fence — so it is a +profiling/debugging tool, not something you leave on in production. + +The kernel +---------- + +The kernel below brackets a ``load`` / ``compute`` / ``store`` sequence. The +``compute`` region runs a 4000-iteration FMA loop so it clearly dominates. Events +are a plain ``enum.Enum`` whose integer values start at 0 and index a names list. + +.. code-block:: python + + from enum import Enum + import numpy as np + import tvm + from tvm.script import tirx as T + from tvm.tirx.bench import CudaProfiler, export_to_perfetto_trace + + NUM_BLOCKS, BLOCK, NUM_GROUPS = 4, 128, 1 + WRITE_STRIDE = NUM_BLOCKS * NUM_GROUPS # >= number of (block, group) lanes + PROF_SIZE = 4096 # uint64 slots in the profiler buffer + N = NUM_BLOCKS * BLOCK + + class Ev(Enum): + Load = 0 + Compute = 1 + Store = 2 + + EV_NAMES = ["load", "compute", "store"] + + @T.prim_func + def profiled_kernel(out_ptr: T.handle, inp_ptr: T.handle, prof_ptr: T.handle): + out = T.match_buffer(out_ptr, (N,), "float32") + inp = T.match_buffer(inp_ptr, (N,), "float32") + prof = T.match_buffer(prof_ptr, (PROF_SIZE,), "uint64") + T.device_entry() + bid = T.cta_id([NUM_BLOCKS]) + tid = T.thread_id([BLOCK]) + idx = bid * BLOCK + tid + + # Construct the profiler inside the kernel; only the leader thread writes. + p = CudaProfiler(prof, write_stride=WRITE_STRIDE, num_groups=NUM_GROUPS, + default_leader=(tid == 0)) + p.init(0) # group_id = 0; also stamps the buffer header at slot 0 + + p.start(Ev.Load) + x: T.f32 = inp[idx] + p.end(Ev.Load) + + p.start(Ev.Compute) + acc: T.f32 = T.float32(0) + for _ in range(4000): + acc = acc * T.float32(1.0001) + x + p.end(Ev.Compute) + + p.start(Ev.Store) + out[idx] = acc + p.end(Ev.Store) + + p.finalize() # mark this (block, group) lane done + +Run it and read the trace +------------------------- + +Allocate a zeroed ``uint64`` buffer, pass it as the last argument, then read it +back. Each record is one ``uint64``: the high 32 bits are the timestamp, the low +32 bits a packed tag, so decoding is plain bit-twiddling on the host. + +.. code-block:: python + + dev = tvm.cuda(0) + exe = tvm.compile(tvm.IRModule({"main": profiled_kernel}), + target=tvm.target.Target("cuda"), tir_pipeline="tirx") + + inp = tvm.runtime.tensor(np.ones(N, "float32"), device=dev) + out = tvm.runtime.tensor(np.zeros(N, "float32"), device=dev) + prof = tvm.runtime.tensor(np.zeros(PROF_SIZE, "uint64"), device=dev) + + exe(out, inp, prof) + dev.sync() + + prof_np = prof.numpy() + opens, spans = {}, {} + for i in range(1, len(prof_np)): + word = int(prof_np[i]) + if word == 0: + continue + ts, tag = word >> 32, word & 0xFFFFFFFF + block = (tag >> 12) // NUM_GROUPS + event_idx, event_type = (tag >> 2) & 0x3FF, tag & 0x3 # 0=start 1=end 2=instant 3=finalize + if event_type == 0: + opens[(block, event_idx)] = ts + elif event_type == 1: + spans.setdefault(block, []).append((EV_NAMES[event_idx], ts - opens[(block, event_idx)])) Review Comment:  If the profiler buffer wraps or contains incomplete trace data, an end event might be processed without a matching start event in `opens`. Using `opens.pop((block, event_idx), None)` avoids a potential `KeyError` and safely ignores unmatched end events. ```suggestion elif event_type == 1: start_ts = opens.pop((block, event_idx), None) if start_ts is not None: spans.setdefault(block, []).append((EV_NAMES[event_idx], ts - start_ts)) ``` ########## docs/tirx/native_basics/cuda/profiling.rst: ########## @@ -0,0 +1,238 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you 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. + +In-kernel profiling with CudaProfiler +===================================== + +Once a kernel is correct and you have seen how it compiles (see +:doc:`compiling`), the next question is usually *where the cycles go*. Host-side +timers and ``nsys`` tell you how long a launch took, but not how that time splits +across the regions *inside* one kernel — the TMA loads, the mainloop MMAs, the +softmax, the epilogue. + +``tvm.tirx.bench.CudaProfiler`` is a lightweight, in-kernel event tracer for +exactly this. You bracket regions of device code with ``start`` / ``end`` +markers; at runtime one leader thread per block stamps the GPU global timer into +a buffer you pass in as an ordinary kernel argument. After the launch you read +the buffer back and decode it into per-region durations or a Perfetto timeline. + +It is *not* zero cost — every event is a ``%globaltimer`` read plus a global +store, and every thread in the region pays a block fence — so it is a +profiling/debugging tool, not something you leave on in production. + +The kernel +---------- + +The kernel below brackets a ``load`` / ``compute`` / ``store`` sequence. The +``compute`` region runs a 4000-iteration FMA loop so it clearly dominates. Events +are a plain ``enum.Enum`` whose integer values start at 0 and index a names list. + +.. code-block:: python + + from enum import Enum Review Comment:  Import `IntEnum` instead of `Enum` to support the integer-based enum definition for `Ev`. ```suggestion from enum import IntEnum ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
