Triton IR question

[milesial]

Hey, let's see how this discussion feature works, as this is not really worth an issue.

On the blog post for the release, there is an example of an add kernel, along with the generated IR:

def void add(i32* X .aligned(16) , i32* Y .aligned(16) , i32* Z .aligned(16) , i32 N .multipleof(2) )
{
entry:
  %0 = get_program_id[0] i32;
  %1 = mul i32 %0, 512;
  %3 = make_range[0 : 512] i32<512>;
  %4 = splat i32<512> %1;
  %6 = add i32<512> %4, %3;
  %9 = splat i32<512> N;
  %11 = icmp_slt i1<512> %6, %9;
  %14 = splat i32*<512> X;
  %16 = getelementptr i32*<512> %14, %6;
  %19 = broadcast i1<512> %11;
  %21 = splat i32<512> undef;
  %22 = masked_load i32<512> %16, %19, %21;
  %26 = splat i32*<512> Y;
  %28 = getelementptr i32*<512> %26, %6;
  %31 = broadcast i1<512> %11;
  %33 = splat i32<512> undef;
  %34 = masked_load i32<512> %28, %31, %33;
  %38 = splat i32*<512> Z;
  %40 = getelementptr i32*<512> %38, %6;
  %43 = add i32<512> %22, %34;
  %46 = broadcast i32<512> %43;
  %48 = broadcast i1<512> %11;
  masked_store void %40, %46, %48;
  ret void;
}

I'm no expert in compiler IR, but I think I get it overall.
I'm curious about the following assignments:

  %19 = broadcast i1<512> %11;

  %31 = broadcast i1<512> %11;

  %48 = broadcast i1<512> %11;

As I understand it, %11 is the result of icmp_slt and is already of type i1<512>, so I have two questions:

• why do we need to broadcast it?
• why do we need to do a new broadcast each time before using the mask, instead of using the already computed value?

I guess that it's optimized down the line to PTX, but I'm just curious about this Triton IR step.

[ptillet]

To answer your question:

• We don't need to broadcast it, since it has the same shape. I have it already fixed (among some other bugs) in a local branch I plan to push in the next few days :)
• Same thing. We don't need broadcast when all the shapes are the same.

I also want to point out that Triton-IR is not really documented on purpose, as I don't think it's stable enough yet to be used by systems other than the Triton python frontend.

[milesial]

Thanks :)