Support for all KernelBench kernels

[rengolin]

This issue is only about conformance: lower correctly and generate output compatible (allclose(0.01)) with PyTorch eager. When using the --torch-compile option, there are fewer errors, so the list below is using that flag. The additional errors when using full export can be tracked elsewhere.

Torch-MLIR issues:

• Dialect tm_tensor not found for custom op 'tm_tensor.scan' torch-mlir#4574: Level1 { 89 ~ 93 }
• error: failed to legalize operation 'torch.operator' that was explicitly marked illegal torch-mlir#4575: Level1 { 12, 96, 99 }, Level 2 { 52 }
• error: failed to legalize operation 'torch.constant.bool' torch-mlir#4578: Level 2 { 18, 22, 42, 43, 45, 51, 58, 64, 92 }
• Inconsistent lowering of Torch tril and triu torch-mlir#4604: Level 1 { 14 }
• Inconsistent lowering of 3D and 4D matmuls torch-mlir#4610 Level 1 { 10, 11 }

MLIR Issues:

• LLVM ERROR: SmallVector unable to grow (9TB): Level 1 { 34 }, Level 2 { 17, 28, 79 }
• RuntimeError: gather(): Expected dtype int32/64 for index: Level 1 { 95 }
• Outputs don't match (expected 0.01, got 1.00 ): Level 1 { 94, 98, 100 }

Look into level 3 and 4.

@adam-smnk @tkarna

[AbdallahRashed]

Hello @rengolin is there something relatively not complex I can start working on? Thank you

[rengolin]

The torch-mlir errors are better described in the torch-mlir issues, which you're already handling some. The issue with the decomposition is that torch-mlir handles it in a central place, and I don't know the code too well, so I rely on the other reviewers to help.

In here, perhaps the two first MLIR issues are simpler than the last one. Here's some hints as to what may be happening with them:

1. SmallVector growing to 9TB doesn't make sense, so it's likely a bad computation on the MLIR side. Inspecting the IR generated at various points could yield a clue.
2. The data type is assumed to be the same for all inputs, which is not true for some kernels. I assume this is the only case that crashes, but not necessarily the only case that is wrong.
3. The outputs do match for small shapes (lower number of operations), but not the ones in the python files. This is likely some legalization issue. Either torch-mlir is lowering wrong, or there's some optimization in MLIR that is broken. Printing the MLIR at every stage (--print-original-module, --print-mlir-after-all, and --print-optimized-module) could yield clues.

[rengolin]

llvm/torch-mlir#4644 merged, fixing the first part of llvm/torch-mlir#4575

[rengolin]

Kernel Level1 / 14 now working and running with matmul.yaml pipeline.

[rengolin]

As of the last Torch-MLIR & MLIR Bindings update, this is the current status of the PyTorch support in KernelBench:

The triangular matrix lowering was fixed and is now running at same (good) performance as the other. 💯

Error message	Level1 kernels	Level2 kernels
LLVM ERROR: SmallVector unable to grow. Requested capacity (...) is larger than maximum value for size type	34_InstanceNorm	None
Dialect 'tm_tensor' not found for custom op 'tm_tensor.scan'	89_cumsum, 90_cumprod, 91_cumsum_reverse, 92_cumsum_exclusive, 93_masked_cumsum	None
expected target dtype to be Long or Byte, but got Float	95_CrossEntropyLoss	None
failed to legalize operation 'torch.operator' that was explicitly marked illegal	96_HuberLoss, 99_TripletMarginLoss	52_Conv2d_Activation_BatchNorm
The output of the compiled model does't match the reference output	94_MSELoss, 98_KLDivLoss, 100_HingeLoss	None
'torch.aten.unsqueeze' op operand #0 must be Any Torch tensor type, but got '!torch.none'	None	17_Conv2d_InstanceNorm_Divide, 79_Conv3d_Multiply_InstanceNorm_Clamp_Multiply_Max
failed to legalize operation 'torch.constant.bool'	None	18_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp, 22_Matmul_Scale_ResidualAdd_Clamp_LogSumExp_Mish, 42_ConvTranspose2d_GlobalAvgPool_BiasAdd_LogSumExp_Sum_Multiply, 43_Conv3d_Max_LogSumExp_ReLU, 45_Gemm_Sigmoid_LogSumExp, 51_Gemm_Subtract_GlobalAvgPool_LogSumExp_GELU_ResidualAdd, 58_ConvTranspose3d_LogSumExp_HardSwish_Subtract_Clamp, 64_Gemm_LogSumExp_LeakyReLU_LeakyReLU_GELU_GELU, 92_Conv2d_GroupNorm_Tanh_HardSwish_ResidualAdd_LogSumExp

I also see a warning on level2/28_BMM_InstanceNorm_Sum_ResidualAdd_Multiply:

• UserWarning: input's size at dim=1 does not match num_features. You can silence this warning by not passing in num_features, which is not used because affine=False

[rengolin]

Now that we updated torch-mlir, level 1 / kernel 12 works well with llvm/torch-mlir#4644. #254

[rengolin]

Need to rebase torch-mlir again, to catch the fix for torch.constant.bool from llvm/torch-mlir#4606