[Codegen] Only use arith.select in padding on linalg.generic

[newling]

This PR updates the logic so that rewrites only happens on linalg.generic. Before, all linalg op's were having the block-rewrite logic applied to them, including named ops like linalg.matmul. Was hitting

error: expected add/mul op in the body
  %0 =  linalg.matmul {lowering_config = #iree_gpu.lowering_config<{partial_reduction = [0, 0, 32]}>}

[newling]

@Groverkss / @kuhar this will help in deprecating warp reduction

[Groverkss]

Do you know where these linalg.matmul ops are coming from? We generalize all named ops other than contractions in vector distribute i think

[kuhar]

Suggested change

	if (linalg::GenericOp genericOp =
	if (auto genericOp =

[kuhar]

We generalize all named ops other than contractions in vector distribute i think

+1, I think this was the case. The tuning infra makes this assumption too.

[Groverkss]

I have a feeling this might be crashing because of linalg.conv_2d ops, which makes sense, but we should probably add a test for conv operations then. Let's check what exactly are the ops that are failing.

[newling]

Do you know where these linalg.matmul ops are coming from? We generalize all named ops other than contractions in vector distribute i think

With a change like this : #21718

One of the e2e tests compiled as

iree-compile --output-format=vm-bytecode --mlir-print-op-on-diagnostic=false --iree-hal-target-backends=rocm --iree-hip-target=gfx942 --iree-opt-data-tiling=false --iree-dispatch-creation-experimental-data-tiling=true --iree-dispatch-creation-set-encoding-strategy=padding --iree-hip-encoding-layout-resolver=pad e2e_matmul_cdna3_pad_i8_rocm_hip_matmul.mlir -o abc.vmfb 

Ends up with

e2e_matmul_cdna3_pad_i8_rocm_hip_matmul.mlir:2:13: error: expected add/mul op in the body
  %result = linalg.matmul ins(%lhs, %rhs: tensor<?x?xi8>, tensor<?x?xi8>) outs(%acc: tensor<?x?xi32>) -> tensor<?x?xi32>
            ^
e2e_matmul_cdna3_pad_i8_rocm_hip_matmul.mlir:1:1: note: called from
util.func @matmul_accumulate_DYNxDYNxi8_times_DYNxDYNxi8_into_DYNxDYNxi32(%lhs: tensor<?x?xi8>, %rhs: tensor<?x?xi8>, %acc: tensor<?x?xi32>) -> tensor<?x?xi32> {

[...}

This error comes from the padding pass this PR touches. The failing pass is here in the pipeline:

iree/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp

Line 922 in f1ee95c

funcPassManager.addPass(createGPUApplyPaddingLevelPass(padOptions));

In the pass run in this failing case, I see

...
GPUGeneralizeNamedOpsPass
...
LLVMGPUSelectLoweringStrategyPass
...
GPUApplyPaddingLevelPass (CRASH!!!!!!!)

The obvious question then is why is it still a named op in the pass that crashes, if there is clearly a generalization pass run before it? Because the pass to generalize ops does not generalize matmul :

iree/compiler/src/iree/compiler/Codegen/Common/GPU/GPUGeneralizeNamedOps.cpp

Line 56 in f1ee95c

if (isa<linalg::BatchMatmulTransposeBOp, linalg::MatmulTransposeBOp,

which seems suspicious, but I assume it was done like that for some undocumented reason?

There is another generalization pass that gets runs later, over here

iree/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp

Line 955 in f1ee95c

funcPassManager.addPass(createLinalgGeneralizeNamedOpsPass());

But we're not getting to that.

[newling]

Maybe we need APIs like isProjectingNamedLinalgOp / isNamedLinalgOp or something, and then passes should be required to use these to make their expectations clear with assertions.

I have a feeling this might be crashing because of linalg.conv_2d ops, which makes sense, but we should probably add a test for conv operations then. Let's check what exactly are the ops that are failing.

So... it makes sense because convolutions aren't easy to specialize? So having this pass work for (named) conv but not matmul makes sense to you? Please spell out your thinking from first principles for me

[kuhar]

Named matmul op is new, probably fell through the cracks. We should generalize it instead of checking for transposed variants.

[kuhar]

The named transpose ops are going away in O(days)

[newling]

Named matmul op is new, probably fell through the cracks.

That's surprising, there was a linalg::BatchMatmulTransposeBOp before a linalg::MatmulOp ?

We should generalize it instead of checking for transposed variants.

What does this mean? I'm assuming it's not just adding it https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/Common/GPU/GPUGeneralizeNamedOps.cpp#L56 ?

[newling]

The named transpose ops are going away in O(days)

Ah ok, I hadn't been following that

[newling]

Generalize matmul: #21720

I still think this PR is a good-to-have. Either this, or an assertion/failure that linalg.matmul is not entering the llvm kernel configuration pass.