Prevent DAGCombine from slicing a vector load from non-default address space

My target machine features several memory address spaces. The default address space supports loads/stores of various sizes, but a special address space (called "flex") supports only v4i32 memory accesses. This has become a problem in one case where DAGCombine tries to simplify (extract_vector_elt (v4i32 (load ...)), k) into a 32-bit load of the specific lane from the flex address space.

The faulty transformation takes place when DAGCombiner::visitEXTRACT_VECTOR_ELT calls DAGCombiner::ReplaceExtractVectorEltOfLoadWithNarrowedLoad.

This routine first checks that narrowing will result in a legal load by calling TLI.isOperationLegalOrCustom(ISD::LOAD, VecEltVT). However, this test is insensitive to the address space of the loaded pointer. Indeed, for my target this test should return true for pointers to the default address space but should return false for pointers to the flex address space. I don't immediately see any way for my target to specify this constraint here.

Is there a clean way to prevent this transformation? I don't see any way to fix this other than adding a new address-space-sensitive method to the TargetLowering interface which targets may override.

Thanks,
Nick Johnson
D. E. Shaw Research

Unfortunately there currently is no way to specify legality of a load per address space, although I really could use this ability. A special new TLI hook is the simplest option here.

-Matt