Refactor LowerAffine to use OpRewritePattern instead of ConversionPattern.

ConversionPattern should ideally only be used when the types of the operands are changing, which in this case they aren't. Using OpRewritePattern also lends to much simpler code.

PiperOrigin-RevId: 258158474

LLVMDialect: still depend on standard types.

The dependecy on standard type is mandated by attribute verifiers on some
operations (e.g., llvm.constant) that use values of those types.

PiperOrigin-RevId: 258153134

Automated rollback of changelist 258149291.

PiperOrigin-RevId: 258151618

Start moving AffineOps to ODS

This CL starts with the AffineTerminatorOp to ODS.

PiperOrigin-RevId: 258151579

Introduce loop coalescing utility and a simple pass

Multiple (perfectly) nested loops with independent bounds can be combined into
a single loop and than subdivided into blocks of arbitrary size for load
balancing or more efficient parallelism exploitation.  However, MLIR wants to
preserve the multi-dimensional multi-loop structure at higher levels of
abstraction. Introduce a transformation that coalesces nested loops with
independent bounds so that they can be further subdivided by tiling.

PiperOrigin-RevId: 258151016

Extend linalg transformations to allow value operands that are not views

This CL extends the linalg ops that can be tiled and fused to operations that take either views, scalar or vector operands.

PiperOrigin-RevId: 258149291

Decouple LLVM dialect from Standard dialect

Due to the absence of ODS support for enum attributes, the implementation of
the LLVM dialect `icmp` operation was reusing the comparison predicate from the
Standard dialect, creating an avoidable library dependency.  With ODS support
and ICmpPredicate attribute recently introduced, the dependency is no longer
justified.  Update the Standard to LLVM convresion to also convert the
CmpIPredicate into LLVM::ICmpPredicate and remove the unnecessary includes.

Note that the MLIRLLVMIR library did not explicitly depend on MLIRStandardOps,
requiring dependees of MLIRLLVMIR to also depend on MLIRStandardOps, which
should no longer be the case.

PiperOrigin-RevId: 258148456

LLVM Dialect: define ICmpPredicate in ODS

Use the recently introduced enum-gen functionality to define the predicate
attribute of the ICmp LLVM dialect operation directly in ODS.  This removes the
need for manually-coded string-to-integer conversion functions and contributes
to the overall homogenization of the operation definitions.

PiperOrigin-RevId: 258143923

Extract std.for std.if and std.terminator in their own dialect

These ops should not belong to the std dialect.
This CL extracts them in their own dialect and updates the corresponding conversions and tests.

PiperOrigin-RevId: 258123853

Fix a bug in DialectConversion when using RewritePattern.

When using a RewritePattern and replacing an operation with an existing value, that value may have already been replaced by something else. This cl ensures that only the final value is used when applying rewrites.

PiperOrigin-RevId: 258058488

Fix typos

Closes tensorflow/mlir#45

PiperOrigin-RevId: 257948893

NFC: Add header blocks to DialectConversion.h to improve readability.
PiperOrigin-RevId: 257903383

NFC: Don't print the location of a diagnostic if it is unknown.

Printing 'loc(unknown)' just clutters the output with unhelpful information.

PiperOrigin-RevId: 257883717

Add serialization and deserialization of FuncOps. To support this the
following SPIRV Instructions serializaiton/deserialization are added
as well

OpFunction
OpFunctionParameter
OpFunctionEnd
OpReturn

PiperOrigin-RevId: 257869806

[spirv] Various small code improvements

* Changed SPIR-V types to all use unsigned for member type count and index
* Used the same method name for getting element type and count
* Improved spv.CompositeExtract verification a bit

PiperOrigin-RevId: 257862580

Update the PatternRewriter constructor to take a context instead of a region.

This will allow for cleanly using a rewriter for multiple different regions.

PiperOrigin-RevId: 257845371

Change the IR printing pass instrumentation to ignore the verifier passes on non-failure.

The verifier passes are NO-OP and are only useful to print after in the case of failure. This removes a lot of unnecessary clutter when printing after/before all passes.

PiperOrigin-RevId: 257836310

Remove the 'region' field from OpBuilder.

This field wasn't updated as the insertion point changed, making it potentially dangerous given the multi-level of MLIR(e.g. 'createBlock' would always insert the new block in 'region'). This also allows for building an OpBuilder with just a context.

