def lazy_tensor_decls(self, func: NativeFunction, schema: LazyIrSchema) -> str: value_args = schema.filtered_args(values=True, scalars=False) # Generates lazy_{name} variables for LazyTensors wrapping input tensors lazy_tensor_decls: List[str] = [] for arg in value_args: if arg.is_wrapped_scalar: # no lazy tensor wrapper for scalars that are promoted to IR values continue elif arg.is_symint_or_list: continue # values are extracted in isValueType elif isinstance(arg.lazy_type, BaseCType): if arg.lazy_type.type is tensorListValueT: lazy_tensor_decls.append( f"auto lazy_{arg.name}_tensorlist = " f"{self.backend_namespace}::{self.get_tensorlist}({arg.name});" ) else: lazy_tensor_decls.append( f"{self.lazy_tensor_ptr} lazy_{arg.name} = " f"{self.backend_namespace}::{self.get_tensor_or_wrap_number}({arg.name}, *common_device);" ) elif isinstance(arg.lazy_type, OptionalCType): # TODO(alanwaketan): Maybe we want to apply GetLtcTensorOrCreateForWrappedNumber here, but hold it # until we encounter a real world example. lazy_tensor_decls.append( f"{self.lazy_tensor_ptr} lazy_{arg.name} = " f"{self.backend_namespace}::{self.try_get_tensor}({arg.name}.value_or(at::Tensor()));" ) else: raise AssertionError( f"TODO not sure if there are other valid types to handle here ({arg.lazy_type})" ) return ("\n ").join(lazy_tensor_decls)
def generate_non_native_lazy_ir_nodes( non_native: List[Dict[str, Any]], gen_lazy_ir: GenLazyIR ) -> List[str]: """Generate the non-native lazy IR node classes""" nodes = [] for op in non_native: # Set default properties for Non-Native IRs properties = LazyIrProperties("ShapeCache", "CanBeReused") for p in op.get("properties", []): setattr(properties, p, True) schema = LazyIrSchema(FunctionSchema.parse(op["func"]), properties) schema.opkind = op.get("opkind") nodes.append(gen_lazy_ir.gen(schema)[0]) return nodes
class ComputeShapeSignature: """ Here we use the base name as the suffix of the signature to avoid generating for in-place variants. """ def __init__(self, kernel_name: str, f: NativeFunction): self.__schema = LazyIrSchema(f.func) self.__dispatch_args = ", ".join( [a.decl() for a in dispatcher.arguments(f.func)] ) self.__call_args = ", ".join( [f"{arg.name}" for arg in self.__schema.filtered_args(generator=True)] ) self.__kernel_name = kernel_name def __decl_suffix(self) -> str: return f"{self.__kernel_name}({self.__dispatch_args})" def __call_suffix(self) -> str: return f"{self.__kernel_name}({self.__call_args})" @property def shape_decl(self) -> str: return f"TORCH_API std::vector<torch::lazy::Shape> compute_shape_{self.__decl_suffix()}" @property def shape_call(self) -> str: return f"torch::lazy::compute_shape_{self.__call_suffix()}"
def return_aten_tensor(self, func: NativeFunction, schema: LazyIrSchema) -> str: returns_length = len(schema.returns) value_args = schema.filtered_args(values=True, scalars=False) value_types_names = [f"{a.name}" for a in value_args if not a.is_wrapped_scalar] first_tensor_name = value_types_names[0] if len(value_types_names) > 0 else None bridge_str = f"""auto result = {self.create_aten_from_ltc_tensor}( {self.create_lazy_tensor(first_tensor_name)}(std::move(node), *common_device));""" if returns_length > 1: assert ( len(value_types_names) > 0 ), "Code below assumes there is at least one tensor arg" bridge_str = f"""std::vector<{self.lazy_tensor_ptr}> lazy_tensors; for (int i = 0; i < {returns_length}; i++) {{ lazy_tensors.push_back({self.create_lazy_tensor(first_tensor_name)}({getValueT()}(node, i), *common_device)); }} auto result = {self.tuple_aten_from_ltc_tensors}<{returns_length}>(lazy_tensors);""" if schema.name.name.inplace or func.func.is_out_fn(): assert returns_length == 1, ( "We assumed there was no such case where an op is an in-place variant " f"and has tuple outputs, but got tuple of len {returns_length}." ) bridge_str = f"""lazy_{first_tensor_name}->SetInPlaceIrValue(node); auto& result = {first_tensor_name};""" bridge_str += """ return result;""" return bridge_str
def get_device(self, func: NativeFunction, schema: LazyIrSchema) -> str: value_args = schema.filtered_args(values=True, scalars=False) value_types_names = [f"{a.name}" for a in value_args if not a.is_wrapped_scalar] assert ( len(value_types_names) > 0 ), "Code below assumes there is at least one tensor arg" return f"""auto common_device = torch::lazy::GetBackendDevice({', '.join(value_types_names)});
def node_ctor_inputs(schema: LazyIrSchema) -> str: """ Produce a formatted string with the arguments as passed into the constructor of a node class. """ node_ctor_values = [ node_ctor_arg_rvalue_string(arg) for arg in schema.filtered_args() ] return ", ".join(node_ctor_values)
def node_base_ctor_call(self, schema: LazyIrSchema) -> str: # backends can customize the way the node base class constructor is called, # as long as all of its arguments can be generated from information available from the schema base_ctor_value_args_list = [] for arg in schema.filtered_args(values=True, scalars=False): if isinstance(arg.lazy_type, BaseCType) or isinstance( arg.lazy_type, VectorCType): base_ctor_value_args_list.append(f"{arg.name}") elif isinstance(arg.lazy_type, OptionalCType): base_ctor_value_args_list.append( f"{arg.name}.value_or(kNullValue)") else: raise AssertionError( f"Unsupported type ({arg.lazy_type}) - add support if necessary" ) base_ctor_value_args = ", ".join(base_ctor_value_args_list) scalar_args = schema.filtered_args(values=False, scalars=True) scalar_hashes = ", ".join([f"{a.name}" for a in scalar_args]) return f"""{self.node_base}(torch::lazy::OpKind({aten_symbol(schema)}),
def node_base_ctor_call(self, schema: LazyIrSchema) -> str: value_args = schema.filtered_args(values=True, scalars=False) # backends can customize the way the node base class constructor is called, # as long as all of its arguments can be generated from information available from the schema base_ctor_value_args_list = [] for arg in value_args: if isinstance(arg.lazy_type, BaseCType) or isinstance( arg.lazy_type, VectorCType): base_ctor_value_args_list.append(f"{arg.name}") elif isinstance(arg.lazy_type, OptionalCType): base_ctor_value_args_list.append( f"{arg.name}.value_or(kNullValue)") else: raise AssertionError( f"Unsupported type ({arg.lazy_type}) - add support if necessary" ) base_ctor_value_args = ", ".join(base_ctor_value_args_list) scalar_args = schema.filtered_args(values=False, scalars=True) # Shape constuction. # Conditionally build shape depending on specified shape property if schema.properties.ShapePrecompute: shape_ctor_arg = "std::move(shapes)," elif schema.properties.ShapeCompute: shape_args = [a.name for a in value_args] shape_args.extend(a.name for a in scalar_args) shape_ctor_arg = f"compute_shape_{schema.name}({', '.join(shape_args)})," elif schema.properties.ShapeCache: shape_args = [f"operand({i})" for i in range(len(value_args))] shape_args.extend(a.name for a in scalar_args) shape_ctor_arg = f"[&](){{ return compute_shape_{schema.name}({', '.join(shape_args)})[0]; }}," else: shape_ctor_arg = "" scalar_hashes = ", ".join(f"{a.name}" for a in scalar_args) return f"""{self.node_base}(
def gen_fallback_code(schema: LazyIrSchema, overload_name: str) -> str: """ Generate code that falls back to eager conditioned on a predicate """ fallback_args = ",\n ".join( [str(arg.name) for arg in schema.filtered_args(generator=True)]) if len(overload_name): aten_op_str = f"ATEN_OP2({schema.aten_name}, {overload_name})" else: aten_op_str = f"ATEN_OP({schema.aten_name})" or_has_generator = "" if schema.generator_arg: # generators are always optional and there is never more than one, at least currently or_has_generator = f" || ({schema.generator_arg.name}.has_value() && {schema.generator_arg.name}->defined())" return f"""
def __call__(self, func: NativeFunction) -> List[str]: sig = kernel_signature(func, self.backend_index) metadata = self.backend_index.get_kernel(func) assert metadata is not None schema = LazyIrSchema(func.func) return [ f"""\ {sig.decl(name=f"{self.class_method_name}::{metadata.kernel}")} {{ {self.force_eager_fallback(func, schema)} {self.metrics(func, schema)} {self.get_device(func, schema)} {self.lazy_tensor_decls(func, schema)} {self.build_ir_node(func, schema)} {self.return_aten_tensor(func, schema)} }};\n """ ]
def shape_inference(self, func: NativeFunction, schema: LazyIrSchema) -> str: metadata = self.backend_index.get_kernel(func) assert metadata is not None all_args = schema.filtered_args() returns_length = len(schema.returns) # call the meta kernel if it exists, to compute output shape/dtype for our IR if func.structured or func.structured_delegate is not None: meta_out = """std::vector<torch::lazy::Shape> shapes{ torch::lazy::Shape(out_meta.scalar_type(), out_meta.sizes().vec())};""" if returns_length > 1: def this_shape(i: int) -> str: return f"torch::lazy::Shape(std::get<{i}>(out_meta).scalar_type(), std::get<{i}>(out_meta).sizes().vec())" shapes_str = ",".join( [this_shape(i) for i in range(returns_length)]) meta_out = "std::vector<torch::lazy::Shape> shapes{" + shapes_str + "};" shape_str = f"""auto out_meta = at::meta::{schema.aten_name}({', '.join(str(a.name) for a in all_args)}); {meta_out}""" else: shape_sig = ComputeShapeSignature(metadata.kernel, func) shape_str = f""" auto shapes = {shape_sig.shape_call};""" shape_str += f""" TORCH_INTERNAL_ASSERT(shapes.size() == {returns_length});""" # Calculating which dimensions are symbolic func_schema_str = "aten::" + str(func.func) shape_str += f""" if(torch::lazy::symbolicShapeEnabled()){{ std::vector<torch::jit::IValue> inputs = {{ {', '.join(str(a.name) for a in all_args)} }}; const char* schema_str = "{func_schema_str}"; applySymbolicShapesOnLT(schema_str, inputs, shapes); }} """ return shape_str
def shape_inference(self, func: NativeFunction, schema: LazyIrSchema) -> str: metadata = self.backend_index.get_kernel(func) assert metadata is not None all_args = schema.filtered_args() returns_length = len(schema.returns) # call the meta kernel if it exists, to compute output shape/dtype for our IR # Note [Generated LTC Shape Functions] # LTC uses meta tensors from core to do shape inference when possible, and otherwise # we generate a shape function declaration that needs to be manually implemented. # How do we detect which ops are eligible to use meta tensors? # In general we should be able to use meta tensors not just on structured operators, # but also on composite operators that are implemented in terms of structured kernels. # We don't currently have a way of knowing at codegen time which ops are implemented that way. # This is the case for all view and view_copy operators however, so we're going to # use them specifically for all of the view_copy ops (instead of manually writing shape rules for all of them). is_view_copy_op = "view_copy" in func.tags is_structured = func.structured or func.structured_delegate is not None if is_structured or is_view_copy_op: meta_out = """ std::vector<torch::lazy::Shape> shapes{torch::lazy::Shape(out_meta.scalar_type(), out_meta.sizes().vec())};""" if returns_length > 1: def this_shape(i: int) -> str: return f"torch::lazy::Shape(std::get<{i}>(out_meta).scalar_type(), std::get<{i}>(out_meta).sizes().vec())" shapes_str = ",".join( [this_shape(i) for i in range(returns_length)]) meta_out = "std::vector<torch::lazy::Shape> shapes{" + shapes_str + "};" # Convert tensor args to the meta device and call it. # (We can't pass in the input tensors directly, because they are "functional wrappers". # If any of the meta kernels call a tensor op and redispatch, we don't want to hit the functionalize kernels.) # Even at::meta:: functions might redispatch, e.g. if they call into view ops. dispatcher_sig = DispatcherSignature.from_schema(func.func) meta_conversion_str, meta_call_ctx = convert_to_meta_tensors( dispatcher_sig) meta_call_args = [ e.expr for e in translate( meta_call_ctx, dispatcher_sig.arguments(), method=False) ] if is_view_copy_op: # view_copy ops always have a CompositeExplicitAutogradNonFunctional kernel assert func.has_composite_explicit_autograd_non_functional_kernel dispatch_ns = "compositeexplicitautogradnonfunctional" else: dispatch_ns = "meta" aten_name = schema.aten_name # TODO: this is trolling if func.func.has_symint(): aten_name += "_symint" shape_str = f"""\ {meta_conversion_str} auto out_meta = at::{dispatch_ns}::{aten_name}({', '.join(meta_call_args)}); {meta_out}""" else: shape_sig = ComputeShapeSignature(metadata.kernel, func) shape_str = f""" auto shapes = {shape_sig.shape_call};""" shape_str += f""" TORCH_INTERNAL_ASSERT(shapes.size() == {returns_length});""" # Calculating which dimensions are symbolic func_schema_str = "aten::" + str(func.func) shape_str += f""" if(torch::lazy::symbolicShapeEnabled()){{ std::vector<torch::jit::IValue> inputs = {{ {', '.join(str(a.name) for a in all_args)} }}; const char* schema_str = "{func_schema_str}"; applySymbolicShapesOnLT(schema_str, inputs, shapes); }} """ return shape_str
def gen(self, schema: LazyIrSchema) -> List[str]: opkind = schema.opkind or aten_symbol(schema) # for now, we just want one IR class decl and soon after also the method defs # and we use the functional version not out/inplace. all_args = schema.filtered_args() value_args = schema.filtered_args(values=True, scalars=False) scalar_args = schema.filtered_args(values=False, scalars=True) ctor_args = [ f"const {i.lazy_type.cpp_type()}& {i.name}" for i in all_args ] reuse_ctor_args = ", ".join(ctor_args) if schema.properties.ShapePrecompute: ctor_args.append("std::vector<torch::lazy::Shape>&& shapes") node_ctor_args = ", ".join(ctor_args) scalar_initializers = ",\n ".join([ # This code is just special casing the mapping from string_view -> strings f"{a.name}({a.name}.has_value() ? c10::make_optional(std::string(*{a.name})) : c10::nullopt)" if a.lazy_type.cpp_type() == "c10::optional<c10::string_view>" else f"{a.name}({a.name})" for a in scalar_args ]) if len(scalar_initializers): scalar_initializers = f",\n {scalar_initializers}" scalar_decls = "\n ".join([ f"std::string {a.name};" if a.lazy_type.cpp_type() == "c10::string_view" else f"c10::optional<std::string> {a.name};" if a.lazy_type.cpp_type() == "c10::optional<c10::string_view>" else f"{a.lazy_type.cpp_type()} {a.name};" for a in scalar_args ]) optional_values = [ arg.name for arg in schema.filtered_args(values=True, scalars=False) if isinstance(arg.lazy_type, OptionalCType) ] has_optional_decls = "\n ".join( [f"bool has_{value}: 1;" for value in optional_values]) has_optional_defs = "\n ".join( [f"has_{value} = !!{value};" for value in optional_values]) members_to_string = [] for arg in scalar_args: if isinstance(arg.lazy_type, OptionalCType): members_to_string.append(f"""if ({arg.name}.has_value()) {{ ss << ", {arg.name}=" << {arg.name}.value(); }} else {{ ss << ", {arg.name}=null"; }}""") else: members_to_string.append( f'ss << ", {arg.name}=" << {arg.name};') members_to_string_str = "\n ".join(members_to_string) return [ f"""\ class {schema.node_name} : public {self.node_base} {{ public: static torch::lazy::OpKind ClassOpKind() {{ return torch::lazy::OpKind({opkind}); }} {schema.node_name}({node_ctor_args}) : {self.node_base_ctor_call(schema)}{scalar_initializers} {{ {has_optional_defs} }} std::string ToString() const override {{ std::stringstream ss; ss << {self.node_base}::ToString(); {members_to_string_str} return ss.str(); }} {self.create_function(schema, reuse_ctor_args)} {self.can_be_reused_function(schema, reuse_ctor_args)} {self.lowering_function(schema)} {scalar_decls} {has_optional_decls} }}; """, ]
def __call__(self, f: Union[NativeFunctionsGroup, NativeFunction]) -> List[str]: func = f.functional.func if isinstance( f, NativeFunctionsGroup) else f.func schema = LazyIrSchema(func) return self.gen(schema)
def gen(self, f: Union[NativeFunctionsGroup, NativeFunction]) -> List[str]: # for now, we just want one IR class decl and soon after also the method defs # and we use the functional version not out/inplace. func = f.functional.func if isinstance( f, NativeFunctionsGroup) else f.func schema = LazyIrSchema(func) all_args = schema.filtered_args() value_args = schema.filtered_args(values=True, scalars=False) scalar_args = schema.filtered_args(values=False, scalars=True) node_ctor_args = ", ".join( [f"const {i.lazy_type.cpp_type()}& {i.name}" for i in all_args]) scalar_initializers = ",\n ".join( [f"{a.name}({a.name})" for a in scalar_args]) comma_if_scalar_initializers = ",\n" if len( scalar_initializers) else "" scalar_decls = "\n ".join([ f"std::string {a.name};" if a.lazy_type.cpp_type() == "c10::string_view" else f"{a.lazy_type.cpp_type()} {a.name};" for a in scalar_args ]) optional_values = [ arg.name for arg in schema.filtered_args(values=True, scalars=False) if isinstance(arg.lazy_type, OptionalCType) ] has_optional_decls = "\n ".join( [f"bool has_{value}: 1;" for value in optional_values]) has_optional_defs = "\n ".join( [f"has_{value} = !!{value};" for value in optional_values]) members_to_string = [] for arg in scalar_args: if isinstance(arg.lazy_type, OptionalCType): members_to_string.append(f"""if ({arg.name}.has_value()) {{ ss << ", {arg.name}=" << {arg.name}.value(); }} else {{ ss << ", {arg.name}=null"; }}""") else: members_to_string.append( f'ss << ", {arg.name}=" << {arg.name};') members_to_string_str = "\n ".join(members_to_string) return [ f"""\ class {schema.node_name} : public {self.node_base} {{ public: {schema.node_name}({node_ctor_args}, std::vector<Shape>&& shapes) : {self.node_base_ctor_call(schema)}{comma_if_scalar_initializers} {scalar_initializers} {{ {has_optional_defs} }} std::string ToString() const override {{ std::stringstream ss; ss << {self.node_base}::ToString(); {members_to_string_str} return ss.str(); }} {self.lowering_function(f)} {scalar_decls} {has_optional_decls} }}; """, ]