[Vulkan] Support passing 64 bit scalar

[masahi]

Implements suggestion in #7457 (comment)

Seems to work on Vulkan. Enabled get_valid_counts and cumsum tests on vulkan, which uses TIR scan and requires passing int64 scalar.

Metal runtime is also updated but not tested at all, since I don't have mac.

@tqchen @tmoreau89 @jwfromm

[jroesch]

Does this work for f64? I think that was the error I triggered the other night

[jroesch]

Can you clarify why we didn't just replace ArgUnion?

[masahi]

ArgUnion is also used by other runtime (cuda/opencl I think)

tvm/src/runtime/pack_args.h

Line 143 in 1831c17

TempArray<ArgUnion, N> holder_(num_args);

I was not sure / haven't look if we can replace this with 64. Since ArgUnion only supports 32 bit and I haven't seen any issue with it, I'm not sure if replacing this with 64 bit union is a good idea. Thoughts?

[tqchen]

cc @vinx13 @zxy844288792 @jwfromm can you please help if the change works for metal GPU backend? Thank you!

[masahi]

Just found that Metal codegen generates hardcoded __TVMArgUnion union:

tvm/src/target/source/codegen_metal.cc

Lines 49 to 51 in 053347c

	decl_stream << "union __TVMArgUnion {\n"
	<< " int v_int;\n"
	<< "};\n\n";

Do we need to change this to something like this?

  decl_stream << "union __TVMArgUnion64 {\n"
              << " int32_t v_int32[2];\n"
              << " int64_t v_int64;\n"    
              << "};\n\n";

And I don't understand why float32 is not there.

[vinx13]

@masahi Actually I don't think int (and float32) should be there, since 32-bit args are captured here

tvm/src/target/source/codegen_metal.cc

Lines 104 to 109 in 053347c

	if (v.dtype().bits() == 32) {
	decl_stream << " ";
	PrintType(v.dtype(), decl_stream);
	decl_stream << " " << vid << ";\n";
	vref << varg << "." << vid;
	} else {

[masahi]

@tqchen @vinx13 @jroesch @jwfromm I can leave the Metal runtime as an old version and only update vulkan runtime, if that sounds better. Unless someone looks at the Metal issue in detail for me there is nothing I can do.

[tqchen]

@vinx13 @jwfromm @junrushao1994 would be great if any of you who have a macbook can help double check the metal impl

I think we only need to change

tvm/src/target/source/codegen_metal.cc

Line 104 in 053347c

if (v.dtype().bits() == 32) {

to print specific types.

@masahi I think your change to __TArgUnion is right. Note that we only use ArgUnion to pass in values that are non-32 bits(e.g. int8 that needs to be passed).

We need to update the generation of the argument buffer struct for 32bit and 64bit values in L104

[vinx13]

@tqchen @masahi

tvm/src/runtime/metal/metal_module.mm

Line 200 in 6cd9626

length:num_pack_args_ * sizeof(ArgUnion64)

here the encoder assumes every argument is 64 bit, I think we need to remove this block

tvm/src/target/source/codegen_metal.cc

Lines 104 to 109 in 053347c

	if (v.dtype().bits() == 32) {
	decl_stream << " ";
	PrintType(v.dtype(), decl_stream);
	decl_stream << " " << vid << ";\n";
	vref << varg << "." << vid;
	} else {

and use TVMArgUnion64 for argument values

[tqchen]

we should still keep the 32 bit path, instead generating a int32[2] type as in our argunion

[tqchen]

@masahi please still incorporate the metal codegen change, otherwise the constant path would results in error(since we already pass in ArgUnion64. Perhaps @vinx13 can followup to suggest a change

[masahi]

I can let @vinx13 directly push to my branch

[vinx13]

I can help with it. @tqchen are you suggesting separating 32bit and 64 bit when passing args?

[tqchen]

@vinx13 we can follow the same convention, e.g. pass int32 as int32[2] and waste the second element

[masahi]

Ready to merge? @vinx13 @tqchen @jroesch

resolved

[tqchen]

Thanks @masahi @vinx13 @jroesch !