PiperOrigin-RevId: 257829135

Hide some public API of QuantizedType inherited from Type.

It is for preventing subtle bugs when writing predicates with wrong APIs. For example,

// BAD example for checking a quant type has an i8 or i16 storage type.
auto qtype = type.dyn_cast<mlir::quant::UniformQuantizedType>()
if (qtype.isSigned() && (qtype.isInteger(8) || qtype.isInteger(16)))

PiperOrigin-RevId: 257823263

[spirv] Add script to auto-generate SPIR-V op template from spec

SPIR-V has a JSON grammar file that defines the syntax of SPIR-V
instructions. However, its lacks fine-grained constraints on
instruction operands; those information is only available as
natural language sentences in the SPIR-V spec, which also contains
the detailed documentation for each SPIR-V instruction.

This CL pulls information from both the JSON grammar and HTML
spec. It right now uses the former to deduce the arguments and
results (with coarse-grained constraints) and the latter for
documentation. In the future we can add the functionality to
match certain natural language sentences for more fine-grained
constraints, but right now the developer is expected to update
the generated op definition. This should serve as a nice
bootstrap step to save efforts.

PiperOrigin-RevId: 257821205

Fix a bug in the canonicalizer when replacing constants via patterns.

The GreedyPatternRewriteDriver currently does not notify the OperationFolder when constants are removed as part of a pattern match. This materializes in a nasty bug where a different operation may be allocated to the same address. This causes an assertion in the OperationFolder when it gets notified of the new operations removal.

PiperOrigin-RevId: 257817627

Lower affine control flow to std control flow to LLVM dialect

This CL splits the lowering of affine to LLVM into 2 parts:
1. affine -> std
2. std -> LLVM

The conversions mostly consists of splitting concerns between the affine and non-affine worlds from existing conversions.
Short-circuiting of affine `if` conditions was never tested or exercised and is removed in the process, it can be reintroduced later if needed.

LoopParametricTiling.cpp is updated to reflect the newly added ForOp::build.

PiperOrigin-RevId: 257794436

[spirv] Add CompositeExtractOp operation.

CompositeExtractOp allows to extract a part of a composite object.

Closes tensorflow/mlir#44

PiperOrigin-RevId: 257790731

LoopsToGPU: use PassRegistration with constructor

PassRegistration with an optional constructor was introduced after the
LoopsToGPUPass, which resorted to deriving one pass from another as a means of
accepting options supplied as command-line arguments. Use PassRegistration with
constructor instead of defining a derived pass for LoopsToGPU.  Also rename the
pass to better reflect its current nature.

PiperOrigin-RevId: 257786923

LinalgTilingPass: use PassRegistration with a pass constructor

Linalg tiling pass was introduced before PassRegistration with an optional pass
constructor.  It resorted to deriving a helper class from the origial pass
class in order to provide a default constructor with values obtained from
command line flags.  Use PassRegistration with the optional pass constructor
instead, which avoids declaring an additional class.

PiperOrigin-RevId: 257786876

NFC: Remove redundant call to registerPassManagerCLOptions from MLIR tutorial

main already calls registerPassManagerCLOptions.

TESTED = not (NFC)
PiperOrigin-RevId: 257722013

Update the dialect attribute verifiers related to functions.

Remove the Function specific attribute verifier in favor of the general operation verifier. This also generalizes the function argument verifier to allow use for an argument attached to any region of any operation.

PiperOrigin-RevId: 257689962

Add support for serialization/deserialization of OpTypeVoid and
OpTypeFunction

PiperOrigin-RevId: 257684235

Rename FunctionAttr to SymbolRefAttr.

This allows for the attribute to hold symbolic references to other operations than FuncOp. This also allows for removing the dependence on FuncOp from the base Builder.

PiperOrigin-RevId: 257650017

FuncOp::eraseBody: drop all references before erasing blocks

Operations in a block can use a value defined in a dominating block.  When a
block, and therefore all its operations, is deleted, the operations are not
allowed to have any remaining uses.  Drop all uses of values in all blocks
before deleting them in FuncOp::eraseBody to avoid deleting an operation before
deleting the users of its results.

PiperOrigin-RevId: 257628002

affine.load/store: check for the right number of operands

