[TorchToLinalg] Unconventional lowering of KernelBench softmax kernel

[tkarna]

Lowering the KernelBench level1/23_Softmax.py kernel to linalg, the produced IR is functionally correct but has some odd features:

#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d0)>
#map2 = affine_map<(d0, d1) -> (d0, 0)>
module {
  func.func @main(%arg0: tensor<4096x393216xf32>) -> tensor<4096x393216xf32> {
    %c0_i64 = arith.constant 0 : i64
    %cst = arith.constant 0xFF800000 : f32
    %cst_0 = arith.constant 0.000000e+00 : f32
    %0 = tensor.empty() : tensor<4096xi64>
    %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<4096xi64>) -> tensor<4096xi64>
    %2 = tensor.empty() : tensor<4096xf32>
    %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<4096xf32>) -> tensor<4096xf32>
    %4:2 = linalg.generic {indexing_maps = [#map, #map1, #map1], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<4096x393216xf32>) outs(%3, %1 : tensor<4096xf32>, tensor<4096xi64>) {
    ^bb0(%in: f32, %out: f32, %out_1: i64):
      %12 = linalg.index 1 : index
      %13 = arith.index_cast %12 : index to i64
      %14 = arith.maximumf %in, %out : f32
      %15 = arith.cmpf ogt, %in, %out : f32
      %16 = arith.select %15, %13, %out_1 : i64
      linalg.yield %14, %16 : f32, i64
    } -> (tensor<4096xf32>, tensor<4096xi64>)
    %expanded = tensor.expand_shape %4#0 [[0, 1]] output_shape [4096, 1] : tensor<4096xf32> into tensor<4096x1xf32>
    %5 = tensor.empty() : tensor<4096x393216xf32>
    %6 = linalg.generic {indexing_maps = [#map, #map2, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %expanded : tensor<4096x393216xf32>, tensor<4096x1xf32>) outs(%5 : tensor<4096x393216xf32>) {
    ^bb0(%in: f32, %in_1: f32, %out: f32):
      %12 = arith.subf %in, %in_1 : f32
      linalg.yield %12 : f32
    } -> tensor<4096x393216xf32>
    %7 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%6 : tensor<4096x393216xf32>) outs(%5 : tensor<4096x393216xf32>) {
    ^bb0(%in: f32, %out: f32):
      %12 = math.exp %in : f32
      linalg.yield %12 : f32
    } -> tensor<4096x393216xf32>
    %8 = tensor.empty() : tensor<4096x1xf32>
    %9 = linalg.fill ins(%cst_0 : f32) outs(%8 : tensor<4096x1xf32>) -> tensor<4096x1xf32>
    %10 = linalg.generic {indexing_maps = [#map, #map2], iterator_types = ["parallel", "reduction"]} ins(%7 : tensor<4096x393216xf32>) outs(%9 : tensor<4096x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %12 = arith.addf %in, %out : f32
      linalg.yield %12 : f32
    } -> tensor<4096x1xf32>
    %11 = linalg.generic {indexing_maps = [#map, #map2, #map], iterator_types = ["parallel", "parallel"]} ins(%7, %10 : tensor<4096x393216xf32>, tensor<4096x1xf32>) outs(%5 : tensor<4096x393216xf32>) {
    ^bb0(%in: f32, %in_1: f32, %out: f32):
      %12 = arith.divf %in, %in_1 : f32
      linalg.yield %12 : f32
    } -> tensor<4096x393216xf32>
    return %11 : tensor<4096x393216xf32>
  }
}

First, the first reduction loop computes max reduction (as expected) but also computes the argmax index which is in fact never used. Canonicalization does not remove unused linalg results.

 %4:2 = linalg.generic {indexing_maps = [#map, #map1, #map1], iterator_types = ["parallel", "reduction"]} ...

Second, the output of the latter reduction loop is initialized with a shape that has a trailing unit dimension. This unit dim effectively prevents using upstream tile-and-fuse transforms on the softmax kernel.

%8 = tensor.empty() : tensor<4096x1xf32>

In principle both of these issues can be cleaned up with suitable code transformations but it would be better if torch-mlir produced canonical IR from the start.

Tested with current torch-mlir nightly build (pip install).

A standalone reproducer.

import importlib.util
from pathlib import Path

from torch_mlir import fx
from torch_mlir.fx import OutputType

KERNELBENCH_SOFTMAX_PATH = Path(
    "/path/to/KernelBench/level1/23_Softmax.py"
)


def _load_python_module(path: Path):
    spec = importlib.util.spec_from_file_location(path.stem, path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Failed to load module spec from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _get_model_and_inputs(module):
    model_cls = getattr(module, "Model")
    get_init_inputs = getattr(module, "get_init_inputs", None)
    get_inputs = getattr(module, "get_inputs")

    init_args = tuple(get_init_inputs()) if callable(get_init_inputs) else tuple()
    model = model_cls(*init_args).eval()

    sample_args = get_inputs()
    return model, sample_args


def main() -> int:
    if not KERNELBENCH_SOFTMAX_PATH.is_file():
        print(f"ERROR: benchmark file not found: {KERNELBENCH_SOFTMAX_PATH}")
        return 2

    module = _load_python_module(KERNELBENCH_SOFTMAX_PATH)
    model, sample_args = _get_model_and_inputs(module)

    imported = fx.export_and_import(
        model,
        *sample_args,
        output_type=OutputType.LINALG_ON_TENSORS,
    )

    print(imported)
    return 0


if __name__ == "__main__":
    main()