Stack from ghstack (oldest at bottom):

• futher scheduler changes for invoke_quant: prologue low prec, (slightly) more aggressive fusion #145104
• -> [Inductor changes] Invoke Quant #139102

Adds a invoke_quant higher order operator as proposed here.

The primary motivations are

• Unifying scattered reasoning for quant operators throughout the code base

• Easy of pattern matching - see this very large pattern match expression [here](

  pytorch/torch/_inductor/fx_passes/post_grad.py

  Lines 390 to 426 in 949fdd2

  	@register_lowering_pattern(
  	CallFunction(
  	aten.mm.default,
  	KeywordArg("mat1"),
  	CallFunction(
  	aten.sub.Tensor,
  	CallFunction(
  	prims.convert_element_type.default,
  	CallFunction(
  	aten.reshape.default,
  	CallFunction(
  	aten.cat.default,
  	ListOf(
  	CallFunction(
  	aten.bitwise_and.Scalar,
  	KeywordArg("mat2"),
  	0xF,
  	),
  	# CallFunction(
  	# aten.__rshift__.Scalar,
  	# KeywordArg("mat2"),
  	# 4,
  	# ),
  	True,
  	),
  	1,
  	),
  	KeywordArg("mat2_mm_shape"),
  	),
  	KeywordArg("mat2_dtype"),
  	),
  	8,
  	),
  	),
  	extra_check=cuda_and_enabled_mixed_mm_and_not_int8,
  	)
  	def uint4x2_mixed_mm(match: Match, mat1, mat2, mat2_mm_shape, mat2_dtype):

  . Compared to the pattern I have in the tests:

        @register_graph_pattern(
            CallFunction(
                torch.ops.aten.mm,
                CallFunction(
                    torch.ops.higher_order.invoke_quant,
                    Ignored(),
                    Ignored(),
                    Ignored(),
                    scheme="nf4",
                ),
                Arg(),
            ),
            pass_dict=test_pass,
        )

• Ability to specify inductor specific logic, like codegen'ing the operators in lower precision, or forcing fusion to a matmul.

Example graph:

 ===== AFTER POST GRAD =====
 /data/users/eellison/pytorch/torch/fx/_lazy_graph_module.py class <lambda>(torch.nn.Module):
    def forward(self, arg0_1: "f32[8][1]cpu", arg1_1: "f32[8][1]cpu"):
         # File: /data/users/eellison/pytorch/torch/_higher_order_ops/invoke_quant.py:87 in __call__, code: return invoke_quant_tracer(*args, **kwargs, quant_options=self)  # type: ignore[call-arg]
        repeated_subgraph0 = self.repeated_subgraph0
        invoke_quant: "f32[8][1]cpu" = torch.ops.higher_order.invoke_quant(repeated_subgraph0, arg0_1, arg1_1, scheme = 'nf4');  repeated_subgraph0 = arg0_1 = arg1_1 = None
        return (invoke_quant,)
        
    class repeated_subgraph0(torch.nn.Module):
        def forward(self, arg0_1: "f32[8][1]cpu", arg1_1: "f32[8][1]cpu"):
             # File: /data/users/eellison/pytorch/torch/_higher_order_ops/invoke_quant.py:87 in __call__, code: return invoke_quant_tracer(*args, **kwargs, quant_options=self)  # type: ignore[call-arg]
            mul: "f32[8][1]cpu" = torch.ops.aten.mul.Tensor(arg0_1, arg1_1);  arg0_1 = None
            add: "f32[8][1]cpu" = torch.ops.aten.add.Tensor(mul, arg1_1);  mul = arg1_1 = None
            return add

The schema for invoke_quant is torch.ops.higher_order.invoke_quant(subgraph, *args, scheme=None) where the scheme will not always be present.

I wasn't sure exactly how the inductor specific configurations like codgen_in_low_precision should be passed through. I didnt want to stuff them all in as kwargs, and I didn't want to have them affect pattern matching. So they will be stored as meta of the node itself. And, following that, I wanted the invocation of the hop to match how it will show up in the graph. So I decided to have it be an object that is then invoked for the tracing.

invoke_quant = InvokeQuant(codegen_low_precision=True)
invoke_quant(gn, (x, y), scheme="nf4") 

Todo - not require the packing of args in a tuple, will do following #139162.

Feedback welcome.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @rec