Affine load and store operations take a variadic number of arguments, most of
which are interpreted as subscripts for the multi-dimensional memref they
access.  Add a verifier check that ensures the number of operands is equal to
the number affine remapping inputs if present and to the rank of the acessed
memref otherwise.  Although it is impossible to obtain such operations by
parsing the custom syntax, it is possible to construct them using the generic
syntax or programmatically.

PiperOrigin-RevId: 257605902

Propagate linalg op attributes in transformations

Also fix the tile_conv test.

PiperOrigin-RevId: 257602321

mcuMemHostRegister: take into account sizeof(float)

cuMemHostRegister expects the size of registered memory in bytes whereas the
memref descriptor in memref_t contains the number of elements.  Get the actual
size in bytes instead.

PiperOrigin-RevId: 257589116

Update the gen_spirv_dialect.py script to add opcodes from the SPIR-V
JSON spec into the SPIRBase.td file. This is done incrementally to
only import those opcodes that are needed, through use of the script
define_opcode.sh added.

PiperOrigin-RevId: 257517343

NFC: Replace Module::getNamedFunction with lookupSymbol<FuncOp>.

This allows for removing the last direct reference to FuncOp from ModuleOp.

PiperOrigin-RevId: 257498296

Refactor the parsing/printing of the top-level module.

This changes the top-level module parser to handle the case where the top-level module is defined with the module operation syntax, i.e:
  module ... {
  }

The printer is also updated to always print the top-level module in this form. This allows for cleanly round-tripping the location and attributes of the top-level module.

PiperOrigin-RevId: 257492069

Fix BufferSizeOp type lowering to LLVM.

This fixes a bug where the result type was incorrect when lowering to LLVM.

PiperOrigin-RevId: 257449384

PassRegistation: use overloads instead of a default argument

Windows build is not happy trying to convert a lambda into a constant reference
to a function in default arguments of the PassRegistration constructor.  Define
two overloaded constructors instead and construct the lambda as a variable in
the constructor body.

PiperOrigin-RevId: 257448633

Drop the trailing newline from the FuncOp syntax.

The ModulePrinter prints the newline now for children of the top-level module. This also fixes the location printing for functions as the location used to be printed on a different line.

PiperOrigin-RevId: 257447633

EDSC: use affine.load/store instead of std.load/store

Standard load and store operations are evolving to be separated from the Affine
constructs.  Special affine.load/store have been introduced to uphold the
restrictions of the Affine control flow constructs on their operands.
EDSC-produced loads and stores were originally intended to uphold those
restrictions as well so they should use affine.load/store instead of
std.load/store.

PiperOrigin-RevId: 257443307

NFC: Remove Function::getModule.

There is already a more general 'getParentOfType' method, and 'getModule' is likely to be misused as functions get placed within different regions than ModuleOp.

PiperOrigin-RevId: 257442243

Delete dead code

AffineIfOp::build is not tested or exercised anywhere. It also perpetuates a questionable choice of encoding an optional region as an empty region which we would like to change in the future.
PiperOrigin-RevId: 257439832

NFC: Rename Module to ModuleOp.

Module is a legacy name that only exists as a typedef of ModuleOp.

PiperOrigin-RevId: 257427248

Update ModuleOp::create(...) to take a Location instead of a context.

This allows for giving a Module a more interesting location than 'Unknown'.

PiperOrigin-RevId: 257310117

NFC: Rename Function to FuncOp.
PiperOrigin-RevId: 257293379

Constrain regex to avoid ordering issues.

PiperOrigin-RevId: 257257173

Relax the restriction that Modules cannot contain operations producing results.

This was an arbitrary restriction caused by the way that modules were printed. Now that that has been fixed, this restriction can be removed.

PiperOrigin-RevId: 257240329

Standardize the value numbering in the AsmPrinter.

Change the AsmPrinter to number values breadth-first so that values in adjacent regions can have the same name. This allows for ModuleOp to contain operations that produce results. This also standardizes the special name of region entry arguments to "arg[0-9+]" now that Functions are also operations.

PiperOrigin-RevId: 257225069

Add parentheses around boolean operators in an assertion in
Linalg/Transforms/Tiling.cpp.  This fixes a warning.

PiperOrigin-RevId: 257191302

Fix a test broken on some systems due to a mis-rebase.

PiperOrigin-RevId: 257190161

Implement parametric tiling on standard for loops

