Skip to content

Instantly share code, notes, and snippets.

@AmosLewis
Created January 30, 2025 18:23
Show Gist options
  • Save AmosLewis/6feff5aefe1066d2402d961d91ef3431 to your computer and use it in GitHub Desktop.
Save AmosLewis/6feff5aefe1066d2402d961d91ef3431 to your computer and use it in GitHub Desktop.
hal.executable public @_initializer_0_dispatch_0 {
hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb", {abi = "hip", iree.gpu.target = #iree_gpu.target<arch = "gfx942", features = "", wgp = <compute = fp64|fp32|fp16|int64|int32|int16|int8, storage = b64|b32|b16|b8, subgroup = shuffle|arithmetic, dot = dp4xi8toi32, mma = [<MFMA_F32_16x16x4_F32>, <MFMA_F32_16x16x16_F16>, <MFMA_F32_32x32x8_F16>, <MFMA_F64_16x16x4_F64>, <MFMA_F32_16x16x16_BF16>, <MFMA_F32_32x32x8_BF16>, <MFMA_F32_16x16x32_F8E5M2FNUZ>, <MFMA_F32_16x16x32_F8E5M2FNUZ_F8E4M3FNUZ>, <MFMA_F32_16x16x32_F8E4M3FNUZ>, <MFMA_F32_16x16x32_F8E4M3FNUZ_F8E5M2FNUZ>, <MFMA_F32_32x32x16_F8E5M2FNUZ>, <MFMA_F32_32x32x16_F8E5M2FNUZ_F8E4M3FNUZ>, <MFMA_F32_32x32x16_F8E4M3FNUZ>, <MFMA_F32_32x32x16_F8E4M3FNUZ_F8E5M2FNUZ>, <MFMA_I32_16x16x32_I8>, <MFMA_I32_32x32x16_I8>], subgroup_size_choices = [64], max_workgroup_sizes = [1024, 1024, 1024], max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 65536, max_workgroup_counts = [2147483647, 2147483647, 2147483647], max_load_instruction_bits = 128, simds_per_wgp = 4, vgpr_space_bits = 16384>>, ukernels = "none"}>) {
hal.executable.export public @_initializer_0_dispatch_0_elementwise_broadcast_64_i64 ordinal(0) layout(#hal.pipeline.layout<constants = 1, bindings = [#hal.pipeline.binding<storage_buffer>], flags = None>) {
^bb0(%arg0: !hal.device):
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
hal.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @_initializer_0_dispatch_0_elementwise_broadcast_64_i64() {
%c2_i64 = arith.constant 2 : i64
%0 = hal.interface.constant.load layout(<constants = 1, bindings = [#hal.pipeline.binding<storage_buffer>], flags = None>) ordinal(0) : i32
%1 = arith.index_castui %0 : i32 to index
%2 = util.assume.int %1[<umin = 0, umax = 0>, <umin = 768, umax = 768, udiv = 768>] : index
%3 = hal.interface.binding.subspan layout(<constants = 1, bindings = [#hal.pipeline.binding<storage_buffer>], flags = None>) binding(0) alignment(64) offset(%2) flags("None") : !flow.dispatch.tensor<writeonly:tensor<64xi64>>
%4 = tensor.empty() : tensor<64xi64>
%5 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} outs(%4 : tensor<64xi64>) {
^bb0(%out: i64):
%6 = linalg.index 0 : index
%7 = arith.index_cast %6 : index to i64
%8 = arith.muli %7, %c2_i64 : i64
linalg.yield %8 : i64
} -> tensor<64xi64>
flow.dispatch.tensor.store %5, %3, offsets = [0], sizes = [64], strides = [1] : tensor<64xi64> -> !flow.dispatch.tensor<writeonly:tensor<64xi64>>
return
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment