tenstorrent / tt-mlir

Tenstorrent MLIR compiler
https://tenstorrent.github.io/tt-mlir/
Apache License 2.0
72 stars 9 forks source link

Support programming metal ComputeConfig #266

Closed nsmithtt closed 1 month ago

nsmithtt commented 3 months ago

We'll need the relevant flatbuffer attributes that the compiler can pass through to program this attribute

nsmithtt commented 2 months ago

We need a flatbuffer type that reflects metal's ComputeConfig:

        ComputeConfig{                                                                                                                                                                                                                        
            .math_fidelity = MathFidelity::HiFi4,                                                                                                                                                                                             
            .fp32_dest_acc_en = fp32_dest_acc_en,                                                                                                                                                                                             
            .math_approx_mode = math_approx_mode,                                                                                                                                                                                             
            .compile_args = compute_kernel_args,                                                                                                                                                                                              
            .defines = get_defines(BinaryOpType::ADD)                                                                                                                                                                                         
        }

At least for now, the fields that we really care about are math_fidelity, fp32_dest_acc_en, and math_approx_mode.

This new, optional, table ComputeDesc, should live in include/ttmlir/Target/TTMetal/program.fbs and become a member of table KernelDesc.

This new flatbuffer attribute, if present, should also be consumed by runtime and for programming metal API CreateKernel, here runtime/lib/ttmetal/command_queue.cpp.

kmabeeTT commented 2 months ago

Thx Nick. Adding @ddilbazTT as asignee too. Defne, let's look at this issue next when previous issue is ~finished.

ddilbazTT commented 2 months ago

So this is the bit I will add right now to /tt-mlir/include/ttmlir/Target/TTMetal/program.fbs, as I haven't figured out how to handle maps in fbs file - which I believe defines is.

table ComputeDesc {
  math_fidelity: MathFidelity;
  fp32_dest_acc_en: bool;
  math_approx_mode: bool;
  compile_args: [uint32];
  // defines: 
}
nsmithtt commented 2 months ago

Let's elide compile_args and defines for now. Not sure how we want to handle these yet.

I haven't figured out how to handle maps in fbs file

For future reference, maps can be implemented as just a list of pairs:

table MapEntry {
  key: ...;
  value: ...;
}

table Foo {
  my_map: [MapEntry];
}
nsmithtt commented 2 months ago

Offline we decided to go a new direction, and make a few additional cleanups that should help us in the long term. On the flatbuffer side of things I think we want to remove SourceType and instead embed a variant config type that serves 2 purposes:

  1. The union will generate enums for the variant type, denoting if it's a Noc/Tensix/Eth program
  2. The underlying table will carry the respective info for specializing the kernel

See attached: program.fbs.zip

On the compiler side we want:

  1. Update TTMetal_DispatchOp to take a TTKernel_KernelConfigArrayAttr instead of TTKernel_ThreadTypeArrayAttr.
  2. Each element of the array above will be TTKernel_KernelConfigAttr defined in TTKernelOpsTypes.td. These kernel configs will work in the same way as the flatbuffer schema definition and will be a variant type, something like:
    def TTKernel_KernelConfig : TTKernel_Attr<"KernelConfig", "kernel_config"> {    
    let arguments = (ins AnyAttrOf<[NocConfig, TensixConfig, EthConfig]>:$attr);    
    }
  3. Each of NocConfig, TensixConfig, EthConfig should mirror the schema type
  4. Update SerializeToBinary.cpp to plumb this info info CreateKernelSourceDirect