Parametric tiling can be used to extract outer loops with fixed number of
iterations.  This in turn enables mapping to GPU kernels on a fixed grid
independently of the range of the original loops, which may be unknown
statically, making the kernel adaptable to different sizes.  Provide a utility
function that also computes the parametric tile size given the range of the
loop.  Exercise the utility function through a simple pass that applies it to
all top-level loop nests.  Permutability or parallelism checks must be
performed before calling this utility function in actual passes.

Note that parametric tiling cannot be implemented in a purely affine way,
although it can be encoded using semi-affine maps.  The choice to implement it
on standard loops is guided by them being the common representation between
Affine loops, Linalg and GPU kernels.

PiperOrigin-RevId: 257180251

Extend AffineToGPU to support Linalg loops

Extend the utility that converts affine loop nests to support other types of
loops by abstracting away common behavior through templates.  This also
slightly simplifies the existing Affine to GPU conversion by always passing in
the loop step as an additional kernel argument even though it is a known
constant.  If it is used, it will be propagated into the loop body by the
existing canonicalization pattern and can be further constant-folded, otherwise
it will be dropped by canonicalization.

This prepares for the common loop abstraction that will be used for converting
to GPU kernels, which is conceptually close to Linalg loops, while maintaining
the existing conversion operational.

PiperOrigin-RevId: 257172216

NFC: Change a usage of StringLiteral to StringRef.

This matches the result type of other functions returning c-string, and fixes a build failure where some compilers fail to construct the StringLiteral.

PiperOrigin-RevId: 257139513

Generalize the symbol table functionality of ModuleOp into a trait 'OpTrait::SymbolTable'.

Operations must only contain a single region.  Once attached, all operations that contain a 'mlir::SymbolTable::getSymbolAttrName()' StringAttr attribute within the child region will be verified to ensure that the names are uniqued. Operations using this trait also gain access to the 'SymbolTable' class, which can be used to manage the symbol table of the operation. This class also provides constant-time lookup of symbols by name, and will automatically rename symbols on insertion.

PiperOrigin-RevId: 257123573

NFC: Remove `Module::getFunctions` in favor of a general `getOps<T>`.

Modules can now contain more than just Functions, this just updates the iteration API to reflect that. The 'begin'/'end' methods have also been updated to iterate over opaque Operations.

PiperOrigin-RevId: 257099084

NFC: Make the 'disable-pass-threading' flag a PassManagerOption.

This also adds the ability to programmatically disable threading.

PiperOrigin-RevId: 257051809

Remove dead code.

PiperOrigin-RevId: 257050780

Add dependencies for standard ops to SPIR-V conversion

PiperOrigin-RevId: 257026374

Add missing override.

PiperOrigin-RevId: 257024265

NFC: Remove the various "::getFunction" methods.

These methods assume that a function is a valid builtin top-level operation, and removing these methods allows for decoupling FuncOp and IR/. Utility "getParentOfType" methods have been added to Operation/OpState to allow for querying the first parent operation of a given type.

PiperOrigin-RevId: 257018913

NFC: Add AbstractOperation as a friend of 'Op'.

This allows for accessing the hidden overload of 'Op::hasTrait'.

PiperOrigin-RevId: 257018255

Add Ops in SPIR-V Dialect corresponding to OpEntryPoint and OpExecutionMode

PiperOrigin-RevId: 257013183

Remove dead code.

PiperOrigin-RevId: 257012369

Add 'hasTrait' support to Operation.

This is accomplished by adding a new callback to AbstractOperation that takes a ClassID for a given trait type and compares that against the ClassIDs for the traits used when registering an operation. This allows for opaquely querying if an operation was registered with a given trait, e.g:

Operation *op = ...;
bool singleResult = op->hasTrait<OpTrait::OneResult>();

PiperOrigin-RevId: 257007710

[spirv] NFC: Remove functions wrapping spv.module in tests

Now both functions and modules are just general ops and we do not require
top-level entities in a module's block to be the old builtin functions
any more. Removing the wrapping functions to simplify the tests.

PiperOrigin-RevId: 257003572

Add missing overrides.

PiperOrigin-RevId: 256995548

Remove std::move of trivially-copyable type.

Address ClangTidy finding:

* std::move of the expression of the trivially-copyable type 'mlir::Module' (aka 'mlir::ModuleOp') has no effect; remove std::move()

PiperOrigin-RevId: 256981849

Include missing header.

This target was failing to build with newer version of libc++.

PiperOrigin-RevId: 256979592

LinalgOp::Model: add missing override specifiers

All non-static functions in LinalgOp::Model override pure virtual functions
from LinalgOp::Concept, but only some of them have the `override` specifier.
This may lead one to erroneously believe the functions without the specified
are not overrides or virtual functions.  Add the specifier to all relevant
functions.

PiperOrigin-RevId: 256948453

Use isa instead of dyn_cast as cast value isn't used.

Avoids unused variable warning.

PiperOrigin-RevId: 256874512

Migrate NativeCodeCall and AllAttrConstraintsOf tests to use TestDialect

PiperOrigin-RevId: 256685260

Migrate pattern attribute matching tests to use TestDialect

This CL also reorders op definitions and tests a bit to make them
group more logically.

PiperOrigin-RevId: 256682105

Move StdForOp to ODS ForOp

PiperOrigin-RevId: 256657155

Add a standard if op

This CL adds an "std.if" op to represent an if-then-else construct whose condition is an arbitrary value of type i1.
This is necessary to lower all the existing examples from affine and linalg to std.for + std.if.

This CL introduces the op and adds the relevant positive and negative unit test. Lowering will be done in a separate followup CL.

PiperOrigin-RevId: 256649138

Simplify launch_func rewrite pattern in mlir-cuda-runner

Use the bulk setOperands on the cloned operation, no need to cast or iterate
over the operand list.

PiperOrigin-RevId: 256643496

ODS: provide a flag to skip generation of default build methods

Some operations need to override the default behavior of builders, in
particular region-holding operations such as affine.for or tf.graph want to
inject default terminators into the region upon construction, which default
builders won't do.  Provide a flag that disables the generation of default
builders so that the custom builders could use the same function signatures.
This is an intentionally low-level and heavy-weight feature that requires the
entire builder to be implemented, and it should be used sparingly.  Injecting
code into the end of a default builder would depend on the naming scheme of the
default builder arguments that is not visible in the ODS.  Checking that the
signature of a custom builder conflicts with that of a default builder to
prevent emission would require teaching ODG to differentiate between types and
(optional) argument names in the generated C++ code.  If this flag ends up
being used a lot, we should consider adding traits that inject specific code
into the default builder.

PiperOrigin-RevId: 256640069

Add ODS accessors for named regions.

PiperOrigin-RevId: 256639471

Trivial typo in comment fix.

PiperOrigin-RevId: 256616187

NFC: Remove Region::getContainingFunction as Functions are now Operations.

PiperOrigin-RevId: 256579717

Fix another incorrect forward declaration of LogicalResult

LogicalResult is a struct and the Windows build is unhappy about it being
forward-declared as class.

PiperOrigin-RevId: 256578230

Forward-declare LogicalResult as struct

LogicalResult is a struct but some headers forward-declare it as a class, which
breaks Windows builds.  Forward-delcare as a struct instead.

PiperOrigin-RevId: 256565699

Make TranslateFromMLIRFunction type return LogicalResult instead of bool

This interface was created before MLIR introduced LogicalResult.

PiperOrigin-RevId: 256554557

Add an mlir-cuda-runner tool.

This tool allows to execute MLIR IR snippets written in the GPU dialect
on a CUDA capable GPU. For this to work, a working CUDA install is required
and the build has to be configured with MLIR_CUDA_RUNNER_ENABLED set to 1.

PiperOrigin-RevId: 256551415

AllocOp: drop redundant verifier check

The equality check between the rank of a memref and the input size of the
layout affine map in AllocOp::verify is subsumed by the well-formedness check
of the memref type itself.  Drop the redundant check from the verifier since it
is never exercised (the type builder does not allow one to construct a type
that would not pass the verifier check).

PiperOrigin-RevId: 256551247

Add missing mlir:: namespace in definition of createConvertToLLVMIRPass.

PiperOrigin-RevId: 256546769

Make ConvertStandardToLLVMPass extendable with other patterns

Extend the LLVM lowering pass to accept callbacks that construct an instance of
(a subclass of) LLVMTypeConverter and populate a list of conversion patterns.
These callbacks will be called when the pass processes a module and their
results will be used to set up the dialect conversion infrastructure.  Clients
can now provide additional conversion patterns to avoid the need of
materializing type conversions between LLVM and other types.

PiperOrigin-RevId: 256532415

NFC: Refactoring to remove code bloat in SPIRV due to handling of Enum
Class Attribute parsing

PiperOrigin-RevId: 256471248

[TableGen] Support creating multi-result ops in result patterns

This CL introduces a new syntax for creating multi-result ops and access their
results in result patterns. Specifically, if a multi-result op is unbound or
bound to a name without a trailing `__N` suffix, it will act as a value pack
and expand to all its values. If a multi-result op is bound to a symbol with
`__N` suffix, only the N-th result will be extracted and used.

PiperOrigin-RevId: 256465208

Replace the implementation of Function and Module with FuncOp and ModuleOp.

This is an important step in allowing for the top-level of the IR to be extensible. FuncOp and ModuleOp contain all of the necessary functionality, while using the existing operation infrastructure. As an interim step, many of the usages of Function and Module, including the name, will remain the same. In the future, many of these will be relaxed to allow for many different types of top-level operations to co-exist.

PiperOrigin-RevId: 256427100

Globally change load/store/dma_start/dma_wait operations over to affine.load/store/dma_start/dma_wait.
In most places, this is just a name change (with the exception of affine.dma_start swapping the operand positions of its tag memref and num_elements operands).
Significant code changes occur here:
*) Vectorization: LoopAnalysis.cpp, Vectorize.cpp
*) Affine Transforms: Transforms/Utils/Utils.cpp

PiperOrigin-RevId: 256395088

More general subview calculation in tiling

This CL refactors tiling to enable tiling of views that are not just specified by a simple permutation. This allows the tiling of convolutions for which a new example is added.

PiperOrigin-RevId: 256346028

Add a generic loop abstraction to the std dialect

This CL is the first step of a refactoring unification of the control flow abstraction used in different dialects. The `std.for` loop accepts unrestricted indices to encode min, max and step and will be used as a common abstraction on the way to lower level dialects.

PiperOrigin-RevId: 256331795

NFC: Allow clearing the functions of a Module and splicing the functions from other modules.

PiperOrigin-RevId: 256280543

Add support for getting and setting the dialect attributes of an operation.

This allows for handling dialect attributes, those with names starting with a dialect prefix, to be handled explicitly.

PiperOrigin-RevId: 256279728

Add a 'getOps<>' method to Block to only iterate operations of a certain kind.

PiperOrigin-RevId: 256277436

NFC: Move Standard to SPIR-V conversion to lib/Conversion

PiperOrigin-RevId: 256271759

NFC: Move the Function/Module/Operation::verify methods out-of-line.

As Functions/Modules becomes operations, these methods will conflict with the 'verify' hook already on derived operation types.

PiperOrigin-RevId: 256246112

Add support for promoting Linalg views into new buffers.

This CL uses the generic CopyOp to promote a subview (constructed during tiling) into a new buffer + copy by:
  1. Creating a new buffer for the subview.
  2. Taking a view into the buffer and copying into it.
  3. Adapting the linalg op to operating on the view from point 2.
Tiling is extended with a boolean flag to enable promoting views (all or nothing for now).

More specifically, the current implementation creates a buffer that is always of the full size of the ranges of the subview. This produces a buffer whose size may be bigger
than the actual size of the `subView` at the boundaries and is related to the full/partial tile problem.

In practice, we introduce a `buffer`, a `fullLocalView` and a `partialLocalView` such that:
  * `buffer` is always the size of the subview in the full tile case.
  * `fullLocalView` is a dense contiguous view into that buffer.
  * `partialLocalView` is a dense non-contiguous slice of `fullLocalView`
     that corresponds to the size of `subView` and accounting for boundary
     effects.
The point of the full tile buffer is that constant static tile sizes are
folded and result in a buffer type with statically known size and alignment
properties.

Padding is introduced on the boundary tiles with a `fill` op followed by a partial `copy` op.

These behaviors will be refined later, on a per-need basis.

PiperOrigin-RevId: 256237319

Add support for SPIR-V Struct Types. Current support is limited to
supporting only Offset decorations

PiperOrigin-RevId: 256216704

[spirv] Various small improvements

* Added comments
* Improved op creation
* Used LogicalResult where suitable

PiperOrigin-RevId: 256203068